diff --git a/.asf.yaml b/.asf.yaml
new file mode 100644
index 00000000000..dd8bbfe04aa
--- /dev/null
+++ b/.asf.yaml
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# https://cwiki.apache.org/confluence/display/INFRA/Git+-+.asf.yaml+features
+
+github:
+ description: "Collaborative Machine-Learning-Centric Data Analytics Using Workflows"
+ homepage: https://texera.io/
+ labels:
+ - workflow
+ - data-science
+ - data
+ - machine-learning
+ - artificial-intelligence
+ - cloud-native
+ - data-analytics
+ - texera
+
+ protected_tags:
+ - "v*.*.*"
+
+ dependabot_alerts: true
+ dependabot_updates: false
+
+ features:
+ # Enable wiki for documentation
+ wiki: true
+ # Enable issue management
+ issues: true
+ # Enable projects for project management boards
+ projects: true
+ # Enable github discussions
+ discussions: true
+
+ pull_requests:
+ # allow auto-merge
+ allow_auto_merge: true
+ # enable updating head branches of pull requests
+ allow_update_branch: true
+ # auto-delete head branches after being merged
+ del_branch_on_merge: true
+
+ enabled_merge_buttons:
+ squash: true
+ squash_commit_message: PR_TITLE_AND_DESC
+ merge: false
+ rebase: false
+
+ protected_branches:
+ main:
+ required_status_checks:
+ # strict means "Require branches to be up to date before merging".
+ strict: true
+ # contexts are the names of checks that must pass
+ contexts:
+ - Required Checks
+ - Check License Headers
+ - Validate PR title
+ required_pull_request_reviews:
+ dismiss_stale_reviews: false
+ require_code_owner_reviews: false
+ required_approving_review_count: 1
+ required_linear_history: true
+
+notifications:
+ commits: commits@texera.apache.org
+ issues: notifications@texera.apache.org
+ pullrequests: notifications@texera.apache.org
+ discussions: dev@texera.apache.org
+ jobs: commits@texera.apache.org
diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 00000000000..0a608d39bbd
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,135 @@
+# Ignore all directories named `user-resources` anywhere in the project
+**/user-resources/
+
+# Ignoring binary/output
+**/target/
+**/out/
+
+# Ignoring packages
+*.jar
+*.war
+*.nar
+*.ear
+*.zip
+*.tar.gz
+*.rar
+
+# Ignoring VSCode related files
+.vscode/
+
+# Ignoring IntelliJ related files
+*.iml
+.idea/
+.idea_modules/
+lib_managed/
+src_managed/
+
+# Ignoring Eclipse files
+.classpath
+.project
+.settings
+
+# Ignoring sublime files
+*.sublime-workspace
+
+# Ignoring index folder and data folder
+index/
+catalog/
+plan/
+plan_files/
+query-results/
+
+# Ignoring Mac OSX specific files
+.DS_Store
+
+# Ignoring jenv related files
+.java-version
+
+# Ignoring scala related files
+hs_err_pid*
+
+# Ignoring Python related files
+venv/
+__pycache__/
+*.py[cod]
+*$py.class
+.ipynb_checkpoints
+.pytype/
+
+# Ignoring Python-generated files
+*.model
+*.pkl
+
+# Ignoring user-generated resources
+user-resources/
+
+# Ignoring Gmail tokens
+gmail/
+
+# Ignoring Maven-related files
+pom.xml.tag
+pom.xml.releaseBackup
+pom.xml.versionsBackup
+pom.xml.next
+release.properties
+dependency-reduced-pom.xml
+buildNumber.properties
+.mvn/timing.properties
+.mvn/wrapper/maven-wrapper.jar
+
+# Ignoring sbt related files
+.bsp/
+sbt.json
+
+# Ignoring rebel related files
+rebel.xml
+
+# Ignoring log files
+*.log
+*.log.gz
+
+# Ignoring the entire log folder
+log/
+
+# Ignoring package-lock.json
+package-lock.json
+
+# Ignoring Protobuf related files
+scalapb/scalapb
+
+# Ignoring credentials
+client_secret_*
+StoredCredential*
+**/apache2/
+**/Apache24/
+**/php/
+Composer-Setup.exe
+
+# Ignoring folders generated by VSCode IDE
+.metals/
+.bloop/
+.ammonite/
+metals.sbt
+
+# === NEW: Ignore frontend-related files ===
+# Ignore node_modules in all subdirectories
+**/node_modules/
+**/.pnp/
+**/.pnp.js
+
+# Ignore Angular build output
+**/dist/
+**/.angular/cache/
+**/.nx/cache/
+
+# Ignore Yarn cache and lock files
+**/.yarn/cache/
+**/.yarn/install-state.gz
+**/.pnp.cjs
+**/.pnp.loader.mjs
+
+# Ignore frontend dependency-related files
+**/yarn-error.log
+**/.turbo/
+**/.next/
+**/coverage/
\ No newline at end of file
diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 00000000000..6313b56c578
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1 @@
+* text=auto eol=lf
diff --git a/.github/ISSUE_TEMPLATE/bug-template.yaml b/.github/ISSUE_TEMPLATE/bug-template.yaml
new file mode 100644
index 00000000000..7a3ebc92022
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug-template.yaml
@@ -0,0 +1,81 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Bug Report
+description: File a bug report.
+labels: ["triage"]
+type: "Bug"
+body:
+ - type: markdown
+ attributes:
+ value: |
+ Thanks for taking the time to fill out this bug report!
+ - type: textarea
+ id: what-happened
+ attributes:
+ label: What happened?
+ description: Also tell us, what did you expect to happen?
+ placeholder: Tell us what you see!
+ value: "A bug happened!"
+ validations:
+ required: true
+ - type: textarea
+ id: reproduce
+ attributes:
+ label: How to reproduce?
+ description: Please include steps for a repro.
+ validations:
+ required: true
+ - type: dropdown
+ id: version
+ attributes:
+ label: Version
+ description: What version of Texera are you running?
+ options:
+ - 1.1.0-incubating (Pre-release/Master)
+ - 1.0.0
+ default: 0
+ validations:
+ required: true
+ - type: input
+ id: commit-hash
+ attributes:
+ label: Commit Hash (Optional)
+ description: If you know the specific commit that has the issue, please provide the commit hash here.
+ placeholder: e.g., a1b2c3d4e5f6g7h8i9j0k1l2m3n4o5p6q7r8s9t0
+ validations:
+ required: false
+ - type: dropdown
+ id: browsers
+ attributes:
+ label: What browsers are you seeing the problem on?
+ multiple: true
+ options:
+ - Chrome
+ - Safari
+ - Firefox
+ - Microsoft Edge
+ - type: textarea
+ id: logs
+ attributes:
+ label: Relevant log output
+ description: Please copy and paste any relevant log output. This will be automatically formatted into code, so no need for backticks.
+ render: shell
+
+ - type: markdown
+ attributes:
+ value: |
+ By submitting this issue, you agree to follow the [Apache Code of Conduct](https://www.apache.org/foundation/policies/conduct).
diff --git a/.github/ISSUE_TEMPLATE/feature-template.yaml b/.github/ISSUE_TEMPLATE/feature-template.yaml
new file mode 100644
index 00000000000..f88d85b0f67
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/feature-template.yaml
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Feature Request
+description: Suggest a new feature or improvement.
+labels: ["triage"]
+type: "Feature"
+body:
+ - type: markdown
+ attributes:
+ value: |
+ Thanks for suggesting a feature! Please provide as much detail as possible to help us evaluate your idea.
+
+ - type: textarea
+ id: summary
+ attributes:
+ label: Feature Summary
+ description: Clearly describe what the feature is and the problem it solves.
+ placeholder: Describe your feature idea and what problem it addresses.
+ validations:
+ required: true
+
+ - type: textarea
+ id: proposal
+ attributes:
+ label: Proposed Solution or Design
+ description: Explain how you imagine this feature working. Include examples, diagrams, or pseudo-code if relevant.
+ placeholder: Describe your proposed solution or design approach.
+ validations:
+ required: true
+
+ - type: dropdown
+ id: impact
+ attributes:
+ label: Impact / Priority
+ description: How important is this feature?
+ options:
+ - (P0)Critical – blocks existing use cases
+ - (P1)High – significantly improves user experience
+ - (P2)Medium – useful enhancement
+ - (P3)Low – nice to have
+ default: 2
+ validations:
+ required: true
+
+ - type: dropdown
+ id: affected-area
+ attributes:
+ label: Affected Area
+ description: Which part of the system does this feature relate to?
+ multiple: true
+ options:
+ - Workflow Engine (Amber)
+ - Workflow UI
+ - Hub
+ - Storage / Metadata
+ - Deployment / Infrastructure
+ - Other
+
+ - type: markdown
+ attributes:
+ value: |
+ By submitting this issue, you agree to follow the [Apache Code of Conduct](https://www.apache.org/foundation/policies/conduct).
\ No newline at end of file
diff --git a/.github/ISSUE_TEMPLATE/task-template.yaml b/.github/ISSUE_TEMPLATE/task-template.yaml
new file mode 100644
index 00000000000..1e299fef7c6
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/task-template.yaml
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Task
+description: Create a new development or maintenance task.
+labels: ["triage"]
+type: "Task"
+body:
+ - type: markdown
+ attributes:
+ value: |
+ Thanks for creating a task! Please describe what needs to be done and why.
+
+ - type: textarea
+ id: task-summary
+ attributes:
+ label: Task Summary
+ description: Briefly describe what needs to be done, try to do a single step in a task.
+ placeholder: Example — Refactor workflow scheduler module for better modularity.
+ validations:
+ required: true
+
+ - type: dropdown
+ id: priority
+ attributes:
+ label: Priority
+ description: How urgent or important is this task?
+ options:
+ - P0 – Critical
+ - P1 – High
+ - P2 – Medium
+ - P3 – Low
+ default: 2
+
+ - type: checkboxes
+ id: checklist
+ attributes:
+ label: Task Type
+ description: Select the type of work involved.
+ options:
+ - label: Code Implementation
+ - label: Documentation
+ - label: Refactor / Cleanup
+ - label: Testing / QA
+ - label: DevOps / Deployment
+
+ - type: markdown
+ attributes:
+ value: |
+ By submitting this issue, you agree to follow the [Apache Code of Conduct](https://www.apache.org/foundation/policies/conduct).
diff --git a/.github/PULL_REQUEST_TEMPLATE b/.github/PULL_REQUEST_TEMPLATE
new file mode 100644
index 00000000000..41287564ff3
--- /dev/null
+++ b/.github/PULL_REQUEST_TEMPLATE
@@ -0,0 +1,53 @@
+
+
+### What changes were proposed in this PR?
+
+
+
+### Any related issues, documentation, discussions?
+
+
+
+### How was this PR tested?
+
+
+
+### Was this PR authored or co-authored using generative AI tooling?
+
diff --git a/.github/labeler.yml b/.github/labeler.yml
new file mode 100644
index 00000000000..0e12f509750
--- /dev/null
+++ b/.github/labeler.yml
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+frontend:
+ - changed-files:
+ - any-glob-to-any-file:
+ - 'frontend/**'
+
+common:
+ - changed-files:
+ - any-glob-to-any-file:
+ - 'common/**'
+
+service:
+ - changed-files:
+ - any-glob-to-any-file:
+ - '*-service/**'
+
+engine:
+ - changed-files:
+ - any-glob-to-any-file:
+ - 'amber/**'
+
+python:
+ - changed-files:
+ - any-glob-to-any-file:
+ - 'amber/src/main/python/**'
+ - '**/*.py'
+
+docs:
+ - changed-files:
+ - any-glob-to-any-file:
+ - 'docs/**'
+ - '**/*.md'
+ - 'NOTICE'
+ - 'LICENSE'
+
+ci:
+ - changed-files:
+ - any-glob-to-any-file:
+ - '.github/workflows/**'
+
+dev:
+ - changed-files:
+ - any-glob-to-any-file:
+ - 'bin/**'
+
+dependencies:
+ - changed-files:
+ - any-glob-to-any-file:
+ - '**/requirements.txt'
+ - '**/package.json'
+ - '**/build.sbt'
+ - '**/project.sbt'
+
+ddl-change:
+ - changed-files:
+ - any-glob-to-any-file:
+ - '**/*.sql'
+
+feature:
+ - head-branch:
+ - '^feat'
+ - 'feature'
+
+fix:
+ - head-branch: '^fix'
+
+refactor:
+ - head-branch: '^refactor'
+
diff --git a/.github/release/vote-email-template.md b/.github/release/vote-email-template.md
new file mode 100644
index 00000000000..f3dd942a1d7
--- /dev/null
+++ b/.github/release/vote-email-template.md
@@ -0,0 +1,61 @@
+Subject: [VOTE] Release Apache Texera (incubating) ${VERSION} RC${RC_NUM}
+
+Hi Texera Community,
+
+This is a call for vote to release Apache Texera (incubating) ${VERSION}.
+
+== Release Candidate Artifacts ==
+
+https://dist.apache.org/repos/dist/dev/incubator/texera/${VERSION}-RC${RC_NUM}/
+
+The directory contains:
+- Source tarball (.tar.gz) with GPG signature (.asc) and SHA512 checksum (.sha512)
+- Docker Compose deployment bundle with GPG signature and SHA512 checksum
+- Helm chart package with GPG signature and SHA512 checksum
+
+== Container Images ==
+
+Container images are available at:
+ ${IMAGE_REGISTRY}/texera-dashboard-service:${VERSION}
+ ${IMAGE_REGISTRY}/texera-workflow-execution-coordinator:${VERSION}
+ ${IMAGE_REGISTRY}/texera-workflow-compiling-service:${VERSION}
+ ${IMAGE_REGISTRY}/texera-file-service:${VERSION}
+ ${IMAGE_REGISTRY}/texera-config-service:${VERSION}
+ ${IMAGE_REGISTRY}/texera-access-control-service:${VERSION}
+ ${IMAGE_REGISTRY}/texera-workflow-computing-unit-managing-service:${VERSION}
+
+These images are built from the source tarball included in this release.
+The Dockerfiles are included in the source for audit and verification.
+
+== Git Tag ==
+
+https://github.com/apache/texera/releases/tag/${TAG_NAME}
+Commit: ${COMMIT_HASH}
+
+== Keys ==
+
+The release was signed with GPG key [${GPG_KEY_ID}] (${GPG_EMAIL})
+KEYS file: https://downloads.apache.org/incubator/texera/KEYS
+
+== Vote ==
+
+The vote will be open for at least 72 hours.
+
+[ ] +1 Approve the release
+[ ] 0 No opinion
+[ ] -1 Disapprove the release (please provide the reason)
+
+== Checklist ==
+
+[ ] Checksums and PGP signatures are valid
+[ ] LICENSE and NOTICE files are correct
+[ ] All files have ASF license headers where appropriate
+[ ] No unexpected binary files
+[ ] Source tarball matches the Git tag
+[ ] Can compile from source successfully
+[ ] Docker Compose bundle deploys successfully with the published images
+[ ] Helm chart deploys successfully (if applicable)
+
+Thanks,
+[Your Name]
+Apache Texera (incubating) PPMC
diff --git a/.github/scripts/prepare-backport-checkout.sh b/.github/scripts/prepare-backport-checkout.sh
new file mode 100644
index 00000000000..5286275516a
--- /dev/null
+++ b/.github/scripts/prepare-backport-checkout.sh
@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+target_branch="${1:?target branch is required}"
+commit_range="${2:?commit range is required}"
+workspace_branch="ci-backport-${target_branch//\//-}"
+
+git fetch --no-tags origin "${target_branch}"
+git config user.name "github-actions[bot]"
+git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
+
+if [[ "${commit_range}" != *..* ]]; then
+ echo "Invalid commit range: ${commit_range}" >&2
+ exit 1
+fi
+start_sha="${commit_range%..*}"
+end_sha="${commit_range##*..}"
+
+if [[ -z "$(git rev-list -n 1 "${commit_range}")" ]]; then
+ echo "No commits found in range ${commit_range}" >&2
+ exit 1
+fi
+
+# Build a single squash commit whose parent is the range start and whose tree
+# matches the range end. Cherry-picking this squash onto the release branch
+# applies the cumulative diff in one 3-way merge, which avoids spurious
+# conflicts when intermediate commits in the range happen to overlap with
+# changes already present (under different SHAs) on the release branch.
+end_tree="$(git rev-parse "${end_sha}^{tree}")"
+squash_sha="$(git commit-tree -p "${start_sha}" -m "ci: squashed backport of ${commit_range}" "${end_tree}")"
+
+git checkout -B "${workspace_branch}" "origin/${target_branch}"
+git cherry-pick -x "${squash_sha}"
diff --git a/.github/workflows/auto-assign.yml b/.github/workflows/auto-assign.yml
new file mode 100644
index 00000000000..48fe10f6c2b
--- /dev/null
+++ b/.github/workflows/auto-assign.yml
@@ -0,0 +1,96 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Auto-assign
+on:
+ pull_request_target:
+ types: [opened, closed]
+
+permissions:
+ issues: write
+ pull-requests: write
+
+jobs:
+ assign-pr-author:
+ if: >-
+ github.event.action == 'opened'
+ && github.event.pull_request.user.type != 'Bot'
+ && github.event.pull_request.assignees[0] == null
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/github-script@v7
+ with:
+ script: |
+ const pr = context.payload.pull_request;
+ await github.rest.issues.addAssignees({
+ ...context.repo,
+ issue_number: pr.number,
+ assignees: [pr.user.login],
+ });
+
+ credit-issue-on-pr-merge:
+ if: github.event.action == 'closed' && github.event.pull_request.merged
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/github-script@v7
+ with:
+ script: |
+ const { owner, repo } = context.repo;
+ const opener = context.payload.pull_request.user;
+ const isHuman = (l) => l && !l.endsWith('[bot]');
+
+ const { repository: { pullRequest: prq } } = await github.graphql(`
+ query($owner: String!, $repo: String!, $pr: Int!) {
+ repository(owner: $owner, name: $repo) {
+ pullRequest(number: $pr) {
+ closingIssuesReferences(first: 50) {
+ nodes {
+ number
+ repository { nameWithOwner }
+ assignees(first: 20) { nodes { login } }
+ }
+ }
+ commits(first: 250) {
+ nodes { commit {
+ parents { totalCount }
+ authors(first: 10) { nodes { user { login } } }
+ } }
+ }
+ }
+ }
+ }`, { owner, repo, pr: context.payload.pull_request.number });
+
+ const authors = new Set();
+ if (opener.type !== 'Bot' && isHuman(opener.login)) authors.add(opener.login);
+ for (const { commit } of prq.commits.nodes) {
+ if (commit.parents.totalCount > 1) continue;
+ for (const a of commit.authors.nodes) {
+ if (isHuman(a.user?.login)) authors.add(a.user.login);
+ }
+ }
+ const credited = [...authors].slice(0, 10);
+ if (!credited.length) return;
+ const creditedSet = new Set(credited);
+
+ for (const issue of prq.closingIssuesReferences.nodes) {
+ if (issue.repository.nameWithOwner !== `${owner}/${repo}`) continue;
+ const current = issue.assignees.nodes.map(n => n.login);
+ const toRemove = current.filter(l => !creditedSet.has(l));
+ const toAdd = credited.filter(l => !current.includes(l));
+ const args = { owner, repo, issue_number: issue.number };
+ if (toRemove.length) await github.rest.issues.removeAssignees({ ...args, assignees: toRemove });
+ if (toAdd.length) await github.rest.issues.addAssignees({ ...args, assignees: toAdd });
+ }
diff --git a/.github/workflows/automatic-email-notif-on-ddl-change.yml b/.github/workflows/automatic-email-notif-on-ddl-change.yml
new file mode 100644
index 00000000000..782deaa2a77
--- /dev/null
+++ b/.github/workflows/automatic-email-notif-on-ddl-change.yml
@@ -0,0 +1,62 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Automatic email notification on DDL change
+
+on:
+ pull_request:
+ types:
+ - closed
+
+jobs:
+ notify:
+ if: >-
+ github.event.pull_request.merged == true &&
+ contains(github.event.pull_request.labels.*.name, 'ddl-change')
+ runs-on: ubuntu-latest
+ steps:
+ - name: Checkout
+ uses: actions/checkout@v4
+ with:
+ fetch-depth: 0
+ sparse-checkout: sql/updates/
+
+ - name: Get added file in sql/updates/
+ id: get_sql_file
+ run: |
+ FILE=$(git diff --name-only --diff-filter=A \
+ ${{ github.event.pull_request.base.sha }} \
+ ${{ github.event.pull_request.merge_commit_sha }} \
+ -- 'sql/updates/')
+ echo "sql_file=$FILE" >> $GITHUB_OUTPUT
+ - name: Send email
+ run: |
+ curl --ssl-reqd \
+ --url "smtps://smtp.gmail.com:465" \
+ --user "${{ secrets.NOREPLY_EMAIL_USERNAME }}:${{ secrets.NOREPLY_EMAIL_PASSWORD }}" \
+ --mail-from "${{ secrets.NOREPLY_EMAIL_USERNAME }}" \
+ --mail-rcpt "dev@texera.apache.org" \
+ --upload-file - <Hi all,
+ We have merged PR #${{ github.event.pull_request.number }} (${{ github.event.pull_request.html_url }} ): ${{ github.event.pull_request.title }}. To incorporate the change, please apply ${{ steps.get_sql_file.outputs.sql_file }} to your local Postgres instance and run sbt jooqGenerate to generate jooq classes.
+ Best, ${{ github.event.pull_request.user.login }}
+ EOF
\ No newline at end of file
diff --git a/.github/workflows/build-and-push-images.yml b/.github/workflows/build-and-push-images.yml
new file mode 100644
index 00000000000..d6688336b8a
--- /dev/null
+++ b/.github/workflows/build-and-push-images.yml
@@ -0,0 +1,549 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+name: Build and push images
+
+on:
+ workflow_dispatch:
+ inputs:
+ branch:
+ description: 'Branch to checkout and build from'
+ required: false
+ default: 'main'
+ type: string
+ image_tag:
+ description: 'Docker image tag (e.g., latest, v1.0.0). Leave empty to use the short commit hash of the branch.'
+ required: false
+ default: ''
+ type: string
+ docker_registry:
+ description: 'Full image registry prefix (e.g., ghcr.io/apache, docker.io/apache)'
+ required: false
+ default: 'ghcr.io/apache'
+ type: string
+ services:
+ description: 'Services to build (comma-separated, "*" for all)'
+ required: false
+ default: '*'
+ type: string
+ platforms:
+ description: 'Target platforms to build'
+ required: false
+ default: 'both'
+ type: choice
+ options:
+ - both
+ - amd64
+ - arm64
+ schedule:
+ # Run nightly at 2:00 AM UTC
+ - cron: '0 2 * * *'
+
+permissions:
+ contents: read
+ packages: write # Required for pushing to ghcr.io
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event.inputs.image_tag || 'nightly' }}
+ cancel-in-progress: false
+
+jobs:
+ # Step 0: Set runtime parameters (handles both manual and scheduled runs)
+ set-parameters:
+ runs-on: ubuntu-latest
+ if: github.event_name != 'schedule' || github.repository == 'apache/texera'
+ outputs:
+ branch: ${{ steps.set-params.outputs.branch }}
+ image_tag: ${{ steps.set-params.outputs.image_tag }}
+ docker_registry: ${{ steps.set-params.outputs.docker_registry }}
+ services: ${{ steps.set-params.outputs.services }}
+ platforms: ${{ steps.set-params.outputs.platforms }}
+ steps:
+ - name: Set build parameters
+ id: set-params
+ env:
+ GH_TOKEN: ${{ github.token }}
+ run: |
+ # Detect if this is a scheduled run
+ if [[ "${{ github.event_name }}" == "schedule" ]]; then
+ echo "Nightly build detected - using nightly defaults"
+ echo "branch=main" >> $GITHUB_OUTPUT
+ echo "image_tag=latest" >> $GITHUB_OUTPUT
+ echo "docker_registry=ghcr.io/apache" >> $GITHUB_OUTPUT
+ echo "services=*" >> $GITHUB_OUTPUT
+ echo "platforms=both" >> $GITHUB_OUTPUT
+ else
+ echo "Manual workflow_dispatch - using user inputs"
+ BRANCH="${{ github.event.inputs.branch || 'main' }}"
+ IMAGE_TAG="${{ github.event.inputs.image_tag }}"
+
+ # If image_tag is empty, resolve to the short commit hash of the branch
+ if [[ -z "$IMAGE_TAG" ]]; then
+ COMMIT_SHORT=$(gh api "repos/${{ github.repository }}/commits/${BRANCH}" --jq '.sha[:9]')
+ IMAGE_TAG="$COMMIT_SHORT"
+ echo "No image tag specified - using short commit hash: $IMAGE_TAG"
+ fi
+
+ echo "branch=$BRANCH" >> $GITHUB_OUTPUT
+ echo "image_tag=$IMAGE_TAG" >> $GITHUB_OUTPUT
+ echo "docker_registry=${{ github.event.inputs.docker_registry || 'ghcr.io/apache' }}" >> $GITHUB_OUTPUT
+ echo "services=${{ github.event.inputs.services || '*' }}" >> $GITHUB_OUTPUT
+ echo "platforms=${{ github.event.inputs.platforms || 'both' }}" >> $GITHUB_OUTPUT
+ fi
+
+ # Step 1: Generate JOOQ code once and share it
+ generate-jooq:
+ needs: [set-parameters]
+ runs-on: ubuntu-latest
+ env:
+ JAVA_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
+ JVM_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
+
+ steps:
+ - name: Checkout Texera
+ uses: actions/checkout@v5
+ with:
+ ref: ${{ needs.set-parameters.outputs.branch }}
+
+ - name: Setup JDK
+ uses: actions/setup-java@v5
+ with:
+ distribution: 'temurin'
+ java-version: 11
+
+ - name: Setup sbt launcher
+ uses: sbt/setup-sbt@3e125ece5c3e5248e18da9ed8d2cce3d335ec8dd # v1.1.14
+
+ - uses: coursier/cache-action@90c37294538be80a558fd665531fcdc2b467b475 # v8.1.0
+ with:
+ extraSbtFiles: '["*.sbt", "project/**.{scala,sbt}", "project/build.properties" ]'
+
+ - name: Install PostgreSQL
+ run: sudo apt-get update && sudo apt-get install -y postgresql
+
+ - name: Start PostgreSQL Service
+ run: sudo systemctl start postgresql
+
+ - name: Configure PostgreSQL authentication
+ run: |
+ sudo -u postgres psql -c "ALTER USER postgres PASSWORD 'postgres';"
+ sudo sed -i 's/local all postgres peer/local all postgres md5/' /etc/postgresql/*/main/pg_hba.conf
+ sudo sed -i 's/host all all 127.0.0.1\/32 scram-sha-256/host all all 127.0.0.1\/32 md5/' /etc/postgresql/*/main/pg_hba.conf
+ sudo systemctl restart postgresql
+ sleep 2
+
+ - name: Create Databases
+ run: |
+ PGPASSWORD=postgres psql -h localhost -U postgres -f sql/texera_ddl.sql
+ PGPASSWORD=postgres psql -h localhost -U postgres -f sql/iceberg_postgres_catalog.sql
+ PGPASSWORD=postgres psql -h localhost -U postgres -f sql/texera_lakefs.sql
+
+ - name: Generate JOOQ code
+ run: sbt DAO/jooqGenerate
+
+ - name: Upload JOOQ generated code
+ uses: actions/upload-artifact@v4
+ with:
+ name: jooq-code
+ path: |
+ common/dao/src/main/scala/org/apache/texera/dao/jooq/generated/
+ retention-days: 1
+
+ # Step 2: Parse services and prepare build matrix
+ prepare-matrix:
+ needs: [set-parameters]
+ runs-on: ubuntu-latest
+ outputs:
+ matrix: ${{ steps.set-matrix.outputs.matrix }}
+ build_amd64: ${{ steps.set-platforms.outputs.build_amd64 }}
+ build_arm64: ${{ steps.set-platforms.outputs.build_arm64 }}
+ need_manifest: ${{ steps.set-platforms.outputs.need_manifest }}
+ steps:
+ - name: Checkout Texera
+ uses: actions/checkout@v5
+ with:
+ ref: ${{ needs.set-parameters.outputs.branch }}
+
+ - name: Set target platforms
+ id: set-platforms
+ run: |
+ PLATFORM_INPUT="${{ needs.set-parameters.outputs.platforms }}"
+
+ case "$PLATFORM_INPUT" in
+ both)
+ echo "build_amd64=true" >> $GITHUB_OUTPUT
+ echo "build_arm64=true" >> $GITHUB_OUTPUT
+ echo "need_manifest=true" >> $GITHUB_OUTPUT
+ echo "Building for both platforms (parallel jobs)"
+ ;;
+ amd64)
+ echo "build_amd64=true" >> $GITHUB_OUTPUT
+ echo "build_arm64=false" >> $GITHUB_OUTPUT
+ echo "need_manifest=false" >> $GITHUB_OUTPUT
+ echo "Building for AMD64 only"
+ ;;
+ arm64)
+ echo "build_amd64=false" >> $GITHUB_OUTPUT
+ echo "build_arm64=true" >> $GITHUB_OUTPUT
+ echo "need_manifest=false" >> $GITHUB_OUTPUT
+ echo "Building for ARM64 only"
+ ;;
+ esac
+
+ - name: Discover and parse services
+ id: set-matrix
+ run: |
+ SERVICES="${{ needs.set-parameters.outputs.services }}"
+
+ # Discover all Dockerfiles in bin/ directory
+ echo "Discovering services from Dockerfiles..."
+ cd bin
+
+ # Standard services from *.dockerfile pattern (excluding postgres17-pgroonga)
+ STANDARD_SERVICES=()
+ for dockerfile in *.dockerfile; do
+ if [[ -f "$dockerfile" ]]; then
+ service_name=$(basename "$dockerfile" .dockerfile)
+ # Skip postgres17-pgroonga
+ if [[ "$service_name" != "postgres17-pgroonga" ]]; then
+ STANDARD_SERVICES+=("$service_name")
+ fi
+ fi
+ done
+
+ # All services are standard services only
+ ALL_SERVICES=("${STANDARD_SERVICES[@]}")
+
+ echo "Found ${#ALL_SERVICES[@]} services: ${ALL_SERVICES[*]}"
+
+ # Filter based on user input
+ if [[ "$SERVICES" == "*" ]]; then
+ SERVICES_LIST=("${ALL_SERVICES[@]}")
+ else
+ IFS=',' read -ra SERVICES_LIST <<< "$SERVICES"
+ # Trim whitespace
+ for i in "${!SERVICES_LIST[@]}"; do
+ SERVICES_LIST[$i]=$(echo "${SERVICES_LIST[$i]}" | xargs)
+ done
+ fi
+
+ # Create JSON matrix with dockerfile info
+ JSON="["
+ FIRST=true
+ for service in "${SERVICES_LIST[@]}"; do
+ # Determine dockerfile path and context
+ if [[ " ${STANDARD_SERVICES[@]} " =~ " ${service} " ]]; then
+ dockerfile="bin/${service}.dockerfile"
+ context="."
+
+ # Map dockerfile service names to Docker image names
+ case "$service" in
+ "texera-web-application")
+ image_name="texera-dashboard-service"
+ ;;
+ "computing-unit-master")
+ image_name="texera-workflow-execution-coordinator"
+ ;;
+ "computing-unit-worker")
+ image_name="texera-workflow-execution-runner"
+ ;;
+ "access-control-service")
+ image_name="texera-access-control-service"
+ ;;
+ "config-service")
+ image_name="texera-config-service"
+ ;;
+ "file-service")
+ image_name="texera-file-service"
+ ;;
+ "workflow-compiling-service")
+ image_name="texera-workflow-compiling-service"
+ ;;
+ "workflow-computing-unit-managing-service")
+ image_name="texera-workflow-computing-unit-managing-service"
+ ;;
+ "agent-service")
+ image_name="texera-agent-service"
+ ;;
+ *)
+ # Default: use service name as-is
+ image_name="$service"
+ ;;
+ esac
+ else
+ echo "WARNING: Unknown service: $service, skipping"
+ continue
+ fi
+
+ if [[ "$FIRST" == "true" ]]; then
+ FIRST=false
+ else
+ JSON+=","
+ fi
+ JSON+="{\"service\":\"$service\",\"image_name\":\"$image_name\",\"dockerfile\":\"$dockerfile\",\"context\":\"$context\"}"
+ done
+ JSON+="]"
+
+ echo "Generated matrix: $JSON"
+ echo "matrix={\"include\":$JSON}" >> $GITHUB_OUTPUT
+
+ # Step 3a: Build AMD64 images (runs in parallel with ARM64)
+ build-amd64:
+ runs-on: ubuntu-latest
+ needs: [set-parameters, generate-jooq, prepare-matrix]
+ if: needs.prepare-matrix.outputs.build_amd64 == 'true'
+ strategy:
+ matrix: ${{ fromJson(needs.prepare-matrix.outputs.matrix) }}
+ fail-fast: false
+ max-parallel: 8 # Higher parallelism for native builds
+ env:
+ DOCKER_REGISTRY: ${{ needs.set-parameters.outputs.docker_registry }}
+ JAVA_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
+ JVM_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
+
+ steps:
+ - name: Checkout Texera
+ uses: actions/checkout@v5
+ with:
+ ref: ${{ needs.set-parameters.outputs.branch }}
+
+ - name: Setup JDK
+ uses: actions/setup-java@v5
+ with:
+ distribution: 'temurin'
+ java-version: 11
+
+ - name: Setup sbt launcher
+ uses: sbt/setup-sbt@508b753e53cb6095967669e0911487d2b9bc9f41 # v1.1.22
+
+ - uses: coursier/cache-action@90c37294538be80a558fd665531fcdc2b467b475 # v8.1.0
+ with:
+ extraSbtFiles: '["*.sbt", "project/**.{scala,sbt}", "project/build.properties" ]'
+
+ - name: Download JOOQ generated code
+ uses: actions/download-artifact@v4
+ with:
+ name: jooq-code
+ path: common/dao/src/main/scala/org/apache/texera/dao/jooq/generated/
+
+ - name: Free up disk space
+ run: |
+ sudo apt-get clean
+ sudo rm -rf /usr/share/dotnet
+ sudo rm -rf /opt/ghc
+ sudo rm -rf /usr/local/share/boost
+ df -h
+
+ - name: Set up Docker Buildx
+ uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.11.1
+
+ - name: Log in to GitHub Container Registry
+ if: startsWith(needs.set-parameters.outputs.docker_registry, 'ghcr.io/')
+ uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.6.0
+ with:
+ registry: ghcr.io
+ username: ${{ github.actor }}
+ password: ${{ secrets.GITHUB_TOKEN }}
+
+ - name: Log in to Docker Hub
+ if: ${{ !startsWith(needs.set-parameters.outputs.docker_registry, 'ghcr.io/') }}
+ uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.6.0
+ with:
+ username: ${{ secrets.DOCKERHUB_USERNAME }}
+ password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+ - name: Build and push AMD64 image
+ uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6.18.0
+ with:
+ context: ${{ matrix.context }}
+ file: ${{ matrix.dockerfile }}
+ platforms: linux/amd64
+ push: true
+ tags: ${{ env.DOCKER_REGISTRY }}/${{ matrix.image_name }}:${{ needs.set-parameters.outputs.image_tag }}-amd64
+ cache-from: type=gha,scope=${{ matrix.image_name }}-amd64
+ cache-to: type=gha,mode=max,scope=${{ matrix.image_name }}-amd64
+ labels: |
+ org.opencontainers.image.title=${{ matrix.image_name }}
+ org.opencontainers.image.description=Apache Texera ${{ matrix.image_name }} (AMD64)
+ org.opencontainers.image.vendor=Apache Texera
+
+ # Step 3b: Build ARM64 images (runs in parallel with AMD64)
+ build-arm64:
+ runs-on: ubuntu-latest
+ needs: [set-parameters, generate-jooq, prepare-matrix]
+ if: needs.prepare-matrix.outputs.build_arm64 == 'true'
+ strategy:
+ matrix: ${{ fromJson(needs.prepare-matrix.outputs.matrix) }}
+ fail-fast: false
+ max-parallel: 4 # Lower for QEMU builds
+ env:
+ DOCKER_REGISTRY: ${{ needs.set-parameters.outputs.docker_registry }}
+ JAVA_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
+ JVM_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
+
+ steps:
+ - name: Checkout Texera
+ uses: actions/checkout@v5
+ with:
+ ref: ${{ needs.set-parameters.outputs.branch }}
+
+ - name: Setup JDK
+ uses: actions/setup-java@v5
+ with:
+ distribution: 'temurin'
+ java-version: 11
+
+ - name: Setup sbt launcher
+ uses: sbt/setup-sbt@508b753e53cb6095967669e0911487d2b9bc9f41 # v1.1.22
+
+ - uses: coursier/cache-action@90c37294538be80a558fd665531fcdc2b467b475 # v8.1.0
+ with:
+ extraSbtFiles: '["*.sbt", "project/**.{scala,sbt}", "project/build.properties" ]'
+
+ - name: Download JOOQ generated code
+ uses: actions/download-artifact@v4
+ with:
+ name: jooq-code
+ path: common/dao/src/main/scala/org/apache/texera/dao/jooq/generated/
+
+ - name: Free up disk space
+ run: |
+ sudo apt-get clean
+ sudo rm -rf /usr/share/dotnet
+ sudo rm -rf /opt/ghc
+ sudo rm -rf /usr/local/share/boost
+ df -h
+
+ # Set up QEMU for ARM64 emulation
+ - name: Set up QEMU
+ uses: docker/setup-qemu-action@29109295f81e9208d7d86ff1c6c12d2833863392 # v3.7.0
+ with:
+ platforms: linux/arm64
+
+ - name: Set up Docker Buildx
+ uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.11.1
+
+ - name: Log in to GitHub Container Registry
+ if: startsWith(needs.set-parameters.outputs.docker_registry, 'ghcr.io/')
+ uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.6.0
+ with:
+ registry: ghcr.io
+ username: ${{ github.actor }}
+ password: ${{ secrets.GITHUB_TOKEN }}
+
+ - name: Log in to Docker Hub
+ if: ${{ !startsWith(needs.set-parameters.outputs.docker_registry, 'ghcr.io/') }}
+ uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.6.0
+ with:
+ username: ${{ secrets.DOCKERHUB_USERNAME }}
+ password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+ - name: Build and push ARM64 image
+ uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6.18.0
+ with:
+ context: ${{ matrix.context }}
+ file: ${{ matrix.dockerfile }}
+ platforms: linux/arm64
+ push: true
+ tags: ${{ env.DOCKER_REGISTRY }}/${{ matrix.image_name }}:${{ needs.set-parameters.outputs.image_tag }}-arm64
+ cache-from: type=gha,scope=${{ matrix.image_name }}-arm64
+ cache-to: type=gha,mode=max,scope=${{ matrix.image_name }}-arm64
+ labels: |
+ org.opencontainers.image.title=${{ matrix.image_name }}
+ org.opencontainers.image.description=Apache Texera ${{ matrix.image_name }} (ARM64)
+ org.opencontainers.image.vendor=Apache Texera
+
+ # Step 4: Create multi-arch manifests (only if building both platforms)
+ create-manifests:
+ runs-on: ubuntu-latest
+ needs: [set-parameters, prepare-matrix, build-amd64, build-arm64]
+ if: always() && needs.prepare-matrix.outputs.need_manifest == 'true'
+ strategy:
+ matrix: ${{ fromJson(needs.prepare-matrix.outputs.matrix) }}
+ fail-fast: false
+ env:
+ DOCKER_REGISTRY: ${{ needs.set-parameters.outputs.docker_registry }}
+
+ steps:
+ - name: Log in to GitHub Container Registry
+ if: startsWith(needs.set-parameters.outputs.docker_registry, 'ghcr.io/')
+ uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.6.0
+ with:
+ registry: ghcr.io
+ username: ${{ github.actor }}
+ password: ${{ secrets.GITHUB_TOKEN }}
+
+ - name: Log in to Docker Hub
+ if: ${{ !startsWith(needs.set-parameters.outputs.docker_registry, 'ghcr.io/') }}
+ uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.6.0
+ with:
+ username: ${{ secrets.DOCKERHUB_USERNAME }}
+ password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+ - name: Create and push multi-arch manifest
+ run: |
+ # Create manifest list combining both architectures
+ docker buildx imagetools create -t \
+ ${{ env.DOCKER_REGISTRY }}/${{ matrix.image_name }}:${{ needs.set-parameters.outputs.image_tag }} \
+ ${{ env.DOCKER_REGISTRY }}/${{ matrix.image_name }}:${{ needs.set-parameters.outputs.image_tag }}-amd64 \
+ ${{ env.DOCKER_REGISTRY }}/${{ matrix.image_name }}:${{ needs.set-parameters.outputs.image_tag }}-arm64
+
+ - name: Inspect multi-arch manifest
+ run: |
+ docker buildx imagetools inspect ${{ env.DOCKER_REGISTRY }}/${{ matrix.image_name }}:${{ needs.set-parameters.outputs.image_tag }}
+
+ # Step 5: Summary report
+ build-summary:
+ runs-on: ubuntu-latest
+ needs: [set-parameters, prepare-matrix, build-amd64, build-arm64, create-manifests]
+ if: always()
+ steps:
+ - name: Generate build summary
+ run: |
+ echo "# Texera Multi-Arch Build Complete (Parallel)" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "## Build Configuration" >> $GITHUB_STEP_SUMMARY
+ echo "- **Trigger:** ${{ github.event_name }}" >> $GITHUB_STEP_SUMMARY
+ echo "- **Branch:** \`${{ needs.set-parameters.outputs.branch }}\`" >> $GITHUB_STEP_SUMMARY
+ echo "- **Registry:** \`${{ needs.set-parameters.outputs.docker_registry }}\`" >> $GITHUB_STEP_SUMMARY
+ echo "- **Tag:** \`${{ needs.set-parameters.outputs.image_tag }}\`" >> $GITHUB_STEP_SUMMARY
+ echo "- **Services:** ${{ needs.set-parameters.outputs.services }}" >> $GITHUB_STEP_SUMMARY
+ echo "- **Platforms:** ${{ needs.set-parameters.outputs.platforms }}" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "## Build Method" >> $GITHUB_STEP_SUMMARY
+ echo "**Parallel platform builds** (faster)" >> $GITHUB_STEP_SUMMARY
+ echo "- AMD64: Native build on \`ubuntu-latest\`" >> $GITHUB_STEP_SUMMARY
+ echo "- ARM64: QEMU emulation on \`ubuntu-latest\` (runs in parallel)" >> $GITHUB_STEP_SUMMARY
+ echo "- Manifests: Combined into multi-arch images" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "> **Performance:** AMD64 and ARM64 now build simultaneously instead of sequentially!" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "## Images Published" >> $GITHUB_STEP_SUMMARY
+ echo "All images are now available as multi-arch manifests at:" >> $GITHUB_STEP_SUMMARY
+ echo "\`\`\`" >> $GITHUB_STEP_SUMMARY
+ echo "docker pull ${{ needs.set-parameters.outputs.docker_registry }}/:${{ needs.set-parameters.outputs.image_tag }}" >> $GITHUB_STEP_SUMMARY
+ echo "\`\`\`" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "### Usage" >> $GITHUB_STEP_SUMMARY
+ echo "The images will automatically use the correct architecture:" >> $GITHUB_STEP_SUMMARY
+ echo "- On x86_64/AMD64: pulls linux/amd64 variant" >> $GITHUB_STEP_SUMMARY
+ echo "- On ARM64/M1/M2: pulls linux/arm64 variant" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "### Build Status" >> $GITHUB_STEP_SUMMARY
+ echo "- AMD64 builds: ${{ needs.build-amd64.result }}" >> $GITHUB_STEP_SUMMARY
+ echo "- ARM64 builds: ${{ needs.build-arm64.result }}" >> $GITHUB_STEP_SUMMARY
+ echo "- Manifest creation: ${{ needs.create-manifests.result }}" >> $GITHUB_STEP_SUMMARY
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
new file mode 100644
index 00000000000..4e1b431169f
--- /dev/null
+++ b/.github/workflows/build.yml
@@ -0,0 +1,319 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+name: Build
+
+on:
+ workflow_call:
+ inputs:
+ checkout_ref:
+ required: false
+ type: string
+ default: ""
+ backport_target_branch:
+ required: false
+ type: string
+ default: ""
+ backport_commit_range:
+ required: false
+ type: string
+ default: ""
+ job_name_suffix:
+ required: false
+ type: string
+ default: ""
+ run_frontend:
+ required: false
+ type: boolean
+ default: true
+ run_scala:
+ required: false
+ type: boolean
+ default: true
+ run_python:
+ required: false
+ type: boolean
+ default: true
+ run_agent_service:
+ required: false
+ type: boolean
+ default: true
+
+env:
+ NX_CLOUD_ACCESS_TOKEN: ${{ secrets.NX_CLOUD_ACCESS_TOKEN }}
+
+jobs:
+ frontend:
+ if: ${{ inputs.run_frontend }}
+ name: ${{ format('frontend{0} ({1}, 18)', inputs.job_name_suffix, matrix.os) }}
+ runs-on: ${{ matrix.os }}
+ strategy:
+ matrix:
+ os: [ubuntu-latest, windows-latest, macos-latest]
+ include:
+ - os: macos-latest
+ arch: arm64
+ - os: ubuntu-latest
+ arch: x64
+ - os: windows-latest
+ arch: x64
+ node-version:
+ - 20.19.0
+ steps:
+ - name: Checkout Texera
+ uses: actions/checkout@v5
+ with:
+ ref: ${{ inputs.checkout_ref || github.sha }}
+ fetch-depth: 0
+ - name: Prepare backport workspace
+ if: ${{ inputs.backport_target_branch != '' }}
+ working-directory: ${{ github.workspace }}
+ run: bash ./.github/scripts/prepare-backport-checkout.sh "${{ inputs.backport_target_branch }}" "${{ inputs.backport_commit_range }}"
+ - name: Setup node
+ uses: actions/setup-node@v5
+ with:
+ node-version: ${{ matrix.node-version }}
+ architecture: ${{ matrix.arch }}
+ - uses: actions/cache@v4
+ with:
+ path: frontend/.yarn/cache
+ key: ${{ runner.os }}-${{ matrix.arch }}-${{ matrix.node-version }}-yarn-cache-v4-${{ hashFiles('**/yarn.lock') }}
+ restore-keys: |
+ ${{ runner.os }}-${{ matrix.arch }}-${{ matrix.node-version }}-yarn-cache-v4-
+ - name: Prepare Yarn 4.14.1
+ run: corepack enable && corepack prepare yarn@4.14.1 --activate
+ - name: Setup Python
+ uses: actions/setup-python@v6
+ with:
+ python-version: "3.12"
+ - name: Install dependency
+ timeout-minutes: 20
+ run: yarn --cwd frontend install --immutable --inline-builds --network-timeout=100000
+ - name: Lint with Prettier & ESLint
+ run: yarn --cwd frontend format:ci
+ - name: Prod build
+ run: yarn --cwd frontend run build:ci
+ - name: Check bundled npm packages against LICENSE-binary
+ if: matrix.os == 'ubuntu-latest'
+ run: ./bin/licensing/check_binary_deps.py npm frontend/dist/3rdpartylicenses.json
+ - name: Run frontend unit tests
+ run: yarn --cwd frontend run test:ci
+
+ scala:
+ if: ${{ inputs.run_scala }}
+ name: ${{ format('scala{0} ({1}, 11)', inputs.job_name_suffix, matrix.os) }}
+ strategy:
+ matrix:
+ os: [ubuntu-22.04]
+ java-version: [11]
+ runs-on: ${{ matrix.os }}
+ env:
+ JAVA_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
+ JVM_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
+ services:
+ postgres:
+ image: postgres
+ env:
+ POSTGRES_PASSWORD: postgres
+ ports:
+ - 5432:5432
+ options: >-
+ --health-cmd="pg_isready -U postgres"
+ --health-interval=10s
+ --health-timeout=5s
+ --health-retries=5
+ steps:
+ - name: Checkout
+ uses: actions/checkout@v5
+ with:
+ ref: ${{ inputs.checkout_ref || github.sha }}
+ fetch-depth: 0
+ - name: Prepare backport workspace
+ if: ${{ inputs.backport_target_branch != '' }}
+ working-directory: ${{ github.workspace }}
+ run: bash ./.github/scripts/prepare-backport-checkout.sh "${{ inputs.backport_target_branch }}" "${{ inputs.backport_commit_range }}"
+ - name: Setup JDK
+ uses: actions/setup-java@v5
+ with:
+ distribution: "temurin"
+ java-version: 11
+ - name: Setup Python for Scala tests
+ uses: actions/setup-python@v6
+ with:
+ python-version: "3.11"
+ - name: Show Python
+ run: python --version || python3 --version
+ - name: Install dependencies
+ run: |
+ python -m pip install --upgrade pip
+ if [ -f amber/requirements.txt ]; then pip install -r amber/requirements.txt; fi
+ if [ -f amber/operator-requirements.txt ]; then pip install -r amber/operator-requirements.txt; fi
+ - name: Setup sbt launcher
+ uses: sbt/setup-sbt@508b753e53cb6095967669e0911487d2b9bc9f41 # v1.1.22
+ - uses: coursier/cache-action@90c37294538be80a558fd665531fcdc2b467b475 # v8.1.0
+ with:
+ extraSbtFiles: '["*.sbt", "project/**.{scala,sbt}", "project/build.properties" ]'
+ - name: Lint with scalafmt
+ run: sbt scalafmtCheckAll
+ - name: Create Databases
+ run: |
+ psql -h localhost -U postgres -f sql/texera_ddl.sql
+ psql -h localhost -U postgres -f sql/iceberg_postgres_catalog.sql
+ psql -h localhost -U postgres -f sql/texera_lakefs.sql
+ env:
+ PGPASSWORD: postgres
+ - name: Build distributable bundles for license check
+ # Build every dist-producing module so the union of bundled jars can
+ # be diffed against LICENSE-binary.
+ run: sbt 'clean; ConfigService/dist; AccessControlService/dist; FileService/dist; ComputingUnitManagingService/dist; WorkflowCompilingService/dist; WorkflowExecutionService/dist'
+ - name: Unzip JVM distributable bundles
+ run: |
+ mkdir -p /tmp/dists
+ for zip in \
+ config-service/target/universal/config-service-*.zip \
+ access-control-service/target/universal/access-control-service-*.zip \
+ file-service/target/universal/file-service-*.zip \
+ computing-unit-managing-service/target/universal/computing-unit-managing-service-*.zip \
+ workflow-compiling-service/target/universal/workflow-compiling-service-*.zip \
+ amber/target/universal/amber-*.zip; do
+ unzip -q "$zip" -d /tmp/dists/
+ done
+ - name: Check bundled jars against LICENSE-binary
+ run: |
+ ./bin/licensing/check_binary_deps.py jar \
+ /tmp/dists/config-service-*/lib \
+ /tmp/dists/access-control-service-*/lib \
+ /tmp/dists/file-service-*/lib \
+ /tmp/dists/computing-unit-managing-service-*/lib \
+ /tmp/dists/workflow-compiling-service-*/lib \
+ /tmp/dists/amber-*/lib
+ - name: Audit per-dep license preservation (advisory)
+ if: always()
+ run: |
+ ./bin/licensing/audit_jar_licenses.py \
+ /tmp/dists/config-service-*/lib \
+ /tmp/dists/access-control-service-*/lib \
+ /tmp/dists/file-service-*/lib \
+ /tmp/dists/computing-unit-managing-service-*/lib \
+ /tmp/dists/workflow-compiling-service-*/lib \
+ /tmp/dists/amber-*/lib
+ - name: Create texera_db_for_test_cases
+ run: psql -h localhost -U postgres -v DB_NAME=texera_db_for_test_cases -f sql/texera_ddl.sql
+ env:
+ PGPASSWORD: postgres
+ - name: Compile with sbt
+ run: sbt clean package
+ - name: Lint with scalafix
+ run: sbt "scalafixAll --check"
+ - name: Set docker-java API version
+ run: |
+ echo "api.version=1.52" >> ~/.docker-java.properties
+ cat ~/.docker-java.properties
+ - name: Run backend tests
+ run: sbt test
+
+ python:
+ if: ${{ inputs.run_python }}
+ name: ${{ format('python{0} ({1}, {2})', inputs.job_name_suffix, matrix.os, matrix.python-version) }}
+ strategy:
+ matrix:
+ os: [ubuntu-latest]
+ python-version: ["3.10", "3.11", "3.12", "3.13"]
+ runs-on: ${{ matrix.os }}
+ steps:
+ - name: Checkout Texera
+ uses: actions/checkout@v5
+ with:
+ ref: ${{ inputs.checkout_ref || github.sha }}
+ fetch-depth: 0
+ - name: Prepare backport workspace
+ if: ${{ inputs.backport_target_branch != '' }}
+ run: bash ./.github/scripts/prepare-backport-checkout.sh "${{ inputs.backport_target_branch }}" "${{ inputs.backport_commit_range }}"
+ - name: Set up Python ${{ matrix.python-version }}
+ uses: actions/setup-python@v6
+ with:
+ python-version: ${{ matrix.python-version }}
+ - name: Install dependencies
+ run: |
+ python -m pip install --upgrade pip
+ if [ -f amber/requirements.txt ]; then pip install -r amber/requirements.txt; fi
+ if [ -f amber/operator-requirements.txt ]; then pip install -r amber/operator-requirements.txt; fi
+ if [ "${{ matrix.python-version }}" = "3.12" ]; then pip install pip-licenses; fi
+ - name: Generate pip-licenses manifest
+ if: matrix.python-version == '3.12'
+ run: pip-licenses --format=csv --ignore-packages pip-licenses prettytable wcwidth > /tmp/pip-licenses.csv
+ - name: Check installed Python packages against LICENSE-binary
+ if: matrix.python-version == '3.12'
+ run: ./bin/licensing/check_binary_deps.py python /tmp/pip-licenses.csv
+ - name: Install PostgreSQL
+ run: sudo apt-get update && sudo apt-get install -y postgresql
+ - name: Start PostgreSQL Service
+ run: sudo systemctl start postgresql
+ - name: Create Database and User
+ run: |
+ cd sql && sudo -u postgres psql -f iceberg_postgres_catalog.sql
+ - name: Lint with Ruff
+ run: |
+ cd amber/src/main/python && ruff check . && ruff format --check .
+ - name: Test with pytest
+ run: |
+ cd amber/src/main/python && pytest -sv
+
+ agent-service:
+ if: ${{ inputs.run_agent_service }}
+ name: ${{ format('agent-service{0} ({1})', inputs.job_name_suffix, matrix.os) }}
+ runs-on: ${{ matrix.os }}
+ strategy:
+ fail-fast: false
+ matrix:
+ os: [ubuntu-latest, macos-latest]
+ bun-version: ["1.3.3"]
+ defaults:
+ run:
+ working-directory: agent-service
+ steps:
+ - name: Checkout Texera
+ uses: actions/checkout@v5
+ with:
+ ref: ${{ inputs.checkout_ref || github.sha }}
+ fetch-depth: 0
+ - name: Prepare backport workspace
+ if: ${{ inputs.backport_target_branch != '' }}
+ working-directory: ${{ github.workspace }}
+ run: bash ./.github/scripts/prepare-backport-checkout.sh "${{ inputs.backport_target_branch }}" "${{ inputs.backport_commit_range }}"
+ - name: Setup Bun
+ run: |
+ curl -fsSL https://bun.sh/install | bash -s -- bun-v${{ matrix.bun-version }}
+ echo "$HOME/.bun/bin" >> $GITHUB_PATH
+ - name: Install production dependencies
+ run: bun install --production --frozen-lockfile
+ - name: Generate agent-service license manifest
+ if: matrix.os == 'ubuntu-latest'
+ run: |
+ mkdir -p dist
+ bun run bin/collect-licenses.ts > dist/3rdpartylicenses.json
+ - name: Check bundled agent-service packages against LICENSE-binary
+ if: matrix.os == 'ubuntu-latest'
+ run: ../bin/licensing/check_binary_deps.py agent-npm dist/3rdpartylicenses.json
+ - name: Install development dependencies
+ run: bun install --frozen-lockfile
+ - name: Lint with Prettier
+ run: bun run format:check
+ - name: Typecheck
+ run: bun run typecheck
+ - name: Run unit tests
+ run: bun test
diff --git a/.github/workflows/check-header.yml b/.github/workflows/check-header.yml
new file mode 100644
index 00000000000..c2cbddc1333
--- /dev/null
+++ b/.github/workflows/check-header.yml
@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Release Auditing
+
+on:
+ push:
+ branches:
+ - 'ci-enable/**'
+ - 'main'
+ pull_request:
+ workflow_dispatch:
+
+jobs:
+ test:
+ name: Check License Headers
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v4
+ - uses: apache/skywalking-eyes@5c5b974209f0de5d905f37deb69369068ebfc15c # v0.7.0
diff --git a/.github/workflows/create-release-candidate.yml b/.github/workflows/create-release-candidate.yml
new file mode 100644
index 00000000000..4b0ff4bb3a7
--- /dev/null
+++ b/.github/workflows/create-release-candidate.yml
@@ -0,0 +1,497 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+name: Create and upload release candidate artifacts
+
+on:
+ workflow_dispatch:
+ inputs:
+ tag:
+ description: 'Existing Git tag (e.g., v1.1.0-incubating-rc1)'
+ required: true
+ type: string
+ rc_number:
+ description: 'Release candidate number for artifacts (e.g., 1 for RC1, 2 for RC2)'
+ required: true
+ type: string
+ default: '1'
+ image_registry:
+ description: 'Container image registry prefix (e.g., ghcr.io/apache, docker.io/apache)'
+ required: false
+ type: string
+ default: 'ghcr.io/apache'
+
+jobs:
+ create-rc:
+ runs-on: ubuntu-latest
+ outputs:
+ version: ${{ steps.vars.outputs.version }}
+ rc_num: ${{ steps.vars.outputs.rc_num }}
+ tag_name: ${{ steps.vars.outputs.tag_name }}
+ rc_dir: ${{ steps.vars.outputs.rc_dir }}
+ commit_hash: ${{ steps.vars.outputs.commit_hash }}
+ src_tarball: ${{ steps.vars.outputs.src_tarball }}
+ compose_tarball: ${{ steps.vars.outputs.compose_tarball }}
+ helm_tarball: ${{ steps.vars.outputs.helm_tarball }}
+
+ steps:
+ - name: Checkout code
+ uses: actions/checkout@v4
+ with:
+ fetch-depth: 0 # Full history for proper tagging
+
+ - name: Validate tag exists
+ run: |
+ TAG_NAME="${{ github.event.inputs.tag }}"
+
+ # Check if tag exists
+ if ! git rev-parse "$TAG_NAME" >/dev/null 2>&1; then
+ echo "Error: Tag '$TAG_NAME' does not exist"
+ echo "Available tags:"
+ git tag -l | tail -10
+ exit 1
+ fi
+
+ echo "✓ Tag validation passed: $TAG_NAME"
+
+ - name: Set up variables
+ id: vars
+ run: |
+ TAG_NAME="${{ github.event.inputs.tag }}"
+ RC_NUM="${{ github.event.inputs.rc_number }}"
+ IMAGE_REGISTRY="${{ github.event.inputs.image_registry }}"
+
+ # Parse version from tag (format: v1.1.0-incubating or v1.1.0-incubating-rcN)
+ # Both formats are accepted, but we use the input rc_number for artifacts
+ if [[ "$TAG_NAME" =~ ^v([0-9]+\.[0-9]+\.[0-9]+-incubating)(-rc[0-9]+)?$ ]]; then
+ VERSION="${BASH_REMATCH[1]}"
+ else
+ echo "Error: Tag must be in format vX.Y.Z-incubating or vX.Y.Z-incubating-rcN (e.g., v1.1.0-incubating-rc1)"
+ exit 1
+ fi
+
+ COMMIT_HASH=$(git rev-parse "$TAG_NAME")
+ COMMIT_SHORT=$(git rev-parse "$TAG_NAME" | cut -c1-9)
+ RC_DIR="${VERSION}-RC${RC_NUM}"
+ SRC_TARBALL="apache-texera-${VERSION}-src.tar.gz"
+ COMPOSE_TARBALL="apache-texera-${VERSION}-docker-compose.tar.gz"
+ HELM_TARBALL="apache-texera-${VERSION}-helm.tgz"
+
+ echo "version=$VERSION" >> $GITHUB_OUTPUT
+ echo "rc_num=$RC_NUM" >> $GITHUB_OUTPUT
+ echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT
+ echo "rc_dir=$RC_DIR" >> $GITHUB_OUTPUT
+ echo "commit_hash=$COMMIT_HASH" >> $GITHUB_OUTPUT
+ echo "commit_short=$COMMIT_SHORT" >> $GITHUB_OUTPUT
+ echo "image_registry=$IMAGE_REGISTRY" >> $GITHUB_OUTPUT
+ echo "src_tarball=$SRC_TARBALL" >> $GITHUB_OUTPUT
+ echo "compose_tarball=$COMPOSE_TARBALL" >> $GITHUB_OUTPUT
+ echo "helm_tarball=$HELM_TARBALL" >> $GITHUB_OUTPUT
+
+ echo "Release Candidate: $TAG_NAME"
+ echo "Version: $VERSION"
+ echo "RC Number: $RC_NUM"
+ echo "Commit: $COMMIT_HASH ($COMMIT_SHORT)"
+ echo "Image Registry: $IMAGE_REGISTRY"
+ echo "Staging directory: dist/dev/incubator/texera/$RC_DIR"
+
+ - name: Create source tarball
+ run: |
+ TAG_NAME="${{ steps.vars.outputs.tag_name }}"
+ SRC_TARBALL="${{ steps.vars.outputs.src_tarball }}"
+ VERSION="${{ steps.vars.outputs.version }}"
+
+ TEMP_DIR=$(mktemp -d)
+
+ # Export the git repository at the tag
+ git archive --format=tar --prefix="apache-texera-${VERSION}-src/" "$TAG_NAME" | tar -x -C "$TEMP_DIR"
+
+ # Create tarball
+ cd "$TEMP_DIR"
+ tar -czf "$GITHUB_WORKSPACE/$SRC_TARBALL" "apache-texera-${VERSION}-src"
+
+ cd "$GITHUB_WORKSPACE"
+
+ # Verify tarball was created
+ if [[ ! -f "$SRC_TARBALL" ]]; then
+ echo "Error: Source tarball was not created"
+ exit 1
+ fi
+
+ # Show tarball info
+ ls -lh "$SRC_TARBALL"
+ echo "✓ Created source tarball: $SRC_TARBALL"
+
+ - name: Create Docker Compose deployment bundle
+ run: |
+ VERSION="${{ steps.vars.outputs.version }}"
+ TAG_NAME="${{ steps.vars.outputs.tag_name }}"
+ IMAGE_REGISTRY="${{ steps.vars.outputs.image_registry }}"
+ COMMIT_SHORT="${{ steps.vars.outputs.commit_short }}"
+ COMPOSE_TARBALL="${{ steps.vars.outputs.compose_tarball }}"
+
+ TEMP_DIR=$(mktemp -d)
+ BUNDLE_DIR="$TEMP_DIR/apache-texera-${VERSION}-docker-compose"
+ mkdir -p "$BUNDLE_DIR"
+
+ # Export the single-node directory from the tagged source
+ mkdir -p "$TEMP_DIR/_raw"
+ git archive --format=tar "$TAG_NAME" -- bin/single-node/ sql/ | tar -x -C "$TEMP_DIR/_raw"
+
+ # Copy deployment files
+ cp "$TEMP_DIR/_raw/bin/single-node/docker-compose.yml" "$BUNDLE_DIR/"
+ cp "$TEMP_DIR/_raw/bin/single-node/nginx.conf" "$BUNDLE_DIR/"
+ cp -r "$TEMP_DIR/_raw/sql" "$BUNDLE_DIR/"
+
+ # Patch the SQL mount path for the self-contained bundle layout
+ # In the repo it's ../../sql (relative to bin/single-node/), in the bundle it's ./sql
+ sed -i 's|../../sql|./sql|g' "$BUNDLE_DIR/docker-compose.yml"
+
+ # Generate a release-pinned .env file with the version tag
+ # Start from the source .env and ensure IMAGE_REGISTRY and IMAGE_TAG are set
+ cp "$TEMP_DIR/_raw/bin/single-node/.env" "$BUNDLE_DIR/.env"
+ # Replace if line exists, otherwise append
+ if grep -q '^IMAGE_REGISTRY=' "$BUNDLE_DIR/.env"; then
+ sed -i "s|^IMAGE_REGISTRY=.*|IMAGE_REGISTRY=${IMAGE_REGISTRY}|" "$BUNDLE_DIR/.env"
+ else
+ echo "IMAGE_REGISTRY=${IMAGE_REGISTRY}" >> "$BUNDLE_DIR/.env"
+ fi
+ if grep -q '^IMAGE_TAG=' "$BUNDLE_DIR/.env"; then
+ sed -i "s|^IMAGE_TAG=.*|IMAGE_TAG=${COMMIT_SHORT}|" "$BUNDLE_DIR/.env"
+ else
+ echo "IMAGE_TAG=${COMMIT_SHORT}" >> "$BUNDLE_DIR/.env"
+ fi
+ if grep -q '^TEXERA_SERVICE_LOG_LEVEL=' "$BUNDLE_DIR/.env"; then
+ sed -i "s|^TEXERA_SERVICE_LOG_LEVEL=.*|TEXERA_SERVICE_LOG_LEVEL=ERROR|" "$BUNDLE_DIR/.env"
+ else
+ echo "TEXERA_SERVICE_LOG_LEVEL=ERROR" >> "$BUNDLE_DIR/.env"
+ fi
+
+ # Include the README from the repo
+ cp "$TEMP_DIR/_raw/bin/single-node/README.md" "$BUNDLE_DIR/"
+
+ # Include example datasets, workflows, and the loader script
+ if [ -d "$TEMP_DIR/_raw/bin/single-node/examples" ]; then
+ cp -r "$TEMP_DIR/_raw/bin/single-node/examples" "$BUNDLE_DIR/"
+ echo "✓ Included examples directory (datasets, workflows, load-examples.sh)"
+ fi
+
+ # Create tarball
+ cd "$TEMP_DIR"
+ tar -czf "$GITHUB_WORKSPACE/$COMPOSE_TARBALL" "apache-texera-${VERSION}-docker-compose"
+
+ cd "$GITHUB_WORKSPACE"
+ ls -lh "$COMPOSE_TARBALL"
+ echo "✓ Created Docker Compose bundle: $COMPOSE_TARBALL"
+
+ - name: Create Helm chart package
+ run: |
+ VERSION="${{ steps.vars.outputs.version }}"
+ TAG_NAME="${{ steps.vars.outputs.tag_name }}"
+ COMMIT_SHORT="${{ steps.vars.outputs.commit_short }}"
+ IMAGE_REGISTRY="${{ steps.vars.outputs.image_registry }}"
+ HELM_TARBALL="${{ steps.vars.outputs.helm_tarball }}"
+
+ # Install Helm
+ curl https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash
+
+ # Export the Helm chart from the tagged source
+ # Include sql/ because bin/k8s/files/ has symlinks to ../../../sql/
+ TEMP_DIR=$(mktemp -d)
+ git archive --format=tar "$TAG_NAME" -- bin/k8s/ sql/ | tar -x -C "$TEMP_DIR"
+
+ CHART_DIR="$TEMP_DIR/bin/k8s"
+
+ # Resolve symlinks in the chart so it packages as a self-contained artifact
+ find "$CHART_DIR" -type l | while read -r link; do
+ target=$(readlink -f "$link")
+ rm "$link"
+ cp "$target" "$link"
+ done
+
+ # Update Chart.yaml with release version
+ sed -i "s/^version:.*/version: ${VERSION}/" "$CHART_DIR/Chart.yaml"
+ sed -i "s/^appVersion:.*/appVersion: \"${VERSION}\"/" "$CHART_DIR/Chart.yaml"
+
+ # Update values.yaml with release image registry and tag
+ sed -i "s|imageRegistry:.*|imageRegistry: ${IMAGE_REGISTRY}|" "$CHART_DIR/values.yaml"
+ sed -i "s|imageTag:.*|imageTag: ${COMMIT_SHORT}|" "$CHART_DIR/values.yaml"
+
+ # Download chart dependencies declared in Chart.yaml
+ # These tarballs are .gitignored so they're not in git archive
+ helm dependency build "$CHART_DIR"
+
+ # Package the Helm chart
+ helm package "$CHART_DIR" \
+ --version "$VERSION" \
+ --app-version "$VERSION" \
+ --destination "$GITHUB_WORKSPACE"
+
+ # Rename to our expected artifact name
+ HELM_PKG=$(ls "$GITHUB_WORKSPACE"/texera-helm-*.tgz 2>/dev/null | head -1)
+ if [[ -n "$HELM_PKG" && "$HELM_PKG" != "$GITHUB_WORKSPACE/$HELM_TARBALL" ]]; then
+ mv "$HELM_PKG" "$GITHUB_WORKSPACE/$HELM_TARBALL"
+ fi
+
+ ls -lh "$GITHUB_WORKSPACE/$HELM_TARBALL"
+ echo "✓ Created Helm chart package: $HELM_TARBALL"
+
+ - name: Import GPG key
+ run: |
+ echo "${{ secrets.GPG_PRIVATE_KEY }}" | gpg --batch --import
+
+ # List imported keys
+ gpg --list-secret-keys
+
+ echo "✓ GPG key imported successfully"
+
+ - name: Sign and checksum all artifacts
+ run: |
+ for artifact in \
+ "${{ steps.vars.outputs.src_tarball }}" \
+ "${{ steps.vars.outputs.compose_tarball }}" \
+ "${{ steps.vars.outputs.helm_tarball }}"; do
+
+ # GPG signature
+ echo "${{ secrets.GPG_PASSPHRASE }}" | gpg --batch --yes --pinentry-mode loopback --passphrase-fd 0 \
+ --armor --detach-sign --output "${artifact}.asc" "$artifact"
+ gpg --verify "${artifact}.asc" "$artifact"
+ echo "✓ Signed: ${artifact}"
+
+ # SHA512 checksum
+ sha512sum "$artifact" > "${artifact}.sha512"
+ echo "✓ Checksum: ${artifact}.sha512"
+ done
+
+ - name: Generate vote email template
+ id: vote_email
+ run: |
+ VERSION="${{ steps.vars.outputs.version }}"
+ RC_NUM="${{ steps.vars.outputs.rc_num }}"
+ TAG_NAME="${{ steps.vars.outputs.tag_name }}"
+ RC_DIR="${{ steps.vars.outputs.rc_dir }}"
+ COMMIT_HASH="${{ steps.vars.outputs.commit_hash }}"
+ IMAGE_REGISTRY="${{ steps.vars.outputs.image_registry }}"
+
+ # Get GPG key ID from the imported key
+ GPG_KEY_ID=$(gpg --list-secret-keys --keyid-format LONG | grep 'sec' | head -n1 | awk '{print $2}' | cut -d'/' -f2)
+ GPG_EMAIL=$(gpg --list-secret-keys | grep 'uid' | head -n1 | grep -oP '[\w\.-]+@[\w\.-]+')
+
+ # Copy template from repository
+ cp .github/release/vote-email-template.md vote-email.txt
+
+ # Substitute variables in the template
+ sed -i "s|\${VERSION}|${VERSION}|g" vote-email.txt
+ sed -i "s|\${RC_NUM}|${RC_NUM}|g" vote-email.txt
+ sed -i "s|\${RC_DIR}|${RC_DIR}|g" vote-email.txt
+ sed -i "s|\${TAG_NAME}|${TAG_NAME}|g" vote-email.txt
+ sed -i "s|\${COMMIT_HASH}|${COMMIT_HASH}|g" vote-email.txt
+ sed -i "s|\${GPG_KEY_ID}|${GPG_KEY_ID}|g" vote-email.txt
+ sed -i "s|\${GPG_EMAIL}|${GPG_EMAIL}|g" vote-email.txt
+ sed -i "s|\${IMAGE_REGISTRY}|${IMAGE_REGISTRY}|g" vote-email.txt
+
+ echo "✓ Vote email template generated!"
+
+ - name: Upload RC artifacts
+ uses: actions/upload-artifact@v4
+ with:
+ name: rc-artifacts
+ path: |
+ ${{ steps.vars.outputs.src_tarball }}
+ ${{ steps.vars.outputs.src_tarball }}.asc
+ ${{ steps.vars.outputs.src_tarball }}.sha512
+ ${{ steps.vars.outputs.compose_tarball }}
+ ${{ steps.vars.outputs.compose_tarball }}.asc
+ ${{ steps.vars.outputs.compose_tarball }}.sha512
+ ${{ steps.vars.outputs.helm_tarball }}
+ ${{ steps.vars.outputs.helm_tarball }}.asc
+ ${{ steps.vars.outputs.helm_tarball }}.sha512
+ vote-email.txt
+ retention-days: 7
+
+ upload-rc:
+ runs-on: ubuntu-latest
+ needs: create-rc
+
+ steps:
+ - name: Download RC artifacts
+ uses: actions/download-artifact@v4
+ with:
+ name: rc-artifacts
+
+ - name: Verify downloaded artifacts
+ run: |
+ SRC_TARBALL="${{ needs.create-rc.outputs.src_tarball }}"
+ COMPOSE_TARBALL="${{ needs.create-rc.outputs.compose_tarball }}"
+ HELM_TARBALL="${{ needs.create-rc.outputs.helm_tarball }}"
+
+ echo "Verifying downloaded artifacts..."
+ ls -lh
+
+ for artifact in "$SRC_TARBALL" "$COMPOSE_TARBALL" "$HELM_TARBALL"; do
+ if [[ ! -f "$artifact" ]] || [[ ! -f "${artifact}.asc" ]] || [[ ! -f "${artifact}.sha512" ]]; then
+ echo "Error: Missing artifact or signature/checksum for: $artifact"
+ exit 1
+ fi
+ done
+
+ echo "✓ All artifacts downloaded successfully"
+
+ - name: Install SVN
+ run: |
+ sudo apt-get update
+ sudo apt-get install -y subversion
+ svn --version
+
+ - name: Checkout SVN dev directory
+ run: |
+ RC_DIR="${{ needs.create-rc.outputs.rc_dir }}"
+
+ # Checkout the dev directory with depth=empty (lightweight)
+ svn co --depth=empty https://dist.apache.org/repos/dist/dev/incubator/texera svn-texera \
+ --username "${{ secrets.SVN_USERNAME }}" \
+ --password "${{ secrets.SVN_PASSWORD }}" \
+ --no-auth-cache
+
+ cd svn-texera
+
+ # Check if RC directory already exists on the remote
+ SVN_BASE="https://dist.apache.org/repos/dist/dev/incubator/texera"
+ if svn info "$SVN_BASE/$RC_DIR" >/dev/null 2>&1; then
+ # Directory exists remotely — update (checkout) it into the working copy
+ svn update --depth=infinity "$RC_DIR" \
+ --username "${{ secrets.SVN_USERNAME }}" \
+ --password "${{ secrets.SVN_PASSWORD }}" \
+ --no-auth-cache || true
+ # If update didn't bring it down (empty parent checkout), do a sparse checkout
+ if [[ ! -d "$RC_DIR" ]]; then
+ svn update --set-depth=infinity "$RC_DIR" \
+ --username "${{ secrets.SVN_USERNAME }}" \
+ --password "${{ secrets.SVN_PASSWORD }}" \
+ --no-auth-cache
+ fi
+ echo "✓ RC directory already exists remotely, checked out: $RC_DIR"
+ else
+ # Directory doesn't exist remotely — create and add it
+ mkdir -p "$RC_DIR"
+ svn add "$RC_DIR"
+ echo "✓ Created new RC directory: $RC_DIR"
+ fi
+
+ - name: Stage artifacts to SVN
+ run: |
+ SRC_TARBALL="${{ needs.create-rc.outputs.src_tarball }}"
+ COMPOSE_TARBALL="${{ needs.create-rc.outputs.compose_tarball }}"
+ HELM_TARBALL="${{ needs.create-rc.outputs.helm_tarball }}"
+ RC_DIR="${{ needs.create-rc.outputs.rc_dir }}"
+
+ cd svn-texera/"$RC_DIR"
+
+ # Copy all artifacts
+ for artifact in "$SRC_TARBALL" "$COMPOSE_TARBALL" "$HELM_TARBALL"; do
+ cp "$GITHUB_WORKSPACE/$artifact" .
+ cp "$GITHUB_WORKSPACE/${artifact}.asc" .
+ cp "$GITHUB_WORKSPACE/${artifact}.sha512" .
+ done
+
+ # Add files to SVN
+ svn add * --force
+
+ # Check status
+ svn status
+
+ echo "✓ Staged all artifacts to SVN"
+
+ - name: Commit artifacts to dist/dev
+ run: |
+ VERSION="${{ needs.create-rc.outputs.version }}"
+ RC_NUM="${{ needs.create-rc.outputs.rc_num }}"
+ RC_DIR="${{ needs.create-rc.outputs.rc_dir }}"
+
+ cd svn-texera
+
+ # Commit with descriptive message
+ svn commit -m "Add Apache Texera ${VERSION} RC${RC_NUM} artifacts (source + docker-compose + helm)" \
+ --username "${{ secrets.SVN_USERNAME }}" \
+ --password "${{ secrets.SVN_PASSWORD }}" \
+ --no-auth-cache
+
+ echo "✓ Committed artifacts to dist/dev/incubator/texera/$RC_DIR"
+
+ - name: Generate release summary
+ run: |
+ VERSION="${{ needs.create-rc.outputs.version }}"
+ RC_NUM="${{ needs.create-rc.outputs.rc_num }}"
+ TAG_NAME="${{ needs.create-rc.outputs.tag_name }}"
+ RC_DIR="${{ needs.create-rc.outputs.rc_dir }}"
+ COMMIT_HASH="${{ needs.create-rc.outputs.commit_hash }}"
+ SRC_TARBALL="${{ needs.create-rc.outputs.src_tarball }}"
+ COMPOSE_TARBALL="${{ needs.create-rc.outputs.compose_tarball }}"
+ HELM_TARBALL="${{ needs.create-rc.outputs.helm_tarball }}"
+
+ echo "## Release Candidate Created Successfully!" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "### Release Information" >> $GITHUB_STEP_SUMMARY
+ echo "- **Version:** ${VERSION}" >> $GITHUB_STEP_SUMMARY
+ echo "- **RC Number:** RC${RC_NUM}" >> $GITHUB_STEP_SUMMARY
+ echo "- **Git Tag:** \`${TAG_NAME}\`" >> $GITHUB_STEP_SUMMARY
+ echo "- **Commit:** \`${COMMIT_HASH}\`" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "### Artifacts Location" >> $GITHUB_STEP_SUMMARY
+ echo "**Staging Directory:** https://dist.apache.org/repos/dist/dev/incubator/texera/${RC_DIR}/" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "### Artifacts Created" >> $GITHUB_STEP_SUMMARY
+ echo "| Artifact | Description |" >> $GITHUB_STEP_SUMMARY
+ echo "|----------|-------------|" >> $GITHUB_STEP_SUMMARY
+ echo "| \`${SRC_TARBALL}\` | Source code |" >> $GITHUB_STEP_SUMMARY
+ echo "| \`${COMPOSE_TARBALL}\` | Docker Compose deployment bundle |" >> $GITHUB_STEP_SUMMARY
+ echo "| \`${HELM_TARBALL}\` | Helm chart for Kubernetes deployment |" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "Each artifact has a corresponding \`.asc\` (GPG signature) and \`.sha512\` (checksum) file." >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "### Next Steps" >> $GITHUB_STEP_SUMMARY
+ echo "1. Build and push container images using the \`Build and push images\` workflow with tag \`${VERSION}\`" >> $GITHUB_STEP_SUMMARY
+ echo "2. Verify the artifacts at the staging directory" >> $GITHUB_STEP_SUMMARY
+ echo "3. Send [VOTE] email to dev@texera.apache.org" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "### Verification" >> $GITHUB_STEP_SUMMARY
+ echo "\`\`\`bash" >> $GITHUB_STEP_SUMMARY
+ echo "# Import KEYS and verify signatures" >> $GITHUB_STEP_SUMMARY
+ echo "gpg --import KEYS" >> $GITHUB_STEP_SUMMARY
+ echo "gpg --verify ${SRC_TARBALL}.asc ${SRC_TARBALL}" >> $GITHUB_STEP_SUMMARY
+ echo "gpg --verify ${COMPOSE_TARBALL}.asc ${COMPOSE_TARBALL}" >> $GITHUB_STEP_SUMMARY
+ echo "gpg --verify ${HELM_TARBALL}.asc ${HELM_TARBALL}" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "# Verify SHA512 checksums" >> $GITHUB_STEP_SUMMARY
+ echo "sha512sum -c ${SRC_TARBALL}.sha512" >> $GITHUB_STEP_SUMMARY
+ echo "sha512sum -c ${COMPOSE_TARBALL}.sha512" >> $GITHUB_STEP_SUMMARY
+ echo "sha512sum -c ${HELM_TARBALL}.sha512" >> $GITHUB_STEP_SUMMARY
+ echo "\`\`\`" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "**KEYS file:** https://downloads.apache.org/incubator/texera/KEYS" >> $GITHUB_STEP_SUMMARY
+
+ echo "✓ Release candidate workflow completed successfully!"
+
+ - name: Display vote email template
+ run: |
+ echo "## Vote Email Template" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "Copy the content below to send to dev@texera.apache.org:" >> $GITHUB_STEP_SUMMARY
+ echo "" >> $GITHUB_STEP_SUMMARY
+ echo "\`\`\`" >> $GITHUB_STEP_SUMMARY
+ cat "$GITHUB_WORKSPACE/vote-email.txt" >> $GITHUB_STEP_SUMMARY
+ echo "\`\`\`" >> $GITHUB_STEP_SUMMARY
diff --git a/.github/workflows/direct-backport-push.yml b/.github/workflows/direct-backport-push.yml
new file mode 100644
index 00000000000..c291fc18879
--- /dev/null
+++ b/.github/workflows/direct-backport-push.yml
@@ -0,0 +1,216 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Direct Backport Push
+
+on:
+ push:
+ branches:
+ - main
+
+permissions:
+ actions: read
+ contents: write
+ pull-requests: read
+
+jobs:
+ discover:
+ name: Discover direct backport targets
+ runs-on: ubuntu-latest
+ outputs:
+ pr_number: ${{ steps.discover.outputs.pr_number }}
+ targets: ${{ steps.discover.outputs.targets }}
+ has_targets: ${{ steps.discover.outputs.has_targets }}
+ steps:
+ - name: Resolve merged PR and green targets
+ id: discover
+ uses: actions/github-script@v8
+ with:
+ script: |
+ const sha = context.sha;
+ const { owner, repo } = context.repo;
+
+ // Strategy 1 (preferred): parse the squash-merge commit message.
+ // ASF .asf.yaml forces squash merges with PR_TITLE_AND_DESC, so the
+ // first line ends with "(#NNNN)". This is deterministic and avoids
+ // the commit↔PR association index, which can lag for tens of seconds
+ // after a merge.
+ async function resolvePrFromMessage() {
+ const message = context.payload?.head_commit?.message ?? "";
+ const firstLine = message.split("\n", 1)[0];
+ const match = firstLine.match(/\(#(\d+)\)\s*$/);
+ if (!match) {
+ core.info('Commit message does not end with "(#N)"; falling back to API.');
+ return null;
+ }
+ const prNumber = Number(match[1]);
+ try {
+ const { data: pr } = await github.rest.pulls.get({
+ owner,
+ repo,
+ pull_number: prNumber,
+ });
+ if (!pr.merged) {
+ core.warning(`PR #${prNumber} extracted from commit message is not merged; falling back to API.`);
+ return null;
+ }
+ core.info(`Resolved PR #${prNumber} from commit message.`);
+ return pr;
+ } catch (e) {
+ core.warning(`Failed to fetch PR #${prNumber}: ${e.message}. Falling back to API.`);
+ return null;
+ }
+ }
+
+ // Strategy 2 (fallback): GET /commits/{sha}/pulls with exponential
+ // backoff. 5 attempts at 0/2/4/8/16s — total worst case ~30s.
+ async function resolvePrFromApi() {
+ const backoffsMs = [0, 2000, 4000, 8000, 16000];
+ for (let i = 0; i < backoffsMs.length; i++) {
+ if (backoffsMs[i] > 0) {
+ core.info(`Retrying commit→PR lookup in ${backoffsMs[i] / 1000}s (attempt ${i + 1}/${backoffsMs.length}).`);
+ await new Promise((resolve) => setTimeout(resolve, backoffsMs[i]));
+ }
+ const response = await github.request(
+ "GET /repos/{owner}/{repo}/commits/{commit_sha}/pulls",
+ {
+ owner,
+ repo,
+ commit_sha: sha,
+ }
+ );
+ const pr = response.data.find((p) => p.merge_commit_sha === sha) ?? response.data[0];
+ if (pr) {
+ core.info(`Resolved PR #${pr.number} from commits/${sha}/pulls on attempt ${i + 1}.`);
+ return pr;
+ }
+ }
+ return null;
+ }
+
+ const pullRequest = (await resolvePrFromMessage()) ?? (await resolvePrFromApi());
+ if (!pullRequest) {
+ core.info(`No merged pull request is associated with ${sha}.`);
+ core.setOutput("pr_number", "");
+ core.setOutput("targets", "[]");
+ core.setOutput("has_targets", "false");
+ return;
+ }
+
+ const requestedTargets = [...new Set(
+ pullRequest.labels
+ .map((label) => label.name)
+ .filter((name) => /^release\/.+$/.test(name))
+ )].sort();
+
+ if (requestedTargets.length === 0) {
+ core.info(`PR #${pullRequest.number} does not request any backports.`);
+ core.setOutput("pr_number", String(pullRequest.number));
+ core.setOutput("targets", "[]");
+ core.setOutput("has_targets", "false");
+ return;
+ }
+
+ const buildRuns = await github.paginate(
+ github.rest.actions.listWorkflowRuns,
+ {
+ owner,
+ repo,
+ workflow_id: "required-checks.yml",
+ head_sha: pullRequest.head.sha,
+ per_page: 100,
+ }
+ );
+
+ let greenTargets = [];
+ if (buildRuns.length === 0) {
+ core.warning(`No Build workflow runs found for ${pullRequest.head.sha}.`);
+ } else {
+ const allJobs = [];
+ for (const run of buildRuns) {
+ const jobs = await github.paginate(
+ github.rest.actions.listJobsForWorkflowRun,
+ {
+ owner,
+ repo,
+ run_id: run.id,
+ per_page: 100,
+ }
+ );
+ allJobs.push(...jobs);
+ }
+
+ greenTargets = requestedTargets.filter((target) => {
+ const prefix = `backport (${target}) / `;
+ const targetJobs = allJobs.filter((job) => job.name.startsWith(prefix));
+ return targetJobs.length > 0 && targetJobs.every((job) => job.conclusion === "success");
+ });
+ }
+
+ const skippedTargets = requestedTargets.filter((target) => !greenTargets.includes(target));
+ if (skippedTargets.length > 0) {
+ core.warning(`Skipping targets without a successful Backport run: ${skippedTargets.join(", ")}`);
+ }
+
+ core.setOutput("pr_number", String(pullRequest.number));
+ core.setOutput("targets", JSON.stringify(greenTargets));
+ core.setOutput("has_targets", greenTargets.length > 0 ? "true" : "false");
+
+ push-backports:
+ needs: discover
+ if: ${{ needs.discover.outputs.has_targets == 'true' }}
+ runs-on: ubuntu-latest
+ strategy:
+ fail-fast: false
+ matrix:
+ target: ${{ fromJson(needs.discover.outputs.targets) }}
+ steps:
+ - name: Checkout main
+ uses: actions/checkout@v5
+ with:
+ fetch-depth: 0
+ - name: Cherry-pick merge commit onto target branch
+ env:
+ MERGE_SHA: ${{ github.sha }}
+ TARGET_BRANCH: ${{ matrix.target }}
+ run: |
+ set -euo pipefail
+
+ parent_count=$(git rev-list --parents -n 1 "${MERGE_SHA}" | awk '{print NF-1}')
+ if [[ "${parent_count}" -ne 1 ]]; then
+ echo "Direct backport expects a squash-merged commit on main. ${MERGE_SHA} has ${parent_count} parents." >&2
+ exit 1
+ fi
+
+ git config user.name "github-actions[bot]"
+ git config user.email "github-actions[bot]@users.noreply.github.com"
+
+ # Reuse the squash commit's full message so the PR title, description,
+ # and any Co-Authored-By trailers GitHub injected stay attached to
+ # the backport commit. The author is the squash commit's author
+ # (the original PR author for squash merges).
+ original_author=$(git log -1 --format='%an <%ae>' "${MERGE_SHA}")
+ merge_message=$(git log -1 --format=%B "${MERGE_SHA}")
+
+ git fetch --no-tags origin "${TARGET_BRANCH}"
+ git checkout -B "${TARGET_BRANCH}" "origin/${TARGET_BRANCH}"
+ git cherry-pick --no-commit "${MERGE_SHA}"
+
+ {
+ printf '%s\n\n(backported from commit %s)\n' "${merge_message}" "${MERGE_SHA}"
+ } | git commit -F - --author="${original_author}"
+
+ git push origin "HEAD:${TARGET_BRANCH}"
diff --git a/.github/workflows/issue-triage.yml b/.github/workflows/issue-triage.yml
new file mode 100644
index 00000000000..0846b98ebe8
--- /dev/null
+++ b/.github/workflows/issue-triage.yml
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Issue triage
+on:
+ issues:
+ types: [assigned, unassigned]
+
+permissions:
+ issues: write
+ pull-requests: read
+
+jobs:
+ # --------------------------------------------------------
+ # 1) Issue triage: add/remove "triage" on assignment changes
+ # --------------------------------------------------------
+ issue-triage:
+ if: github.event_name == 'issues'
+ runs-on: ubuntu-latest
+ steps:
+ - name: Remove 'triage' label when issue is assigned
+ if: github.event.action == 'assigned'
+ uses: actions/github-script@v7
+ with:
+ github-token: ${{ secrets.GITHUB_TOKEN }}
+ script: |
+ const { owner, repo } = context.repo;
+ const issue_number = context.payload.issue.number;
+
+ try {
+ await github.rest.issues.removeLabel({
+ owner,
+ repo,
+ issue_number,
+ name: 'triage',
+ });
+ core.info(`Removed 'triage' from issue #${issue_number}`);
+ } catch (e) {
+ if (e.status === 404) {
+ core.info(`Issue #${issue_number} has no 'triage' label, nothing to remove.`);
+ } else {
+ core.warning(`Failed to remove 'triage' from issue #${issue_number}: ${e.message}`);
+ throw e;
+ }
+ }
+
+ - name: Add 'triage' label when issue is unassigned and has no assignees
+ if: github.event.action == 'unassigned'
+ uses: actions/github-script@v7
+ with:
+ github-token: ${{ secrets.GITHUB_TOKEN }}
+ script: |
+ const issue = context.payload.issue;
+ const { owner, repo } = context.repo;
+ const issue_number = issue.number;
+
+ const assignees = issue.assignees || [];
+ if (assignees.length > 0) {
+ core.info(
+ `Issue #${issue_number} still has ${assignees.length} assignee(s), not adding 'triage'.`
+ );
+ return;
+ }
+
+ core.info(`Issue #${issue_number} has no assignees, adding 'triage' label.`);
+ try {
+ await github.rest.issues.addLabels({
+ owner,
+ repo,
+ issue_number,
+ labels: ['triage'],
+ });
+ } catch (e) {
+ core.warning(`Failed to add 'triage' to issue #${issue_number}: ${e.message}`);
+ throw e;
+ }
diff --git a/.github/workflows/lint-pr.yml b/.github/workflows/lint-pr.yml
new file mode 100644
index 00000000000..fcc5d46f0ca
--- /dev/null
+++ b/.github/workflows/lint-pr.yml
@@ -0,0 +1,36 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Lint PR
+
+on:
+ pull_request_target:
+ types:
+ - opened
+ - edited
+ - reopened
+ - synchronize
+
+jobs:
+ main:
+ name: Validate PR title
+ runs-on: ubuntu-latest
+ permissions:
+ pull-requests: read
+ steps:
+ - uses: amannn/action-semantic-pull-request@0723387faaf9b38adef4775cd42cfd5155ed6017 # v5.5.3
+ env:
+ GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/workflows/pr-labeler.yml b/.github/workflows/pr-labeler.yml
new file mode 100644
index 00000000000..d074cc9b5ae
--- /dev/null
+++ b/.github/workflows/pr-labeler.yml
@@ -0,0 +1,31 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: "Pull Request Labeler"
+
+on:
+ - pull_request_target
+
+jobs:
+ labeler:
+ permissions:
+ contents: read
+ pull-requests: write
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/labeler@v6
+ with:
+ sync-labels: true
diff --git a/.github/workflows/required-checks.yml b/.github/workflows/required-checks.yml
new file mode 100644
index 00000000000..2f858cd951d
--- /dev/null
+++ b/.github/workflows/required-checks.yml
@@ -0,0 +1,232 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+name: Required Checks
+
+on:
+ push:
+ branches:
+ - 'ci-enable/**'
+ - 'main'
+ - 'release/**'
+ pull_request:
+ types:
+ - opened
+ - reopened
+ - synchronize
+ - labeled
+ - unlabeled
+ workflow_dispatch:
+
+permissions:
+ checks: write
+ contents: read
+ pull-requests: read
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.ref }}
+ cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+jobs:
+ # Precheck decides which downstream jobs run for this event:
+ # - On PR events, wait for the Pull Request Labeler workflow to finish so
+ # the labels it applies (frontend, docs, dev, ...) are available, then
+ # gate run_* outputs on those labels.
+ # - run_frontend / run_scala / run_python / run_agent_service: gate the
+ # main build stacks. PRs labelled exclusively with docs and/or dev skip
+ # every stack; otherwise frontend skips when no `frontend` label is
+ # present (the other stacks always run when at least one non-docs/dev
+ # label exists). Push and dispatch events run every stack.
+ # - backport_targets: JSON array of release/* labels currently on the PR.
+ # Drives the backport matrix; empty array means no backport runs.
+ precheck:
+ name: Precheck
+ runs-on: ubuntu-latest
+ outputs:
+ run_frontend: ${{ steps.decide.outputs.run_frontend }}
+ run_scala: ${{ steps.decide.outputs.run_scala }}
+ run_python: ${{ steps.decide.outputs.run_python }}
+ run_agent_service: ${{ steps.decide.outputs.run_agent_service }}
+ backport_targets: ${{ steps.decide.outputs.backport_targets }}
+ steps:
+ - name: Wait for Pull Request Labeler
+ if: github.event_name == 'pull_request'
+ uses: actions/github-script@v8
+ with:
+ script: |
+ const ref = context.payload.pull_request.head.sha;
+ const maxAttempts = 30;
+ for (let i = 0; i < maxAttempts; i++) {
+ const { data } = await github.rest.checks.listForRef({
+ owner: context.repo.owner,
+ repo: context.repo.repo,
+ ref,
+ check_name: "labeler",
+ });
+ const check = data.check_runs[0];
+ if (check && check.status === "completed") {
+ core.info(`labeler ${check.conclusion}`);
+ return;
+ }
+ core.info(`labeler not ready (attempt ${i + 1}/${maxAttempts})`);
+ await new Promise((r) => setTimeout(r, 10000));
+ }
+ core.warning("labeler did not complete within 5 minutes; proceeding with current labels.");
+
+ - name: Decide which jobs to run
+ id: decide
+ uses: actions/github-script@v8
+ with:
+ script: |
+ const eventName = context.eventName;
+ let labels = [];
+
+ if (eventName === "pull_request") {
+ // Re-fetch labels: the labeler may have just added some.
+ const { data: pr } = await github.rest.pulls.get({
+ owner: context.repo.owner,
+ repo: context.repo.repo,
+ pull_number: context.payload.pull_request.number,
+ });
+ labels = pr.labels.map((l) => l.name);
+ core.info(`PR labels: ${labels.join(", ") || "(none)"}`);
+ }
+
+ const SKIP_ONLY = new Set(["docs", "dev"]);
+ const onlySkippable =
+ eventName === "pull_request" &&
+ labels.length > 0 &&
+ labels.every((l) => SKIP_ONLY.has(l));
+
+ let runFrontend = true;
+ let runScala = true;
+ let runPython = true;
+ let runAgentService = true;
+
+ if (onlySkippable) {
+ runFrontend = runScala = runPython = runAgentService = false;
+ core.info("Labels are docs/dev only; skipping all build stacks.");
+ } else if (eventName === "pull_request" && !labels.includes("frontend")) {
+ runFrontend = false;
+ core.info("No frontend label; skipping frontend stack.");
+ }
+
+ core.setOutput("run_frontend", runFrontend ? "true" : "false");
+ core.setOutput("run_scala", runScala ? "true" : "false");
+ core.setOutput("run_python", runPython ? "true" : "false");
+ core.setOutput("run_agent_service", runAgentService ? "true" : "false");
+
+ // Backport targets: all current release/* labels on the PR.
+ const targets = [...new Set(labels.filter((n) => /^release\/.+$/.test(n)))].sort();
+ if (targets.length === 0) {
+ core.info("No backport targets on PR.");
+ } else {
+ core.info(`Backport targets: ${targets.join(", ")}`);
+ }
+ core.setOutput("backport_targets", JSON.stringify(targets));
+
+ cleanup-stale-backport:
+ if: ${{ github.event_name == 'pull_request' && github.event.action == 'unlabeled' && startsWith(github.event.label.name, 'release/') }}
+ runs-on: ubuntu-latest
+ steps:
+ - name: Cancel obsolete backport check_runs for the removed target
+ uses: actions/github-script@v8
+ with:
+ script: |
+ const { owner, repo } = context.repo;
+ const target = context.payload.label.name;
+ const headSha = context.payload.pull_request.head.sha;
+ const prefix = `backport (${target}) `;
+
+ const checks = await github.paginate(
+ github.rest.checks.listForRef,
+ { owner, repo, ref: headSha, per_page: 100 }
+ );
+
+ for (const check of checks) {
+ if (!check.name.startsWith(prefix)) continue;
+ if (check.status === "completed" && check.conclusion === "cancelled") continue;
+ try {
+ await github.rest.checks.update({
+ owner,
+ repo,
+ check_run_id: check.id,
+ status: "completed",
+ conclusion: "cancelled",
+ });
+ core.info(`Cancelled check ${check.name}`);
+ } catch (e) {
+ core.warning(`Failed to update check ${check.id} (${check.name}): ${e.message}`);
+ }
+ }
+
+ build:
+ needs: precheck
+ uses: ./.github/workflows/build.yml
+ with:
+ run_frontend: ${{ needs.precheck.outputs.run_frontend == 'true' }}
+ run_scala: ${{ needs.precheck.outputs.run_scala == 'true' }}
+ run_python: ${{ needs.precheck.outputs.run_python == 'true' }}
+ run_agent_service: ${{ needs.precheck.outputs.run_agent_service == 'true' }}
+ secrets: inherit
+
+ backport:
+ needs: precheck
+ if: ${{ needs.precheck.outputs.backport_targets != '[]' }}
+ strategy:
+ fail-fast: false
+ matrix:
+ target: ${{ fromJson(needs.precheck.outputs.backport_targets) }}
+ uses: ./.github/workflows/build.yml
+ with:
+ checkout_ref: refs/pull/${{ github.event.pull_request.number }}/head
+ backport_target_branch: ${{ matrix.target }}
+ backport_commit_range: ${{ format('{0}..{1}', github.event.pull_request.base.sha, github.event.pull_request.head.sha) }}
+ job_name_suffix: ""
+ run_frontend: ${{ needs.precheck.outputs.run_frontend == 'true' }}
+ run_scala: ${{ needs.precheck.outputs.run_scala == 'true' }}
+ run_python: ${{ needs.precheck.outputs.run_python == 'true' }}
+ run_agent_service: ${{ needs.precheck.outputs.run_agent_service == 'true' }}
+ secrets: inherit
+
+ required-checks:
+ # Do not rename this job — its display name is referenced in .asf.yaml.
+ name: Required Checks
+ needs: [precheck, build, backport]
+ if: always()
+ runs-on: ubuntu-latest
+ steps:
+ - name: Verify all required checks succeeded or were skipped
+ run: |
+ declare -A results=(
+ [precheck]="${{ needs.precheck.result }}"
+ [build]="${{ needs.build.result }}"
+ [backport]="${{ needs.backport.result }}"
+ )
+ failed=0
+ for job in "${!results[@]}"; do
+ r="${results[$job]}"
+ echo "${job}: ${r}"
+ if [[ "$r" != "success" && "$r" != "skipped" ]]; then
+ failed=1
+ fi
+ done
+ if (( failed )); then
+ echo "::error::One or more required checks did not succeed."
+ exit 1
+ fi
+ echo "All required checks succeeded or were skipped."
diff --git a/.gitignore b/.gitignore
index fd0f230a2d9..d283fc70137 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,25 +1,37 @@
-*.iml
-/.idea
-/.idea_modules
-/.project
-/.settings
+# Ignoring binary/output
target/
out/
-user-resources/
+# Ignoring packages
+*.jar
+*.war
+*.nar
+*.ear
+*.zip
+*.tar.gz
+*.rar
-# Ignoring IntelliJ related files
-.idea
+# Ignoring VSCode related files
+.vscode/
+# Ignoring IntelliJ related files
+*.iml
+.idea/
+.idea_modules/
+lib_managed/
+src_managed/
# Ignoring Eclipse files
.classpath
.project
.settings
+# Ignoring sublime files
+*.sublime-workspace
+
# Ignoring index folder and data folder
-index
-catalog
+index/
+catalog/
plan/
plan_files/
query-results/
@@ -27,3 +39,96 @@ query-results/
# Ignoring Mac OSX specific files
.DS_Store
+# Ignoring jenv related files
+.java-version
+
+# Ignoring scala related files
+hs_err_pid*
+
+# Ignoring python related files
+venv/
+__pycache__/
+*.py[cod]
+*$py.class
+.ipynb_checkpoints
+.pytype/
+
+# Ignoring python generated files
+*.model
+*.pkl
+
+# Ingoring user generated resources
+user-resources/
+
+# Ingoring gmail tokens
+gmail/
+
+# Ignoring maven related files
+pom.xml.tag
+pom.xml.releaseBackup
+pom.xml.versionsBackup
+pom.xml.next
+release.properties
+dependency-reduced-pom.xml
+buildNumber.properties
+.mvn/timing.properties
+# https://github.com/takari/maven-wrapper#usage-without-binary-jar
+.mvn/wrapper/maven-wrapper.jar
+
+# Ignoring sbt related files
+.bsp/
+sbt.json
+
+# Ignoring rebel related files
+rebel.xml
+
+# Ignoring log files
+*.log
+*.log.gz
+
+# Ignoring the entire log folder
+logs/
+
+# Ignoring package-lock.json
+package-lock.json
+
+# Ignoring protobuf related files
+scalapb/scalapb
+
+# Ignoring credentials
+client_secret_*
+StoredCredential*
+**/apache2/
+**/Apache24/
+**/php/
+Composer-Setup.exe
+
+# Ignoring folders generated by vscode IDE
+.metals/
+.bloop/
+.ammonite/
+metals.sbt
+
+# Ignoring Helm related files
+**/charts/*.tgz
+**/texera-helmchart/charts/
+**/texera-helmchart/*/requirements.lock
+**/texera-helmchart/Chart.lock
+**/.helm/
+
+# Additional Helm/Kubernetes related files
+.kubeconfig
+*.kubeconfig
+**/.kube/
+values-*.yaml # Any environment-specific value overrides
+
+# Ignore nested node modules
+**/node_modules/
+**/package-lock.json
+.env
+
+# agent-service is Bun-based; yarn/npm artifacts aren't needed
+agent-service/.yarn/
+agent-service/.yarnrc.yml
+agent-service/yarn.lock
+agent-service/CLAUDE.md
diff --git a/.licenserc.yaml b/.licenserc.yaml
new file mode 100644
index 00000000000..565ef344f80
--- /dev/null
+++ b/.licenserc.yaml
@@ -0,0 +1,51 @@
+header:
+ license:
+ spdx-id: Apache-2.0
+ copyright-owner: Apache Software Foundation
+
+ paths-ignore:
+ - 'licenses'
+ - '**/*.md'
+ - '**/*.csv'
+ - '**/*.txt'
+ - '**/*.json'
+ - '**/*.jsonl'
+ - 'DESCRIPTION'
+ - 'DISCLAIMER-WIP'
+ - 'LICENSE'
+ - 'LICENSE-binary'
+ - 'NOTICE'
+ - 'NOTICE-binary'
+ - '.dockerignore'
+ - '.gitattributes'
+ - '.github/PULL_REQUEST_TEMPLATE'
+ - 'yarn.lock'
+ - '.nvmrc'
+ - '.htaccess'
+ - '.gitkeep'
+ - 'site.webmanifest'
+ - '.gitignore'
+ - '.licenserc.yaml'
+ - 'frontend/.yarn/**'
+ - 'amber/src/main/python/proto/**'
+ - '**/.env.example'
+ - '**/.prettierrc'
+ - '**/bun.lock'
+ # Third-party code with MIT license - see LICENSE file for attribution
+ - 'common/workflow-operator/src/main/scala/com/kjetland/**'
+ # TypeFox monaco-languageclient derived files (MIT License)
+ - 'pyright-language-service/src/main.ts'
+ - 'pyright-language-service/src/language-server-runner.ts'
+ - 'pyright-language-service/src/server-commons.ts'
+ - 'frontend/src/app/common/formly/array.type.ts'
+ - 'frontend/src/app/common/formly/object.type.ts'
+ - 'frontend/src/app/common/formly/multischema.type.ts'
+ - 'frontend/src/app/common/formly/null.type.ts'
+ # Third-party SVG assets - see LICENSE file for attribution
+ - 'frontend/src/assets/svg/operator-view-result.svg'
+ - 'frontend/src/assets/svg/operator-reuse-cache-invalid.svg'
+ - 'frontend/src/assets/svg/operator-reuse-cache-valid.svg'
+ - 'frontend/src/app/common/type/proto/org/apache/texera/amber/core/virtualidentity.ts'
+ - 'frontend/src/app/common/type/proto/org/apache/texera/amber/core/workflow.ts'
+ - 'frontend/src/app/common/type/proto/google/protobuf/descriptor.ts'
+ - 'frontend/src/app/common/type/proto/scalapb/scalapb.ts'
diff --git a/.run/AccessControlService.run.xml b/.run/AccessControlService.run.xml
new file mode 100644
index 00000000000..56ffbc50fbb
--- /dev/null
+++ b/.run/AccessControlService.run.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/ComputingUnitManagingService.run.xml b/.run/ComputingUnitManagingService.run.xml
new file mode 100644
index 00000000000..8bbb5c1966f
--- /dev/null
+++ b/.run/ComputingUnitManagingService.run.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/ComputingUnitMaster.run.xml b/.run/ComputingUnitMaster.run.xml
new file mode 100644
index 00000000000..5ffffb79e88
--- /dev/null
+++ b/.run/ComputingUnitMaster.run.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/ComputingUnitWorker.run.xml b/.run/ComputingUnitWorker.run.xml
new file mode 100644
index 00000000000..854da672287
--- /dev/null
+++ b/.run/ComputingUnitWorker.run.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/ConfigService.run.xml b/.run/ConfigService.run.xml
new file mode 100644
index 00000000000..395688afca4
--- /dev/null
+++ b/.run/ConfigService.run.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/FileService.run.xml b/.run/FileService.run.xml
new file mode 100644
index 00000000000..0a54e70aac8
--- /dev/null
+++ b/.run/FileService.run.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/TexeraWebApplication.run.xml b/.run/TexeraWebApplication.run.xml
new file mode 100644
index 00000000000..d9e3829fc22
--- /dev/null
+++ b/.run/TexeraWebApplication.run.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/WorkflowCompilingService.run.xml b/.run/WorkflowCompilingService.run.xml
new file mode 100644
index 00000000000..d41f60d9c54
--- /dev/null
+++ b/.run/WorkflowCompilingService.run.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/frontend.run.xml b/.run/frontend.run.xml
new file mode 100644
index 00000000000..8132f576027
--- /dev/null
+++ b/.run/frontend.run.xml
@@ -0,0 +1,30 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/texera micro services.run.xml b/.run/texera micro services.run.xml
new file mode 100644
index 00000000000..6cdb0a02a56
--- /dev/null
+++ b/.run/texera micro services.run.xml
@@ -0,0 +1,31 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.run/texera-lakefs.run.xml b/.run/texera-lakefs.run.xml
new file mode 100644
index 00000000000..e26ec964fd2
--- /dev/null
+++ b/.run/texera-lakefs.run.xml
@@ -0,0 +1,28 @@
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.scalafix.conf b/.scalafix.conf
new file mode 100644
index 00000000000..238028c0ce4
--- /dev/null
+++ b/.scalafix.conf
@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+rules = [
+ ProcedureSyntax,
+ RemoveUnused,
+]
+RemoveUnused.imports = true
+RemoveUnused.privates = true
+RemoveUnused.locals = false
+RemoveUnused.patternvars = false
+RemoveUnused.params = false
\ No newline at end of file
diff --git a/.scalafmt.conf b/.scalafmt.conf
new file mode 100644
index 00000000000..8ed0d7e717f
--- /dev/null
+++ b/.scalafmt.conf
@@ -0,0 +1,19 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+version=2.6.4
+maxColumn = 100
diff --git a/.travis.yml b/.travis.yml
deleted file mode 100644
index 0604d247907..00000000000
--- a/.travis.yml
+++ /dev/null
@@ -1,26 +0,0 @@
-# configuration
-dist: trusty
-sudo: required
-language: java
-cache:
- - apt
- - $HOME/.m2
-
-# build matrix
-jdk:
- - oraclejdk8
-
-before_install:
- - sudo apt-get update
- - sudo apt-get install python3
- - sudo apt-get install python3-setuptools
- - sudo easy_install3 pip
- - export LC_ALL=C
- - sudo pip3 install -U nltk
-
-# run steps
-install:
-
-script:
- - cd core
- - mvn test
diff --git a/AGENTS.md b/AGENTS.md
new file mode 100644
index 00000000000..70e31175d35
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1,235 @@
+# AGENTS.md
+
+Guidance for coding agents working in the Apache Texera repository.
+
+## Project Overview
+
+Apache Texera is a collaborative data science and AI/ML workflow system. The
+repo is a multi-language monorepo with Scala/sbt backend services, Python worker
+runtime code, an Angular frontend, and a TypeScript/Bun agent service.
+
+Major areas:
+
+- `amber/`: workflow execution engine, Scala tests, Python worker runtime, and
+ Python operator dependencies.
+- `common/`: shared Scala modules for auth, config, DAO, workflow core,
+ workflow operators, and Python-template building.
+- `config-service/`, `access-control-service/`, `file-service/`,
+ `computing-unit-managing-service/`, `workflow-compiling-service/`: backend
+ services wired through `build.sbt`.
+- `frontend/`: Angular application. Uses Yarn 4.14.1, Node >= 20.19.0, Nx,
+ Prettier, ESLint, Karma/Jasmine, and ng-zorro.
+- `agent-service/`: TypeScript Elysia service for Texera LLM agents. CI uses
+ Bun 1.3.3.
+- `pyright-language-service/`: TypeScript service for Python language support.
+- `sql/`: database DDL used by local runs and CI.
+- `bin/`: shell scripts and Dockerfiles for services, local deployment, and
+ generated protobuf assets.
+
+## Ground Rules
+
+- Keep changes narrowly scoped. Do not rewrite unrelated files or move code
+ between services unless the task explicitly requires it.
+- Preserve local user changes. Check `git status --short` before editing and do
+ not revert unrelated dirty files.
+- Follow existing module boundaries and naming patterns. Prefer local helpers
+ and service abstractions over introducing new framework-level utilities.
+- Add or update tests when behavior changes. For small UI-only fixes where unit
+ tests are not practical, document the manual test steps.
+- Never commit secrets, local config, generated build output, caches, or binary
+ artifacts. Examples to avoid include `python_udf.conf`, `.env` files, `target/`,
+ `dist/`, `.pytest_cache/`, `.ruff_cache/`, and local logs.
+
+## Licensing
+
+- New source/config files should include the Apache 2.0 ASF license header unless
+ `.licenserc.yaml` excludes that file type or path.
+- Markdown files are excluded from the license-header check.
+- Keep third-party/vendored-code attribution intact. `common/workflow-operator`
+ has special license handling in `project/AddMetaInfLicenseFiles.scala`.
+- GitHub Actions in ASF repositories should use approved actions and preferably
+ pinned SHAs, matching the existing workflow style.
+
+## Scala / Backend
+
+- Scala version: 2.13.18.
+- Java in CI: Temurin JDK 11.
+- Formatting: `.scalafmt.conf` uses scalafmt 2.6.4 with `maxColumn = 100`.
+- Lint rules live in `.scalafix.conf` and include `ProcedureSyntax` and
+ `RemoveUnused`.
+
+Useful root commands:
+
+```bash
+sbt scalafmtCheckAll
+sbt scalafmtAll
+sbt "scalafixAll --check"
+sbt scalafixAll
+sbt clean package
+sbt test
+```
+
+Targeted tests are preferred while iterating. Examples:
+
+```bash
+sbt "WorkflowExecutionService/testOnly org.apache.texera.amber.engine.e2e.ReconfigurationSpec"
+sbt "WorkflowCompilingService/testOnly *SomeSpec"
+```
+
+CI creates PostgreSQL databases from:
+
+```bash
+psql -h localhost -U postgres -f sql/texera_ddl.sql
+psql -h localhost -U postgres -f sql/iceberg_postgres_catalog.sql
+psql -h localhost -U postgres -f sql/texera_lakefs.sql
+psql -h localhost -U postgres -v DB_NAME=texera_db_for_test_cases -f sql/texera_ddl.sql
+```
+
+## Python Runtime
+
+Python worker code lives primarily under `amber/src/main/python`.
+
+- Supported CI Python versions: 3.10, 3.11, 3.12, 3.13.
+- Ruff config is in `amber/src/main/python/pyproject.toml`.
+- Ruff line length is 88 and target version is `py310`.
+- Generated protobuf code under `amber/src/main/python/proto` is excluded from
+ Ruff.
+
+Useful commands:
+
+```bash
+cd amber/src/main/python
+ruff check .
+ruff format --check .
+pytest -sv
+python -m pytest core/runnables/test_main_loop.py -v
+```
+
+Install dependencies from `amber/requirements.txt` and
+`amber/operator-requirements.txt` when running the Python runtime or tests
+outside CI.
+
+## Frontend
+
+The Angular frontend lives in `frontend/`.
+
+- Node engine: `>=20.19.0`.
+- Package manager: Yarn 4.14.1 via Corepack.
+- Formatting is Prettier plus prettier-eslint. Prettier uses 2 spaces,
+ semicolons, double quotes, `printWidth: 120`, and LF endings.
+- Unit tests are Karma/Jasmine. Specs should live next to frontend code as
+ `.spec.ts` files.
+
+Useful commands:
+
+```bash
+cd frontend
+corepack enable
+corepack prepare yarn@4.14.1 --activate
+yarn install --immutable --inline-builds --network-timeout=100000
+yarn format:ci
+yarn format:fix
+yarn lint
+yarn test --watch=false
+yarn test:ci
+yarn build:ci
+yarn start
+```
+
+For UI changes, include screenshots/GIFs or clear manual verification steps in
+the PR description when the behavior is visual or interactive.
+
+## Agent Service
+
+The standalone LLM agent service lives in `agent-service/`.
+
+- Runtime/package tool in CI: Bun 1.3.3.
+- Source is TypeScript ESM.
+
+Useful commands:
+
+```bash
+cd agent-service
+bun install --frozen-lockfile
+bun run format:check
+bun run typecheck
+bun test
+bun run dev
+```
+
+## GitHub PR Writing
+
+Texera requires Conventional Commit PR titles and commit messages. Closed PRs
+commonly use titles like:
+
+- `feat(agent-service): enable Texera Agent to do workflow editing and execution`
+- `fix(amber): Python internal marker replay during reconfiguration`
+- `fix(frontend): version history timestamp display`
+- `test(amber-python): add unit tests for evaluate-expression and retry-current-tuple handlers`
+- `chore(deps): upgrade frontend to Angular 21`
+- `ci: bump coursier/cache-action to v8.1.0`
+
+Use the existing `.github/PULL_REQUEST_TEMPLATE` sections:
+
+- `What changes were proposed in this PR?`
+- `Any related issues, documentation, discussions?`
+- `How was this PR tested?`
+- `Was this PR authored or co-authored using generative AI tooling?`
+
+PR description conventions from recent closed PRs:
+
+- Start with the reason for the change, not just the files touched.
+- For bugs, state the root cause and the before/after behavior.
+- For features, describe the user-facing capability and key implementation
+ pieces.
+- Link issues with `Closes #1234`, `Fixes #1234`, or `Resolves #1234` when the
+ PR should close the issue.
+- Include exact test commands and, when useful, the specific test names or pass
+ counts.
+- For UI work, add screenshots/GIFs or explicit manual verification notes.
+- If no automated tests were added, explain why and list manual tests.
+- Answer the AI tooling question explicitly. If AI was used, use the
+ `Generated-by: ` wording from the template or a similarly
+ explicit disclosure. If not, write `No`.
+
+## GitHub Issue Writing
+
+Use the issue templates in `.github/ISSUE_TEMPLATE`.
+
+Bug reports should include:
+
+- What happened and what was expected.
+- Reproduction steps that another contributor can run.
+- Texera version, usually `1.1.0-incubating (Pre-release/Master)` for current
+ main.
+- Commit hash when known.
+- Browser information for frontend bugs.
+- Relevant logs or stack traces in fenced code blocks.
+
+Task and feature issues should include:
+
+- A concise task/feature summary.
+- Motivation or user impact.
+- Proposed action or scope, ideally as concrete bullets.
+- Priority (`P0` through `P3`) and task type.
+- File paths, classes, or modules when the work is already localized.
+
+Recent closed issues are usually specific and actionable: they name the failing
+test, exact command, affected files/classes, observable symptoms, and expected
+fix direction. Preserve that style for future issues.
+
+## Before Opening a PR
+
+Run the narrowest checks that cover the change, then broaden when touching shared
+behavior:
+
+- Scala/backend: targeted `testOnly`, then `sbt scalafmtCheckAll`,
+ `sbt "scalafixAll --check"`, and `sbt test` as appropriate.
+- Python runtime: `ruff check .`, `ruff format --check .`, and targeted/full
+ `pytest` from `amber/src/main/python`.
+- Frontend: `yarn format:ci`, targeted/full `yarn test:ci`, and
+ `yarn build:ci`.
+- Agent service: `bun run format:check`, `bun run typecheck`, and `bun test`.
+
+If a full check is too expensive or cannot run locally, state exactly what was
+run and why the omitted check was skipped.
diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 100644
index 00000000000..5c64d0f0e3f
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1,3 @@
+# CLAUDE.md
+
+Use the project guidance in [AGENTS.md](AGENTS.md).
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 00000000000..f31b0052e09
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,130 @@
+# Contributing to Texera
+
+Thank you for your interest in contributing to Texera! Please follow the steps below to submit your contributions effectively. We follow a **fork-based development workflow** and adopt the [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) specification for commit messages and pull request titles.
+
+---
+## Different roles in the project
+
+| Role | Key Permissions | How to Join
+| -------- | ------- | ------- |
+| Contributor | Submit issues & PRs, join discussions | Start contributing — no formal process |
+| Committer | Merge PRs, push code, vote on code changes | Voted by PPMC based on quality contributions |
+| PPMC Member | Governance, vote on releases & new committers/PPMC | Voted by PPMC members |
+| Mentor | Guide project, oversee releases, ensure Apache policies followed | Appointed by Incubator PMC — must be an experienced Apache member |
+
+## 🛠 Contribution Steps
+
+### 1. Fork the Repo
+- Fork the [Texera repository](https://github.com/Texera/texera) to your own GitHub account.
+
+### 2. Find an Existing Issue or Open an Issue
+- Find an existing issue that you want to work on, or create one issue for new proposal/bug description.
+- Have a discussion on the issue with Texera Committers (@Committer).
+- Reach a consensus before you work on the development related to the issue.
+
+### 3. Open a Pull Request (PR)
+- Create a new branch in your fork for your contribution.
+- Once you are done with the development, submit a PR from your fork to the original Texera repository.
+- **Check** the option **"Allow edits from maintainers"** so that Texera Committers can make minor edits to your PR if needed.
+
+#### PR Title and Commit Messages
+- We require all PR titles and commit messages to follow the [Conventional Commits spec](https://www.conventionalcommits.org/en/v1.0.0/).
+- All PR titles will be used as the **squashed commit message** when merged into the `master` branch.
+- Example PR titles:
+ - `feat: add a new join operator`
+ - `fix(ui): prevent racing of requests`
+ - `chore(deps): bump numpy to version 2.0.0`
+
+> 💡 You can use the [Conventional Commits plugin](https://plugins.jetbrains.com/plugin/13389-conventional-commit) in IntelliJ to help format commit messages correctly.
+
+#### PR Description
+Your pull request description should include:
+
+- **Purpose** of the PR:
+ - If your PR addresses an issue, use `Closes #1234` to automatically close it.
+ - If it relates to an issue or another PR, reference it with `#` or `#`.
+- **Summary** of changes.
+- Optional **design proposal** created based on the [template](https://docs.google.com/document/d/1ih6jLni4GgKETxOAlTOPjarlbeY5ccB2g9y1vK-Xhck/edit?usp=sharing).
+- Optional **technical design diagram** or description.
+- Optional **GIFs or screenshots** for UI-related changes.
+
+#### Avoid Including Sensitive Information
+Do not include any of the following in your PR:
+
+- Local configuration files (e.g., `python_udf.conf`)
+- Secrets or credentials (e.g., passwords, tokens)
+- Build artifacts or binary files
+
+### Final Steps Before Review
+#### Your PR should pass scalafix check (lint) and scalafmt check.
+- To check lint, under the root directory run command `sbt "scalafixAll --check"`; to fix lint issues, run `sbt scalafixAll`.
+- To check format, under the root directory run command `sbt scalafmtCheckAll`; to fix format, run `sbt scalafmtAll`.
+- When you need to execute both, scalafmt is supposed to be executed after scalafix.
+#### Testing the backend
+1. The test framework is `scalatest`, for the amber engine, tests are located under `amber/src/test`; for `WorkflowCompilingService`, tests are located under `workflow-compiling-service/src/test`. You can find unit tests and e2e tests.
+2. To execute it, navigate to the root directory in the command line and execute `sbt test`.
+3. If using IntelliJ to execute the test cases please make sure to be at the correct working directory.
+* For the amber engine's tests, the working directory should be `amber`
+* For the other services' tests, the working directory should be the root directory
+#### Testing the frontend
+Before merging your code to the master branch, you need to pass the existing unit tests first.
+1. Open a command line. Navigate to the `frontend` directory.
+2. Start the test:
+```
+ng test --watch=false
+```
+3. Wait for some time and the test will get started.
+You should also write some unit tests to cover your code. When others need to change your code, they will have to pass these unit tests so that you can keep your features safe.
+The unit tests should be written inside `.spec.ts` file.
+4. Run the following command to fix the formatting of the frontend code.
+```
+yarn format:fix
+```
+
+### 4. PR Review
+- [ ] Ask a Texera Committer (by commenting on the PR) to triage your PR, i.e., request a reviewer, and assign the PR to you.
+- [ ] Add appropriate labels such as `fix`, `enhancement`, `docs`, etc.
+- [ ] If the change should also land in a release branch, add the matching `release/` label (e.g. `release/v1.1.0-incubating`); the change will be backported to that branch automatically.
+- [ ] Ensure that all CI checks pass (see [GitHub Actions](https://github.com/Texera/texera/actions)).
+- [ ] Fully test your changes locally.
+
+> ℹ️ If your PR is not ready for review, please mark it as a draft. You can change it to “Ready for review” when it is complete.
+
+### 5. After PR Approval
+- [ ] Wait for a Texera Committer, usually the reviewer, to merge the PR once it is approved.
+- [ ] Close the related issue once the PR is merged (if it is not automatically closed).
+
+---
+
+## 📝 Apache License Header
+
+All new files must include the Apache License header.
+
+If you are modifying existing files, you may skip this step. For new files, you can automate this in IntelliJ by setting up a Copyright profile.
+
+### Steps in IntelliJ:
+
+1. Go to **Settings → Editor → Copyright → Copyright Profiles**.
+2. Create a new profile and name it **Apache**.
+3. Use the following license text:
+ ```
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+ ```
+4. Go to "Editor" → "Copyright" and choose the "Apache" profile as the default profile for this
+ project.
+5. Click "Apply".
diff --git a/DISCLAIMER-WIP b/DISCLAIMER-WIP
new file mode 100644
index 00000000000..cbc718569a6
--- /dev/null
+++ b/DISCLAIMER-WIP
@@ -0,0 +1,23 @@
+Apache Texera is an effort undergoing incubation at The Apache Software
+Foundation (ASF), sponsored by the Apache Incubator PMC. Incubation is
+required of all newly accepted projects until a further review indicates
+that the infrastructure, communications, and decision-making process have
+stabilized in a manner consistent with other successful ASF projects.
+While incubation status is not necessarily a reflection of the
+completeness or stability of the code, it does indicate that the project
+has yet to be fully endorsed by the ASF.
+
+Some of the incubating project's releases may not be fully compliant
+with ASF policy. For example, releases may have incomplete or
+un-reviewed licensing conditions. What follows is a list of issues
+the project is currently aware of (this list is likely to be incomplete):
+
+- LICENSE and NOTICE files in binary artifacts (Docker images, JARs) may
+ not yet fully account for all bundled third-party dependencies. A
+ comprehensive dependency license audit is in progress.
+
+If you are planning to incorporate this work into your product/project,
+please be aware that you will need to conduct a thorough licensing
+review to determine the overall implications of including this work.
+For the current status of this project through the Apache Incubator,
+visit: https://incubator.apache.org/projects/texera.html
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 00000000000..61dc3049460
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+FROM node:18-alpine AS nodegui
+
+WORKDIR /gui
+COPY core/gui/package.json core/gui/yarn.lock ./
+RUN corepack enable && corepack prepare yarn@4.5.1 --activate && yarn set version --yarn-path 4.5.1
+# Fake git-version.js during yarn install to prevent git from causing cache
+# invalidation of dependencies
+RUN touch git-version.js && YARN_NODE_LINKER=node-modules yarn install
+
+COPY core/gui .
+# Position of .git doesn't matter since it's only there for the revision hash
+COPY .git ./.git
+RUN apk add --no-cache git && \
+ node git-version.js && \
+ apk del git && \
+ yarn run build
+
+FROM sbtscala/scala-sbt:eclipse-temurin-jammy-11.0.17_8_1.9.3_2.13.11
+
+# copy all projects under core to /core
+WORKDIR /core
+COPY core/ .
+
+RUN apt-get update
+RUN apt-get install -y netcat unzip python3-pip
+RUN pip3 install python-lsp-server python-lsp-server[websockets]
+RUN pip3 install -r requirements.txt
+RUN pip3 install -r operator-requirements.txt
+
+WORKDIR /core
+# Add .git for runtime calls to jgit from OPversion
+COPY .git ../.git
+COPY --from=nodegui /gui/dist ./gui/dist
+
+RUN ../bin/build-services.sh
+
+CMD ["../bin/deploy-docker.sh"]
+
+EXPOSE 8080
+
+EXPOSE 9090
+
+EXPOSE 8085
\ No newline at end of file
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 00000000000..6734840ae8d
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,239 @@
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ 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.
+
+--------------------------------------------------------------------------------
+THIRD-PARTY DEPENDENCIES
+--------------------------------------------------------------------------------
+
+This product bundles source code and assets from third-party projects.
+See licenses/ for the full text of these licenses:
+ - licenses/LICENSE-MIT.txt
+
+MIT License (licenses/LICENSE-MIT.txt)
+--------------------------------------
+
+This product bundles code derived from mbknor-jackson-jsonschema:
+ - common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/
+ Copyright (c) 2016 Kjell Tore Eliassen (mbknor)
+ Source: https://github.com/mbknor/mbknor-jackson-jsonschema
+
+This product bundles code derived from Google Angular formly examples:
+ - frontend/src/app/common/formly/array.type.ts
+ - frontend/src/app/common/formly/object.type.ts
+ - frontend/src/app/common/formly/multischema.type.ts
+ - frontend/src/app/common/formly/null.type.ts
+ Copyright (c) 2018 Google Inc. All Rights Reserved.
+ Source: https://angular.io
+
+This product bundles code derived from TypeFox monaco-languageclient:
+ - pyright-language-service/src/main.ts
+ - pyright-language-service/src/language-server-runner.ts
+ - pyright-language-service/src/server-commons.ts
+ Copyright (c) 2024 TypeFox and others.
+ Source: https://github.com/TypeFox/monaco-languageclient
+
+This product includes SVG icons from SVGRepo:
+ - frontend/src/assets/svg/operator-view-result.svg
+ - frontend/src/assets/svg/operator-reuse-cache-valid.svg
+ - frontend/src/assets/svg/operator-reuse-cache-invalid.svg
+ Source: https://www.svgrepo.com
+ License: MIT License
diff --git a/LICENSE-binary b/LICENSE-binary
new file mode 100644
index 00000000000..18127728eea
--- /dev/null
+++ b/LICENSE-binary
@@ -0,0 +1,1242 @@
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for describing the origin of the Work and
+ reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Support. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or support.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ 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.
+
+================================================================================
+THIRD-PARTY COMPONENTS
+================================================================================
+
+Apache Texera's binary distributions and source tree include the following
+third-party components, grouped by license. Each section references
+licenses/ for the full text of the applicable license. Components under
+the Apache License, Version 2.0 are governed by the same license terms as
+Apache Texera itself and are listed for completeness.
+
+Locations within the distribution:
+
+ - Scala/Java jars listed below ship under the lib/ directory of each
+ service's Universal zip (amber, access-control-service,
+ config-service, computing-unit-managing-service, file-service, and
+ workflow-compiling-service).
+
+ - Python packages are installed into the Python site-packages directory
+ of the computing-unit-master and computing-unit-worker Docker images.
+
+ - Angular / npm packages are sourced from frontend/node_modules at
+ build time and shipped compiled into the texera-web-application
+ Docker image's frontend assets.
+
+ - Source files derived from third-party projects live in the Apache
+ Texera source tree at the paths listed against each entry.
+
+--------------------------------------------------------------------------------
+Dependencies under the Apache License, Version 2.0
+--------------------------------------------------------------------------------
+
+Scala/Java jars:
+ - ch.qos.reload4j.reload4j-1.2.18.3.jar
+ - com.fasterxml.classmate-1.3.1.jar
+ - com.fasterxml.classmate-1.7.0.jar
+ - com.fasterxml.jackson.core.jackson-annotations-2.18.6.jar
+ - com.fasterxml.jackson.core.jackson-core-2.18.6.jar
+ - com.fasterxml.jackson.core.jackson-databind-2.18.6.jar
+ - com.fasterxml.jackson.dataformat.jackson-dataformat-yaml-2.16.1.jar
+ - com.fasterxml.jackson.dataformat.jackson-dataformat-yaml-2.17.0.jar
+ - com.fasterxml.jackson.dataformat.jackson-dataformat-yaml-2.9.10.jar
+ - com.fasterxml.jackson.datatype.jackson-datatype-guava-2.16.1.jar
+ - com.fasterxml.jackson.datatype.jackson-datatype-guava-2.9.10.jar
+ - com.fasterxml.jackson.datatype.jackson-datatype-jdk8-2.11.4.jar
+ - com.fasterxml.jackson.datatype.jackson-datatype-jdk8-2.16.1.jar
+ - com.fasterxml.jackson.datatype.jackson-datatype-joda-2.9.10.jar
+ - com.fasterxml.jackson.datatype.jackson-datatype-jsr310-2.16.0.jar
+ - com.fasterxml.jackson.datatype.jackson-datatype-jsr310-2.16.1.jar
+ - com.fasterxml.jackson.datatype.jackson-datatype-jsr310-2.17.0.jar
+ - com.fasterxml.jackson.jakarta.rs.jackson-jakarta-rs-base-2.16.1.jar
+ - com.fasterxml.jackson.jakarta.rs.jackson-jakarta-rs-json-provider-2.16.1.jar
+ - com.fasterxml.jackson.jaxrs.jackson-jaxrs-base-2.10.5.jar
+ - com.fasterxml.jackson.jaxrs.jackson-jaxrs-json-provider-2.10.5.jar
+ - com.fasterxml.jackson.module.jackson-module-afterburner-2.9.10.jar
+ - com.fasterxml.jackson.module.jackson-module-blackbird-2.16.1.jar
+ - com.fasterxml.jackson.module.jackson-module-jakarta-xmlbind-annotations-2.16.1.jar
+ - com.fasterxml.jackson.module.jackson-module-jaxb-annotations-2.10.5.jar
+ - com.fasterxml.jackson.module.jackson-module-jsonSchema-2.18.6.jar
+ - com.fasterxml.jackson.module.jackson-module-no-ctor-deser-2.18.6.jar
+ - com.fasterxml.jackson.module.jackson-module-parameter-names-2.16.1.jar
+ - com.fasterxml.jackson.module.jackson-module-parameter-names-2.9.10.jar
+ - com.fasterxml.jackson.module.jackson-module-scala_2.13-2.18.6.jar
+ - com.fasterxml.woodstox.woodstox-core-5.3.0.jar
+ - com.flipkart.zjsonpatch.zjsonpatch-0.4.13.jar
+ - com.github.ben-manes.caffeine.caffeine-2.9.3.jar
+ - com.github.ben-manes.caffeine.caffeine-3.1.8.jar
+ - com.github.dirkraft.dropwizard.dropwizard-file-assets-0.0.2.jar
+ - com.github.nscala-time.nscala-time_2.13-2.32.0.jar
+ - com.github.sisyphsu.dateparser-1.0.11.jar
+ - com.github.sisyphsu.retree-1.0.4.jar
+ - com.github.stephenc.jcip.jcip-annotations-1.0-1.jar
+ - com.github.toastshaman.dropwizard-auth-jwt-1.1.2-0.jar
+ - com.github.tototoshi.scala-csv_2.13-1.3.10.jar
+ - com.google.android.annotations-4.1.1.4.jar
+ - com.google.api-client.google-api-client-2.2.0.jar
+ - com.google.api.grpc.proto-google-common-protos-2.22.0.jar
+ - com.google.api.grpc.proto-google-common-protos-2.29.0.jar
+ - com.google.code.findbugs.jsr305-3.0.2.jar
+ - com.google.code.gson.gson-2.10.1.jar
+ - com.google.code.gson.gson-2.11.0.jar
+ - com.google.errorprone.error_prone_annotations-2.23.0.jar
+ - com.google.errorprone.error_prone_annotations-2.25.0.jar
+ - com.google.errorprone.error_prone_annotations-2.27.0.jar
+ - com.google.flatbuffers.flatbuffers-java-23.5.26.jar
+ - com.google.guava.failureaccess-1.0.2.jar
+ - com.google.guava.guava-33.0.0-jre.jar
+ - com.google.guava.listenablefuture-9999.0-empty-to-avoid-conflict-with-guava.jar
+ - com.google.http-client.google-http-client-1.42.3.jar
+ - com.google.http-client.google-http-client-apache-v2-1.42.3.jar
+ - com.google.http-client.google-http-client-gson-1.42.3.jar
+ - com.google.inject.extensions.guice-servlet-4.0.jar
+ - com.google.inject.guice-4.0.jar
+ - com.google.j2objc.j2objc-annotations-2.8.jar
+ - com.google.oauth-client.google-oauth-client-1.34.1.jar
+ - com.google.oauth-client.google-oauth-client-java6-1.34.1.jar
+ - com.google.oauth-client.google-oauth-client-jetty-1.34.1.jar
+ - com.googlecode.javaewah.JavaEWAH-1.1.12.jar
+ - com.helger.profiler-1.1.1.jar
+ - com.hierynomus.asn-one-0.6.0.jar
+ - com.nimbusds.nimbus-jose-jwt-9.8.1.jar
+ - com.papertrail.profiler-1.0.2.jar
+ - com.softwaremill.common.tagging_2.13-2.3.5.jar
+ - com.softwaremill.macwire.proxy_2.13-2.6.7.jar
+ - com.softwaremill.macwire.util_2.13-2.6.7.jar
+ - com.softwaremill.sttp.client4.core_2.13-4.0.0-M6.jar
+ - com.softwaremill.sttp.model.core_2.13-1.7.2.jar
+ - com.softwaremill.sttp.shared.core_2.13-1.3.16.jar
+ - com.softwaremill.sttp.shared.ws_2.13-1.3.16.jar
+ - com.squareup.okhttp.okhttp-2.7.5.jar
+ - com.squareup.okhttp3.logging-interceptor-4.12.0.jar
+ - com.squareup.okhttp3.okhttp-4.12.0.jar
+ - com.squareup.okio.okio-3.6.0.jar
+ - com.squareup.okio.okio-jvm-3.6.0.jar
+ - com.thesamet.scalapb.lenses_2.13-0.11.20.jar
+ - com.thesamet.scalapb.scalapb-json4s_2.13-0.12.0.jar
+ - com.thesamet.scalapb.scalapb-runtime-grpc_2.13-0.11.20.jar
+ - com.thesamet.scalapb.scalapb-runtime_2.13-0.11.20.jar
+ - com.twitter.util-core_2.13-22.12.0.jar
+ - com.twitter.util-function_2.13-22.12.0.jar
+ - com.typesafe.config-1.4.6.jar
+ - com.typesafe.play.play-functional_2.13-2.10.6.jar
+ - com.typesafe.play.play-functional_2.13-2.9.4.jar
+ - com.typesafe.play.play-json_2.13-2.10.6.jar
+ - com.typesafe.play.play-json_2.13-2.9.4.jar
+ - com.typesafe.scala-logging.scala-logging_2.13-3.9.5.jar
+ - com.typesafe.ssl-config-core_2.13-0.6.1.jar
+ - com.univocity.univocity-parsers-2.9.1.jar
+ - commons-beanutils.commons-beanutils-1.9.4.jar
+ - commons-cli.commons-cli-1.2.jar
+ - commons-codec.commons-codec-1.17.1.jar
+ - commons-collections.commons-collections-3.2.2.jar
+ - commons-io.commons-io-2.16.1.jar
+ - commons-logging.commons-logging-1.2.jar
+ - commons-net.commons-net-3.6.jar
+ - commons-pool.commons-pool-1.6.jar
+ - dev.failsafe.failsafe-3.3.2.jar
+ - io.airlift.aircompressor-0.27.jar
+ - io.altoo.pekko-kryo-serialization_2.13-1.3.0.jar
+ - io.altoo.scala-kryo-serialization_2.13-1.3.0.jar
+ - io.dropwizard-bundles.dropwizard-redirect-bundle-1.0.5.jar
+ - io.dropwizard.dropwizard-auth-1.3.23.jar
+ - io.dropwizard.dropwizard-auth-4.0.7.jar
+ - io.dropwizard.dropwizard-client-1.3.23.jar
+ - io.dropwizard.dropwizard-configuration-1.3.23.jar
+ - io.dropwizard.dropwizard-configuration-4.0.7.jar
+ - io.dropwizard.dropwizard-core-1.3.23.jar
+ - io.dropwizard.dropwizard-core-4.0.7.jar
+ - io.dropwizard.dropwizard-health-4.0.7.jar
+ - io.dropwizard.dropwizard-jackson-1.3.23.jar
+ - io.dropwizard.dropwizard-jackson-4.0.7.jar
+ - io.dropwizard.dropwizard-jersey-1.3.23.jar
+ - io.dropwizard.dropwizard-jersey-4.0.7.jar
+ - io.dropwizard.dropwizard-jetty-1.3.23.jar
+ - io.dropwizard.dropwizard-jetty-4.0.7.jar
+ - io.dropwizard.dropwizard-lifecycle-1.3.23.jar
+ - io.dropwizard.dropwizard-lifecycle-4.0.7.jar
+ - io.dropwizard.dropwizard-logging-1.3.23.jar
+ - io.dropwizard.dropwizard-logging-4.0.7.jar
+ - io.dropwizard.dropwizard-metrics-1.3.23.jar
+ - io.dropwizard.dropwizard-metrics-4.0.7.jar
+ - io.dropwizard.dropwizard-request-logging-1.3.23.jar
+ - io.dropwizard.dropwizard-request-logging-4.0.7.jar
+ - io.dropwizard.dropwizard-servlets-1.3.23.jar
+ - io.dropwizard.dropwizard-servlets-4.0.7.jar
+ - io.dropwizard.dropwizard-util-1.3.23.jar
+ - io.dropwizard.dropwizard-util-4.0.7.jar
+ - io.dropwizard.dropwizard-validation-1.3.23.jar
+ - io.dropwizard.dropwizard-validation-4.0.7.jar
+ - io.dropwizard.logback.logback-throttling-appender-1.4.2.jar
+ - io.dropwizard.metrics.metrics-annotation-4.0.5.jar
+ - io.dropwizard.metrics.metrics-annotation-4.2.25.jar
+ - io.dropwizard.metrics.metrics-caffeine-4.2.25.jar
+ - io.dropwizard.metrics.metrics-core-4.0.5.jar
+ - io.dropwizard.metrics.metrics-core-4.2.25.jar
+ - io.dropwizard.metrics.metrics-healthchecks-4.0.5.jar
+ - io.dropwizard.metrics.metrics-healthchecks-4.2.25.jar
+ - io.dropwizard.metrics.metrics-httpclient-4.0.5.jar
+ - io.dropwizard.metrics.metrics-jakarta-servlets-4.2.25.jar
+ - io.dropwizard.metrics.metrics-jersey2-4.0.5.jar
+ - io.dropwizard.metrics.metrics-jersey3-4.2.25.jar
+ - io.dropwizard.metrics.metrics-jetty11-4.2.25.jar
+ - io.dropwizard.metrics.metrics-jetty9-4.0.5.jar
+ - io.dropwizard.metrics.metrics-jmx-4.0.5.jar
+ - io.dropwizard.metrics.metrics-jmx-4.2.25.jar
+ - io.dropwizard.metrics.metrics-json-4.0.5.jar
+ - io.dropwizard.metrics.metrics-json-4.2.25.jar
+ - io.dropwizard.metrics.metrics-jvm-4.0.5.jar
+ - io.dropwizard.metrics.metrics-jvm-4.2.25.jar
+ - io.dropwizard.metrics.metrics-logback-4.0.5.jar
+ - io.dropwizard.metrics.metrics-logback-4.2.25.jar
+ - io.dropwizard.metrics.metrics-servlets-4.0.5.jar
+ - io.fabric8.kubernetes-client-6.12.1.jar
+ - io.fabric8.kubernetes-client-api-6.12.1.jar
+ - io.fabric8.kubernetes-httpclient-okhttp-6.12.1.jar
+ - io.fabric8.kubernetes-model-admissionregistration-6.12.1.jar
+ - io.fabric8.kubernetes-model-apiextensions-6.12.1.jar
+ - io.fabric8.kubernetes-model-apps-6.12.1.jar
+ - io.fabric8.kubernetes-model-autoscaling-6.12.1.jar
+ - io.fabric8.kubernetes-model-batch-6.12.1.jar
+ - io.fabric8.kubernetes-model-certificates-6.12.1.jar
+ - io.fabric8.kubernetes-model-common-6.12.1.jar
+ - io.fabric8.kubernetes-model-coordination-6.12.1.jar
+ - io.fabric8.kubernetes-model-core-6.12.1.jar
+ - io.fabric8.kubernetes-model-discovery-6.12.1.jar
+ - io.fabric8.kubernetes-model-events-6.12.1.jar
+ - io.fabric8.kubernetes-model-extensions-6.12.1.jar
+ - io.fabric8.kubernetes-model-flowcontrol-6.12.1.jar
+ - io.fabric8.kubernetes-model-gatewayapi-6.12.1.jar
+ - io.fabric8.kubernetes-model-metrics-6.12.1.jar
+ - io.fabric8.kubernetes-model-networking-6.12.1.jar
+ - io.fabric8.kubernetes-model-node-6.12.1.jar
+ - io.fabric8.kubernetes-model-policy-6.12.1.jar
+ - io.fabric8.kubernetes-model-rbac-6.12.1.jar
+ - io.fabric8.kubernetes-model-resource-6.12.1.jar
+ - io.fabric8.kubernetes-model-scheduling-6.12.1.jar
+ - io.fabric8.kubernetes-model-storageclass-6.12.1.jar
+ - io.fabric8.zjsonpatch-0.3.0.jar
+ - io.github.kostaskougios.cloning-1.10.3.jar
+ - io.grpc.grpc-api-1.60.0.jar
+ - io.grpc.grpc-api-1.62.2.jar
+ - io.grpc.grpc-context-1.60.0.jar
+ - io.grpc.grpc-context-1.62.2.jar
+ - io.grpc.grpc-core-1.60.0.jar
+ - io.grpc.grpc-core-1.62.2.jar
+ - io.grpc.grpc-netty-1.60.0.jar
+ - io.grpc.grpc-protobuf-1.60.0.jar
+ - io.grpc.grpc-protobuf-1.62.2.jar
+ - io.grpc.grpc-protobuf-lite-1.60.0.jar
+ - io.grpc.grpc-protobuf-lite-1.62.2.jar
+ - io.grpc.grpc-stub-1.60.0.jar
+ - io.grpc.grpc-stub-1.62.2.jar
+ - io.grpc.grpc-util-1.60.0.jar
+ - io.gsonfire.gson-fire-1.8.5.jar
+ - io.gsonfire.gson-fire-1.9.0.jar
+ - io.kamon.sigar-loader-1.6.6-rev002.jar
+ - io.kubernetes.client-java-21.0.0.jar
+ - io.kubernetes.client-java-api-21.0.0.jar
+ - io.kubernetes.client-java-proto-21.0.0.jar
+ - io.lakefs.sdk-1.51.0.jar
+ - io.netty.netty-3.10.6.Final.jar
+ - io.netty.netty-buffer-4.1.104.Final.jar
+ - io.netty.netty-buffer-4.1.96.Final.jar
+ - io.netty.netty-codec-4.1.104.Final.jar
+ - io.netty.netty-codec-4.1.96.Final.jar
+ - io.netty.netty-codec-http-4.1.100.Final.jar
+ - io.netty.netty-codec-http-4.1.96.Final.jar
+ - io.netty.netty-codec-http2-4.1.100.Final.jar
+ - io.netty.netty-codec-http2-4.1.96.Final.jar
+ - io.netty.netty-codec-socks-4.1.100.Final.jar
+ - io.netty.netty-common-4.1.104.Final.jar
+ - io.netty.netty-common-4.1.96.Final.jar
+ - io.netty.netty-handler-4.1.104.Final.jar
+ - io.netty.netty-handler-4.1.96.Final.jar
+ - io.netty.netty-handler-proxy-4.1.100.Final.jar
+ - io.netty.netty-resolver-4.1.104.Final.jar
+ - io.netty.netty-resolver-4.1.96.Final.jar
+ - io.netty.netty-tcnative-boringssl-static-2.0.61.Final-linux-aarch_64.jar
+ - io.netty.netty-tcnative-boringssl-static-2.0.61.Final-linux-x86_64.jar
+ - io.netty.netty-tcnative-boringssl-static-2.0.61.Final-osx-aarch_64.jar
+ - io.netty.netty-tcnative-boringssl-static-2.0.61.Final-osx-x86_64.jar
+ - io.netty.netty-tcnative-boringssl-static-2.0.61.Final-windows-x86_64.jar
+ - io.netty.netty-tcnative-boringssl-static-2.0.61.Final.jar
+ - io.netty.netty-tcnative-classes-2.0.61.Final.jar
+ - io.netty.netty-transport-4.1.104.Final.jar
+ - io.netty.netty-transport-4.1.96.Final.jar
+ - io.netty.netty-transport-native-unix-common-4.1.104.Final.jar
+ - io.netty.netty-transport-native-unix-common-4.1.96.Final.jar
+ - io.opencensus.opencensus-api-0.31.1.jar
+ - io.opencensus.opencensus-contrib-http-util-0.31.1.jar
+ - io.perfmark.perfmark-api-0.26.0.jar
+ - io.r2dbc.r2dbc-spi-0.9.0.RELEASE.jar
+ - io.reactivex.rxjava3.rxjava-3.1.6.jar
+ - io.swagger.swagger-annotations-1.6.14.jar
+ - jakarta.inject.jakarta.inject-api-2.0.1.jar
+ - jakarta.validation.jakarta.validation-api-3.0.2.jar
+ - javax.inject.javax.inject-1.jar
+ - javax.validation.validation-api-2.0.1.Final.jar
+ - joda-time.joda-time-2.11.0.jar
+ - log4j.log4j-1.2.17.jar
+ - net.minidev.accessors-smart-2.4.2.jar
+ - net.minidev.accessors-smart-2.4.7.jar
+ - net.minidev.json-smart-2.4.2.jar
+ - net.minidev.json-smart-2.4.7.jar
+ - org.agrona.agrona-1.22.0.jar
+ - org.apache.arrow.arrow-format-15.0.2.jar
+ - org.apache.arrow.arrow-memory-core-15.0.2.jar
+ - org.apache.arrow.arrow-memory-netty-15.0.2.jar
+ - org.apache.arrow.arrow-vector-15.0.2.jar
+ - org.apache.arrow.flight-core-15.0.2.jar
+ - org.apache.arrow.flight-grpc-15.0.2.jar
+ - org.apache.avro.avro-1.12.0.jar
+ - org.apache.commons.commons-collections4-4.2.jar
+ - org.apache.commons.commons-collections4-4.4.jar
+ - org.apache.commons.commons-compress-1.26.2.jar
+ - org.apache.commons.commons-compress-1.27.1.jar
+ - org.apache.commons.commons-configuration2-2.1.1.jar
+ - org.apache.commons.commons-jcs3-core-3.2.jar
+ - org.apache.commons.commons-lang3-3.13.0.jar
+ - org.apache.commons.commons-lang3-3.14.0.jar
+ - org.apache.commons.commons-lang3-3.16.0.jar
+ - org.apache.commons.commons-math3-3.1.1.jar
+ - org.apache.commons.commons-text-1.11.0.jar
+ - org.apache.commons.commons-text-1.4.jar
+ - org.apache.commons.commons-vfs2-2.9.0.jar
+ - org.apache.curator.curator-client-4.2.0.jar
+ - org.apache.curator.curator-framework-4.2.0.jar
+ - org.apache.curator.curator-recipes-4.2.0.jar
+ - org.apache.hadoop.hadoop-annotations-3.3.1.jar
+ - org.apache.hadoop.hadoop-annotations-3.3.3.jar
+ - org.apache.hadoop.hadoop-auth-3.3.1.jar
+ - org.apache.hadoop.hadoop-auth-3.3.3.jar
+ - org.apache.hadoop.hadoop-common-3.3.1.jar
+ - org.apache.hadoop.hadoop-common-3.3.3.jar
+ - org.apache.hadoop.hadoop-hdfs-client-3.3.1.jar
+ - org.apache.hadoop.hadoop-mapreduce-client-core-3.3.1.jar
+ - org.apache.hadoop.hadoop-yarn-api-3.3.1.jar
+ - org.apache.hadoop.hadoop-yarn-client-3.3.1.jar
+ - org.apache.hadoop.hadoop-yarn-common-3.3.1.jar
+ - org.apache.hadoop.thirdparty.hadoop-shaded-guava-1.1.1.jar
+ - org.apache.hadoop.thirdparty.hadoop-shaded-protobuf_3_7-1.1.1.jar
+ - org.apache.htrace.htrace-core4-4.1.0-incubating.jar
+ - org.apache.httpcomponents.client5.httpclient5-5.4.jar
+ - org.apache.httpcomponents.core5.httpcore5-5.3.jar
+ - org.apache.httpcomponents.core5.httpcore5-h2-5.3.jar
+ - org.apache.httpcomponents.httpasyncclient-4.1.5.jar
+ - org.apache.httpcomponents.httpclient-4.5.13.jar
+ - org.apache.httpcomponents.httpclient-4.5.14.jar
+ - org.apache.httpcomponents.httpcore-4.4.16.jar
+ - org.apache.httpcomponents.httpcore-nio-4.4.13.jar
+ - org.apache.httpcomponents.httpmime-4.5.13.jar
+ - org.apache.iceberg.iceberg-api-1.7.1.jar
+ - org.apache.iceberg.iceberg-aws-1.7.1.jar
+ - org.apache.iceberg.iceberg-bundled-guava-1.7.1.jar
+ - org.apache.iceberg.iceberg-common-1.7.1.jar
+ - org.apache.iceberg.iceberg-core-1.7.1.jar
+ - org.apache.iceberg.iceberg-data-1.7.1.jar
+ - org.apache.iceberg.iceberg-parquet-1.7.1.jar
+ - org.apache.kerby.kerb-admin-1.0.1.jar
+ - org.apache.kerby.kerb-client-1.0.1.jar
+ - org.apache.kerby.kerb-common-1.0.1.jar
+ - org.apache.kerby.kerb-core-1.0.1.jar
+ - org.apache.kerby.kerb-crypto-1.0.1.jar
+ - org.apache.kerby.kerb-identity-1.0.1.jar
+ - org.apache.kerby.kerb-server-1.0.1.jar
+ - org.apache.kerby.kerb-simplekdc-1.0.1.jar
+ - org.apache.kerby.kerb-util-1.0.1.jar
+ - org.apache.kerby.kerby-asn1-1.0.1.jar
+ - org.apache.kerby.kerby-config-1.0.1.jar
+ - org.apache.kerby.kerby-pkix-1.0.1.jar
+ - org.apache.kerby.kerby-util-1.0.1.jar
+ - org.apache.kerby.kerby-xdr-1.0.1.jar
+ - org.apache.kerby.token-provider-1.0.1.jar
+ - org.apache.lucene.lucene-analyzers-common-8.11.4.jar
+ - org.apache.lucene.lucene-core-8.11.4.jar
+ - org.apache.lucene.lucene-memory-8.7.0.jar
+ - org.apache.lucene.lucene-queries-8.7.0.jar
+ - org.apache.lucene.lucene-queryparser-8.7.0.jar
+ - org.apache.lucene.lucene-sandbox-8.7.0.jar
+ - org.apache.orc.orc-core-1.9.4-nohive.jar
+ - org.apache.orc.orc-shims-1.9.4.jar
+ - org.apache.parquet.parquet-avro-1.13.1.jar
+ - org.apache.parquet.parquet-column-1.13.1.jar
+ - org.apache.parquet.parquet-common-1.13.1.jar
+ - org.apache.parquet.parquet-encoding-1.13.1.jar
+ - org.apache.parquet.parquet-format-structures-1.13.1.jar
+ - org.apache.parquet.parquet-hadoop-1.13.1.jar
+ - org.apache.parquet.parquet-jackson-1.13.1.jar
+ - org.apache.pekko.pekko-actor_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-cluster-metrics_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-cluster-tools_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-cluster_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-coordination_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-persistence_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-pki_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-protobuf-v3_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-remote_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-slf4j_2.13-1.2.1.jar
+ - org.apache.pekko.pekko-stream_2.13-1.2.1.jar
+ - org.apache.yetus.audience-annotations-0.13.0.jar
+ - org.apache.zookeeper.zookeeper-3.5.6.jar
+ - org.apache.zookeeper.zookeeper-jute-3.5.6.jar
+ - org.bitbucket.b_c.jose4j-0.9.6.jar
+ - org.eclipse.jetty.jetty-annotations-9.4.18.v20190429.jar
+ - org.eclipse.jetty.jetty-client-9.4.40.v20210413.jar
+ - org.eclipse.jetty.jetty-continuation-9.4.18.v20190429.jar
+ - org.eclipse.jetty.jetty-http-11.0.20.jar
+ - org.eclipse.jetty.jetty-http-9.4.20.v20190813.jar
+ - org.eclipse.jetty.jetty-io-11.0.20.jar
+ - org.eclipse.jetty.jetty-io-9.4.40.v20210413.jar
+ - org.eclipse.jetty.jetty-jndi-9.4.18.v20190429.jar
+ - org.eclipse.jetty.jetty-plus-9.4.18.v20190429.jar
+ - org.eclipse.jetty.jetty-security-11.0.20.jar
+ - org.eclipse.jetty.jetty-security-9.4.20.v20190813.jar
+ - org.eclipse.jetty.jetty-server-11.0.20.jar
+ - org.eclipse.jetty.jetty-server-9.4.20.v20190813.jar
+ - org.eclipse.jetty.jetty-servlet-11.0.20.jar
+ - org.eclipse.jetty.jetty-servlet-9.4.20.v20190813.jar
+ - org.eclipse.jetty.jetty-servlets-11.0.20.jar
+ - org.eclipse.jetty.jetty-servlets-9.4.18.v20190429.jar
+ - org.eclipse.jetty.jetty-util-11.0.20.jar
+ - org.eclipse.jetty.jetty-util-9.4.40.v20210413.jar
+ - org.eclipse.jetty.jetty-webapp-9.4.18.v20190429.jar
+ - org.eclipse.jetty.jetty-xml-9.4.18.v20190429.jar
+ - org.eclipse.jetty.toolchain.jetty-jakarta-servlet-api-5.0.2.jar
+ - org.eclipse.jetty.toolchain.setuid.jetty-setuid-java-1.0.3.jar
+ - org.eclipse.jetty.toolchain.setuid.jetty-setuid-java-1.0.4.jar
+ - org.eclipse.jetty.websocket.javax-websocket-client-impl-9.4.18.v20190429.jar
+ - org.eclipse.jetty.websocket.javax-websocket-server-impl-9.4.18.v20190429.jar
+ - org.eclipse.jetty.websocket.websocket-api-9.4.40.v20210413.jar
+ - org.eclipse.jetty.websocket.websocket-client-9.4.40.v20210413.jar
+ - org.eclipse.jetty.websocket.websocket-common-9.4.40.v20210413.jar
+ - org.eclipse.jetty.websocket.websocket-server-9.4.18.v20190429.jar
+ - org.eclipse.jetty.websocket.websocket-servlet-9.4.18.v20190429.jar
+ - org.ehcache.sizeof-0.4.3.jar
+ - org.hibernate.hibernate-validator-5.4.3.Final.jar
+ - org.hibernate.validator.hibernate-validator-7.0.5.Final.jar
+ - org.jasypt.jasypt-1.9.3.jar
+ - org.javassist.javassist-3.30.2-GA.jar
+ - org.jboss.logging.jboss-logging-3.3.0.Final.jar
+ - org.jboss.logging.jboss-logging-3.5.3.Final.jar
+ - org.jetbrains.annotations-17.0.0.jar
+ - org.jetbrains.kotlin.kotlin-stdlib-1.9.10.jar
+ - org.jetbrains.kotlin.kotlin-stdlib-common-1.9.10.jar
+ - org.jetbrains.kotlin.kotlin-stdlib-jdk7-1.9.10.jar
+ - org.jetbrains.kotlin.kotlin-stdlib-jdk8-1.9.10.jar
+ - org.jheaps.jheaps-0.11.jar
+ - org.joda.joda-convert-2.2.2.jar
+ - org.jooq.jooq-3.16.23.jar
+ - org.json4s.json4s-ast_2.13-4.0.1.jar
+ - org.json4s.json4s-jackson-core_2.13-4.0.1.jar
+ - org.lz4.lz4-java-1.8.0.jar
+ - org.objenesis.objenesis-3.4.jar
+ - org.openapitools.jackson-databind-nullable-0.2.6.jar
+ - org.playframework.play-functional_2.13-3.1.0-M1.jar
+ - org.playframework.play-json_2.13-3.1.0-M1.jar
+ - org.roaringbitmap.RoaringBitmap-1.3.0.jar
+ - org.scala-lang.modules.scala-collection-compat_2.13-2.13.0.jar
+ - org.scala-lang.modules.scala-collection-contrib_2.13-0.3.0.jar
+ - org.scala-lang.modules.scala-parser-combinators_2.13-1.1.2.jar
+ - org.scala-lang.scala-library-2.13.18.jar
+ - org.scala-lang.scala-reflect-2.13.18.jar
+ - org.scalactic.scalactic_2.13-3.2.15.jar
+ - org.slf4j.jcl-over-slf4j-2.0.12.jar
+ - org.slf4j.log4j-over-slf4j-2.0.12.jar
+ - org.slf4j.log4j-over-slf4j-2.0.16.jar
+ - org.snakeyaml.snakeyaml-engine-2.7.jar
+ - org.typelevel.cats-effect-kernel_2.13-3.6.3.jar
+ - org.typelevel.cats-effect-std_2.13-3.6.3.jar
+ - org.typelevel.cats-effect_2.13-3.6.3.jar
+ - org.typelevel.cats-mtl_2.13-1.3.1.jar
+ - org.xerial.snappy.snappy-java-1.1.8.3.jar
+ - org.yaml.snakeyaml-1.23.jar
+ - org.yaml.snakeyaml-2.2.jar
+ - software.amazon.awssdk.annotations-2.29.51.jar
+ - software.amazon.awssdk.apache-client-2.29.51.jar
+ - software.amazon.awssdk.arns-2.29.51.jar
+ - software.amazon.awssdk.auth-2.29.51.jar
+ - software.amazon.awssdk.aws-core-2.29.51.jar
+ - software.amazon.awssdk.aws-query-protocol-2.29.51.jar
+ - software.amazon.awssdk.aws-xml-protocol-2.29.51.jar
+ - software.amazon.awssdk.checksums-2.29.51.jar
+ - software.amazon.awssdk.checksums-spi-2.29.51.jar
+ - software.amazon.awssdk.crt-core-2.29.51.jar
+ - software.amazon.awssdk.endpoints-spi-2.29.51.jar
+ - software.amazon.awssdk.http-auth-2.29.51.jar
+ - software.amazon.awssdk.http-auth-aws-2.29.51.jar
+ - software.amazon.awssdk.http-auth-aws-eventstream-2.29.51.jar
+ - software.amazon.awssdk.http-auth-spi-2.29.51.jar
+ - software.amazon.awssdk.http-client-spi-2.29.51.jar
+ - software.amazon.awssdk.identity-spi-2.29.51.jar
+ - software.amazon.awssdk.json-utils-2.29.51.jar
+ - software.amazon.awssdk.metrics-spi-2.29.51.jar
+ - software.amazon.awssdk.netty-nio-client-2.29.51.jar
+ - software.amazon.awssdk.profiles-2.29.51.jar
+ - software.amazon.awssdk.protocol-core-2.29.51.jar
+ - software.amazon.awssdk.regions-2.29.51.jar
+ - software.amazon.awssdk.retries-2.29.51.jar
+ - software.amazon.awssdk.retries-spi-2.29.51.jar
+ - software.amazon.awssdk.s3-2.29.51.jar
+ - software.amazon.awssdk.sdk-core-2.29.51.jar
+ - software.amazon.awssdk.sts-2.29.51.jar
+ - software.amazon.awssdk.third-party-jackson-core-2.29.51.jar
+ - software.amazon.awssdk.utils-2.29.51.jar
+ - software.amazon.eventstream.eventstream-1.0.1.jar
+
+Python packages:
+ - aiobotocore==2.25.1
+ - aiohttp==3.13.5
+ - aiosignal==1.4.0
+ - boto3==1.40.53
+ - botocore==1.40.53
+ - frozenlist==1.8.0
+ - hf-xet==1.4.3
+ - huggingface-hub==0.36.2
+ - multidict==6.7.1
+ - overrides==7.4.0
+ - packaging==26.2
+ - propcache==0.4.1
+ - pyarrow==21.0.0
+ - pyiceberg==0.11.1
+ - pympler==1.1
+ - python-dateutil==2.8.2
+ - regex==2026.4.4
+ - requests==2.33.1
+ - s3transfer==0.14.0
+ - safetensors==0.7.0
+ - tenacity==8.5.0
+ - tokenizers==0.22.2
+ - transformers==4.57.3
+ - tzdata==2026.2
+ - websocket-client==1.9.0
+ - yarl==1.23.0
+
+Angular / npm packages:
+ - fuse.js@6.5.3
+ - jschardet@3.1.3
+ - rxjs@7.8.1
+
+Agent service npm packages:
+ - @ai-sdk/gateway@2.0.18
+ - @ai-sdk/openai@2.0.79
+ - @ai-sdk/provider-utils@3.0.18
+ - @ai-sdk/provider@2.0.0
+ - @opentelemetry/api@1.9.0
+ - @vercel/oidc@3.0.5
+ - ai@5.0.108
+ - rxjs@7.8.2
+ - typescript@5.9.3
+
+--------------------------------------------------------------------------------
+Dependencies under the MIT License
+--------------------------------------------------------------------------------
+
+Source files derived from third-party MIT-licensed projects:
+ - mbknor-jackson-jsonschema
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/JsonSchemaDraft.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/JsonSchemaGenerator.scala
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaArrayWithUniqueItems.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaBool.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaDefault.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaDescription.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaExamples.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaFormat.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaInject.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaInt.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaOptions.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaString.java
+ common/workflow-operator/src/main/scala/com/kjetland/jackson/jsonSchema/annotations/JsonSchemaTitle.java
+ https://github.com/mbknor/mbknor-jackson-jsonschema
+ - Google Angular formly examples
+ frontend/src/app/common/formly/array.type.ts
+ frontend/src/app/common/formly/object.type.ts
+ frontend/src/app/common/formly/multischema.type.ts
+ frontend/src/app/common/formly/null.type.ts
+ https://angular.io
+ - TypeFox monaco-languageclient
+ pyright-language-service/src/main.ts
+ pyright-language-service/src/language-server-runner.ts
+ pyright-language-service/src/server-commons.ts
+ https://github.com/TypeFox/monaco-languageclient
+ - SVGRepo icons
+ frontend/src/assets/svg/operator-view-result.svg
+ frontend/src/assets/svg/operator-reuse-cache-valid.svg
+ frontend/src/assets/svg/operator-reuse-cache-invalid.svg
+ https://www.svgrepo.com
+
+Scala/Java jars:
+ - co.fs2.fs2-core_2.13-3.12.2.jar
+ - com.konghq.unirest-java-3.14.2.jar
+ - com.liveperson.dropwizard-websockets-1.3.14.jar
+ - io.github.classgraph.classgraph-4.8.157.jar
+ - net.sourceforge.argparse4j.argparse4j-0.8.1.jar
+ - net.sourceforge.argparse4j.argparse4j-0.9.0.jar
+ - org.bouncycastle.bcpkix-jdk18on-1.78.1.jar
+ - org.bouncycastle.bcprov-jdk18on-1.78.1.jar
+ - org.bouncycastle.bcutil-jdk18on-1.78.1.jar
+ - org.checkerframework.checker-qual-3.42.0.jar
+ - org.codehaus.mojo.animal-sniffer-annotations-1.23.jar
+ - org.projectlombok.lombok-1.18.24.jar
+ - org.reactivestreams.reactive-streams-1.0.4.jar
+ - org.slf4j.jcl-over-slf4j-1.7.26.jar
+ - org.slf4j.jul-to-slf4j-1.7.26.jar
+ - org.slf4j.jul-to-slf4j-2.0.12.jar
+ - org.slf4j.slf4j-api-1.7.26.jar
+ - org.slf4j.slf4j-api-2.0.12.jar
+ - org.slf4j.slf4j-api-2.0.16.jar
+ - org.typelevel.cats-core_2.13-2.11.0.jar
+ - org.typelevel.cats-kernel_2.13-2.11.0.jar
+ - org.typelevel.fs2-grpc-runtime_2.13-2.11.0.jar
+
+Python packages:
+ - aioitertools==0.13.0
+ - annotated-types==0.7.0
+ - appdirs==1.4.4
+ - asn1crypto==1.5.1
+ - attrs==26.1.0
+ - betterproto==2.0.0b7
+ - cachetools==6.2.6
+ - charset-normalizer==3.4.7
+ - deprecated==1.2.14
+ - filelock==3.29.0
+ - fonttools==4.62.1
+ - fs==2.4.16
+ - greenlet==3.5.0
+ - h2==4.3.0
+ - hpack==4.1.0
+ - hyperframe==6.1.0
+ - iniconfig==1.1.1
+ - jmespath==1.1.0
+ - loguru==0.7.0
+ - markdown-it-py==4.0.0
+ - mdurl==0.1.2
+ - mmh3==5.2.1
+ - pampy==0.3.0
+ - plotly==5.24.1
+ - pluggy==1.6.0
+ - pydantic==2.13.3
+ - pydantic-core==2.46.3
+ - pyparsing==3.3.2
+ - pyroaring==1.1.0
+ - pytest==7.4.0
+ - pytest-reraise==2.1.2
+ - pytest-timeout==2.2.0
+ - pytz==2026.1.post1
+ - pyyaml==6.0.3
+ - readerwriterlock==1.0.9
+ - rich==14.3.4
+ - ruff==0.14.7
+ - scramp==1.4.8
+ - six==1.17.0
+ - sqlalchemy==2.0.37
+ - strictyaml==1.7.3
+ - typing-inspection==0.4.2
+ - tzlocal==2.1
+ - urllib3==2.6.3
+ - wordcloud==1.9.3
+
+Angular / npm packages:
+ - @abacritt/angularx-social-login@2.3.0
+ - @ali-hm/angular-tree-component@12.0.5
+ - @angular/animations@21.2.10
+ - @angular/cdk@21.2.8
+ - @angular/common@21.2.10
+ - @angular/core@21.2.10
+ - @angular/forms@21.2.10
+ - @angular/platform-browser@21.2.10
+ - @angular/router@21.2.10
+ - @ant-design/colors@7.2.1
+ - @ant-design/fast-color@2.0.6
+ - @ant-design/icons-angular@21.0.0
+ - @auth0/angular-jwt@5.1.0
+ - @babel/runtime@7.29.2
+ - @codingame/monaco-vscode-api@8.0.4
+ - @codingame/monaco-vscode-base-service-override@8.0.4
+ - @codingame/monaco-vscode-configuration-service-override@8.0.4
+ - @codingame/monaco-vscode-editor-api@8.0.4
+ - @codingame/monaco-vscode-environment-service-override@8.0.4
+ - @codingame/monaco-vscode-extensions-service-override@8.0.4
+ - @codingame/monaco-vscode-files-service-override@8.0.4
+ - @codingame/monaco-vscode-host-service-override@8.0.4
+ - @codingame/monaco-vscode-java-default-extension@8.0.4
+ - @codingame/monaco-vscode-languages-service-override@8.0.4
+ - @codingame/monaco-vscode-layout-service-override@8.0.4
+ - @codingame/monaco-vscode-model-service-override@8.0.4
+ - @codingame/monaco-vscode-monarch-service-override@8.0.4
+ - @codingame/monaco-vscode-python-default-extension@8.0.4
+ - @codingame/monaco-vscode-quickaccess-service-override@8.0.4
+ - @codingame/monaco-vscode-r-default-extension@8.0.4
+ - @codingame/monaco-vscode-textmate-service-override@8.0.4
+ - @codingame/monaco-vscode-theme-defaults-default-extension@8.0.4
+ - @codingame/monaco-vscode-theme-service-override@8.0.4
+ - @ctrl/tinycolor@3.6.1
+ - @ngneat/until-destroy@8.1.4
+ - @ngx-formly/core@6.3.12
+ - @ngx-formly/ng-zorro-antd@6.3.12
+ - @vscode/iconv-lite-umd@0.7.0
+ - ajv@8.10.0
+ - backbone@1.4.1
+ - balanced-match@1.0.2
+ - base64-js@1.5.1
+ - brace-expansion@2.1.0
+ - buffer@5.7.1
+ - content-disposition@0.5.4
+ - css-loader@6.11.0
+ - dagre@0.8.5
+ - date-fns@2.30.0
+ - fast-deep-equal@3.1.3
+ - fflate@0.7.4
+ - file-saver@2.0.5
+ - graphlib@2.1.8
+ - html2canvas@1.4.1
+ - java@1.0.0
+ - jquery@3.6.4
+ - json-schema-traverse@1.0.0
+ - jszip@3.10.1
+ - lib0@0.2.117
+ - lodash@4.17.23
+ - lodash-es@4.17.21
+ - marked@17.0.1
+ - mobx@4.14.1
+ - monaco-breakpoints@0.2.0
+ - monaco-editor-wrapper@5.5.3
+ - monaco-languageclient@8.8.3
+ - ng-zorro-antd@21.2.2
+ - ngx-color-picker@12.0.1
+ - ngx-file-drop@16.0.0
+ - ngx-json-viewer@3.2.1
+ - ngx-markdown@21.2.0
+ - papaparse@5.4.1
+ - path-browserify@1.0.1
+ - plotly.js-basic-dist-min@2.29.0
+ - point-in-polygon@1.1.0
+ - python@1.0.0
+ - quill-cursors@3.1.2
+ - r@1.0.0
+ - rbush@4.0.1
+ - read-excel-file@5.7.1
+ - ring-buffer-ts@1.0.3
+ - safe-buffer@5.2.1
+ - style-loader@3.3.4
+ - theme-defaults@1.0.0
+ - underscore@1.13.8
+ - uuid@8.3.2
+ - vscode-jsonrpc@8.2.0
+ - vscode-languageclient@9.0.1
+ - vscode-languageserver-protocol@3.17.5
+ - vscode-languageserver-types@3.17.5
+ - vscode-oniguruma@1.7.0
+ - vscode-textmate@9.0.0
+ - vscode-ws-jsonrpc@3.3.2
+ - y-monaco@0.1.5
+ - y-protocols@1.0.7
+ - y-quill@0.1.5
+ - y-websocket@1.5.4
+ - yjs@13.5.41
+ - zone.js@0.15.1
+
+Agent service npm packages:
+ - @borewit/text-codec@0.1.1
+ - @elysiajs/cors@1.4.0
+ - @pinojs/redact@0.4.0
+ - @sinclair/typebox@0.34.41
+ - @standard-schema/spec@1.0.0
+ - @tokenizer/inflate@0.4.1
+ - @tokenizer/token@0.3.0
+ - @types/bun@1.3.3
+ - @types/node@24.10.1
+ - ajv@8.17.1
+ - atomic-sleep@1.0.0
+ - bun-types@1.3.3
+ - cookie@1.1.1
+ - dagre@0.8.5
+ - debug@4.4.3
+ - elysia@1.4.18
+ - eventsource-parser@3.0.6
+ - exact-mirror@0.2.5
+ - fast-decode-uri-component@1.0.1
+ - fast-deep-equal@3.1.3
+ - file-type@21.1.1
+ - graphlib@2.1.8
+ - json-schema-traverse@1.0.0
+ - lodash@4.18.1
+ - memoirist@0.4.0
+ - ms@2.1.3
+ - on-exit-leak-free@2.1.2
+ - openapi-types@12.1.3
+ - pino-abstract-transport@3.0.0
+ - pino-std-serializers@7.1.0
+ - pino@10.3.1
+ - process-warning@5.0.0
+ - quick-format-unescaped@4.0.4
+ - real-require@0.2.0
+ - require-from-string@2.0.2
+ - safe-stable-stringify@2.5.0
+ - sonic-boom@4.2.1
+ - strtok3@10.3.4
+ - thread-stream@4.0.0
+ - token-types@6.1.1
+ - uint8array-extras@1.5.0
+ - undici-types@7.16.0
+ - zod@3.25.76
+
+--------------------------------------------------------------------------------
+Dependencies under the BSD 3-Clause License
+--------------------------------------------------------------------------------
+
+Scala/Java jars:
+ - com.esotericsoftware.kryo-5.6.2.jar
+ - com.esotericsoftware.kryo.kryo5-5.6.2.jar
+ - com.esotericsoftware.kryo5-5.6.0.jar
+ - com.esotericsoftware.minlog-1.3.1.jar
+ - com.esotericsoftware.reflectasm-1.11.9.jar
+ - com.google.protobuf.protobuf-java-3.25.8.jar
+ - com.google.protobuf.protobuf-java-4.27.1.jar
+ - com.google.re2j.re2j-1.1.jar
+ - com.jcraft.jsch-0.1.55.jar
+ - com.thoughtworks.paranamer.paranamer-2.8.jar
+ - org.fusesource.leveldbjni.leveldbjni-all-1.8.jar
+ - org.jline.jline-3.9.0.jar
+ - org.ow2.asm.asm-8.0.1.jar
+ - org.ow2.asm.asm-9.1.jar
+ - org.ow2.asm.asm-analysis-7.0.jar
+ - org.ow2.asm.asm-commons-7.0.jar
+ - org.ow2.asm.asm-tree-7.0.jar
+ - org.scodec.scodec-bits_2.13-1.1.38.jar
+ - org.threeten.threeten-extra-1.7.1.jar
+
+Python packages:
+ - cached-property==1.5.2
+ - click==8.3.3
+ - contourpy==1.3.3
+ - cycler==0.12.1
+ - fsspec==2025.9.0
+ - grpclib==0.4.9
+ - idna==3.13
+ - jinja2==3.1.6
+ - joblib==1.5.3
+ - kiwisolver==1.5.0
+ - lazy-loader==0.5
+ - markupsafe==3.0.3
+ - mpmath==1.3.0
+ - networkx==3.6.1
+ - numpy==2.1.0
+ - pandas==2.2.3
+ - pg8000==1.31.5
+ - protobuf==4.25.8
+ - psutil==5.9.0
+ - s3fs==2025.9.0
+ - scikit-image==0.25.2
+ - scikit-learn==1.5.0
+ - scipy==1.17.1
+ - sympy==1.14.0
+ - threadpoolctl==3.6.0
+ - tifffile==2026.4.11
+ - torch==2.8.0
+ - zstandard==0.25.0
+
+Angular / npm packages:
+ - d3-shape@2.1.0
+ - ieee754@1.2.1
+ - quill@1.3.7
+
+Agent service npm packages:
+ - fast-uri@3.1.0
+ - ieee754@1.2.1
+ - json-schema@0.4.0
+
+--------------------------------------------------------------------------------
+Dependencies under the BSD 2-Clause License
+--------------------------------------------------------------------------------
+
+Scala/Java jars:
+ - com.github.luben.zstd-jni-1.5.0-1.jar
+ - com.github.marianobarrios.lbmq-0.6.0.jar
+ - dnsjava.dnsjava-2.1.7.jar
+ - org.codehaus.woodstox.stax2-api-4.2.1.jar
+ - org.postgresql.postgresql-42.7.4.jar
+
+Python packages:
+ - imageio==2.37.3
+ - praw==7.6.1
+ - prawcore==2.4.0
+ - pybase64==1.3.2
+ - pygments==2.20.0
+ - update-checker==0.18.0
+ - wrapt==1.17.3
+
+Angular / npm packages:
+ - uri-js@4.4.1
+
+--------------------------------------------------------------------------------
+Dependencies under the ISC License
+--------------------------------------------------------------------------------
+
+Scala/Java jars:
+ - org.mindrot.jbcrypt-0.4.jar
+
+Angular / npm packages:
+ - concaveman@2.0.0
+ - d3-path@2.0.0
+ - minimatch@5.1.9
+ - quickselect@3.0.0
+ - tinyqueue@2.0.3
+
+Agent service npm packages:
+ - split2@4.2.0
+
+--------------------------------------------------------------------------------
+Dependencies under the Mozilla Public License, Version 2.0
+--------------------------------------------------------------------------------
+
+Python packages:
+ - bidict==0.22.0
+ - certifi==2026.4.22
+ - tqdm==4.67.3
+
+Angular / npm packages:
+ - jointjs@3.5.4
+
+--------------------------------------------------------------------------------
+Dependencies under the Eclipse Public License, Version 2.0 (some are dual
+licensed with GPL-2.0 with Classpath Exception)
+--------------------------------------------------------------------------------
+
+Scala/Java jars:
+ - jakarta.annotation.jakarta.annotation-api-2.1.1.jar
+ - jakarta.annotation.jakarta.annotation-api-3.0.0.jar
+ - jakarta.el.jakarta.el-api-4.0.0.jar
+ - jakarta.servlet.jakarta.servlet-api-5.0.0.jar
+ - jakarta.ws.rs.jakarta.ws.rs-api-3.0.0.jar
+ - jakarta.ws.rs.jakarta.ws.rs-api-3.1.0.jar
+ - javax.ws.rs.javax.ws.rs-api-2.1.1.jar
+ - org.glassfish.hk2.external.aopalliance-repackaged-3.0.6.jar
+ - org.glassfish.hk2.hk2-api-3.0.6.jar
+ - org.glassfish.hk2.hk2-locator-3.0.3.jar
+ - org.glassfish.hk2.hk2-utils-3.0.6.jar
+ - org.glassfish.hk2.osgi-resource-locator-1.0.3.jar
+ - org.glassfish.jakarta.el-4.0.2.jar
+ - org.glassfish.jersey.containers.jersey-container-servlet-3.0.12.jar
+ - org.glassfish.jersey.containers.jersey-container-servlet-core-3.0.12.jar
+ - org.glassfish.jersey.core.jersey-client-3.0.12.jar
+ - org.glassfish.jersey.core.jersey-common-3.0.12.jar
+ - org.glassfish.jersey.core.jersey-server-3.0.12.jar
+ - org.glassfish.jersey.ext.jersey-bean-validation-3.0.12.jar
+ - org.glassfish.jersey.ext.jersey-metainf-services-3.0.12.jar
+ - org.glassfish.jersey.inject.jersey-hk2-3.0.12.jar
+ - org.jgrapht.jgrapht-core-1.4.0.jar
+
+--------------------------------------------------------------------------------
+Dependencies under the Eclipse Public License, Version 1.0 (Logback is dual
+licensed with LGPL-2.1)
+--------------------------------------------------------------------------------
+
+Scala/Java jars:
+ - ch.qos.logback.logback-access-1.2.3.jar
+ - ch.qos.logback.logback-access-1.4.14.jar
+ - ch.qos.logback.logback-classic-1.2.3.jar
+ - ch.qos.logback.logback-classic-1.4.14.jar
+ - ch.qos.logback.logback-core-1.2.3.jar
+ - ch.qos.logback.logback-core-1.4.14.jar
+
+--------------------------------------------------------------------------------
+Dependencies under the Common Development and Distribution License (CDDL)
+(some are dual licensed with GPL-2.0 with Classpath Exception)
+--------------------------------------------------------------------------------
+
+CDDL 1.0
+~~~~~~~~
+
+Scala/Java jars:
+ - com.sun.mail.javax.mail-1.6.2.jar
+ - javax.activation.activation-1.1.1.jar
+ - javax.annotation.javax.annotation-api-1.3.2.jar
+ - javax.servlet.javax.servlet-api-3.1.0.jar
+ - javax.ws.rs.jsr311-api-1.1.1.jar
+ - org.glassfish.hk2.external.javax.inject-2.5.0-b32.jar
+ - org.glassfish.hk2.hk2-api-2.5.0-b32.jar
+ - org.glassfish.hk2.hk2-locator-2.5.0-b32.jar
+ - org.glassfish.hk2.hk2-utils-2.5.0-b32.jar
+ - org.glassfish.hk2.osgi-resource-locator-1.0.1.jar
+ - org.glassfish.javax.el-3.0.0.jar
+ - org.glassfish.jersey.bundles.repackaged.jersey-guava-2.25.1.jar
+ - org.glassfish.jersey.connectors.jersey-apache-connector-2.25.1.jar
+ - org.glassfish.jersey.containers.jersey-container-servlet-2.25.1.jar
+ - org.glassfish.jersey.containers.jersey-container-servlet-core-2.25.1.jar
+ - org.glassfish.jersey.core.jersey-client-2.25.1.jar
+ - org.glassfish.jersey.core.jersey-common-2.25.1.jar
+ - org.glassfish.jersey.core.jersey-server-2.25.1.jar
+ - org.glassfish.jersey.ext.jersey-bean-validation-2.25.1.jar
+ - org.glassfish.jersey.ext.jersey-metainf-services-2.25.1.jar
+ - org.glassfish.jersey.ext.rx.jersey-rx-client-2.25.1.jar
+ - org.glassfish.jersey.media.jersey-media-jaxb-2.25.1.jar
+
+CDDL 1.1
+~~~~~~~~
+
+Scala/Java jars:
+ - com.sun.jersey.contribs.jersey-guice-1.19.jar
+ - javax.websocket.javax.websocket-api-1.0.jar
+ - javax.websocket.javax.websocket-client-api-1.0.jar
+ - javax.xml.bind.jaxb-api-2.3.0.jar
+ - org.glassfish.hk2.external.aopalliance-repackaged-2.5.0-b32.jar
+
+--------------------------------------------------------------------------------
+Dependencies under the Eclipse Distribution License, Version 1.0
+--------------------------------------------------------------------------------
+
+Scala/Java jars:
+ - com.sun.activation.jakarta.activation-2.0.0.jar
+ - com.sun.activation.jakarta.activation-2.0.1.jar
+ - jakarta.activation.jakarta.activation-api-1.2.1.jar
+ - jakarta.activation.jakarta.activation-api-2.1.0.jar
+ - jakarta.xml.bind.jakarta.xml.bind-api-3.0.0.jar
+ - jakarta.xml.bind.jakarta.xml.bind-api-3.0.1.jar
+ - org.eclipse.collections.eclipse-collections-11.1.0.jar
+ - org.eclipse.collections.eclipse-collections-api-11.1.0.jar
+ - org.eclipse.jgit.org.eclipse.jgit-5.13.0.202109080827-r.jar
+
+--------------------------------------------------------------------------------
+Dependencies under the Python Software Foundation License
+--------------------------------------------------------------------------------
+
+Python packages:
+ - aiohappyeyeballs==2.6.1
+ - matplotlib==3.10.9
+ - typing-extensions==4.14.1
+
+--------------------------------------------------------------------------------
+Dependencies under the MIT-CMU License
+--------------------------------------------------------------------------------
+
+Python packages:
+ - pillow==12.1.1
+
+--------------------------------------------------------------------------------
+Dependencies under the BSD Zero Clause License
+--------------------------------------------------------------------------------
+
+Agent service npm packages:
+ - tslib@2.8.1
+
+--------------------------------------------------------------------------------
+Dependencies in the Public Domain (CC0)
+--------------------------------------------------------------------------------
+
+Scala/Java jars:
+ - aopalliance.aopalliance-1.0.jar
+ - org.reactivestreams.reactive-streams-1.0.3.jar
+ - org.tukaani.xz-1.9.jar
+
+--------------------------------------------------------------------------------
+Dependencies under the Unlicense
+--------------------------------------------------------------------------------
+
+Angular / npm packages:
+ - robust-predicates@3.0.3
+
+Individual jars may contain their own META-INF/LICENSE and META-INF/NOTICE
+files that apply to their specific contents; those files continue to govern
+the use of those components.
diff --git a/NOTICE b/NOTICE
new file mode 100644
index 00000000000..d0729cd83cd
--- /dev/null
+++ b/NOTICE
@@ -0,0 +1,5 @@
+Apache Texera (Incubating)
+Copyright 2025 The Apache Software Foundation
+
+This product includes software developed at
+The Apache Software Foundation (http://www.apache.org/).
diff --git a/NOTICE-binary b/NOTICE-binary
new file mode 100644
index 00000000000..a350e8896c7
--- /dev/null
+++ b/NOTICE-binary
@@ -0,0 +1,2037 @@
+Apache Texera (Incubating)
+Copyright 2025-2026 The Apache Software Foundation
+
+This product includes software developed at
+The Apache Software Foundation (http://www.apache.org/).
+
+--------------------------------------------------------------------------------
+Apache Hadoop
+--------------------------------------------------------------------------------
+
+Apache Hadoop
+Copyright 2006 and onwards The Apache Software Foundation.
+
+Export Control Notice
+---------------------
+
+This distribution includes cryptographic software. The country in which
+you currently reside may have restrictions on the import, possession, use,
+and/or re-export to another country, of encryption software. BEFORE using
+any encryption software, please check your country's laws, regulations and
+policies concerning the import, possession, or use, and re-export of
+encryption software, to see if this is permitted. See
+ for more information.
+
+The U.S. Government Department of Commerce, Bureau of Industry and
+Security (BIS), has classified this software as Export Commodity Control
+Number (ECCN) 5D002.C.1, which includes information security software
+using or performing cryptographic functions with asymmetric algorithms.
+The form and manner of this Apache Software Foundation distribution makes
+it eligible for export under the License Exception ENC Technology Software
+Unrestricted (TSU) exception (see the BIS Export Administration
+Regulations, Section 740.13) for both object code and source code.
+
+The following provides more details on the included cryptographic software:
+
+ This software uses the SSL libraries from the Jetty project written
+ by mortbay.org.
+
+ Hadoop Yarn Server Web Proxy uses the BouncyCastle Java cryptography
+ APIs written by the Legion of the Bouncy Castle Inc.
+
+--------------------------------------------------------------------------------
+Apache Lucene
+--------------------------------------------------------------------------------
+
+Apache Lucene
+Copyright 2001-2021 The Apache Software Foundation
+
+Includes software from other Apache Software Foundation projects,
+including, but not limited to Apache Ant, Apache Jakarta Regexp,
+Apache Commons, and Apache Xerces.
+
+ICU4J (under analysis/icu) is licensed under an MIT-style license and
+Copyright (c) 1995-2008 International Business Machines Corporation and
+others.
+
+Some data files (under analysis/icu/src/data) are derived from Unicode
+data such as the Unicode Character Database. See
+http://unicode.org/copyright.html for more details.
+
+Brics Automaton (under core/src/java/org/apache/lucene/util/automaton) is
+BSD-licensed, created by Anders Moller. See http://www.brics.dk/automaton/
+
+The levenshtein automata tables (under core/src/java/org/apache/lucene/util/automaton)
+were automatically generated with the moman/finenight FSA library, created
+by Jean-Philippe Barrette-LaPierre. This library is available under an
+MIT license.
+
+The class org.apache.lucene.util.WeakIdentityMap was derived from the
+Apache CXF project and is Apache License 2.0.
+
+The class org.apache.lucene.util.compress.LZ4 is a Java rewrite of the LZ4
+compression library (https://github.com/lz4/lz4/tree/dev/lib) that is
+licensed under the 2-clause BSD license.
+
+The Google Code Prettify is Apache License 2.0.
+
+This product includes code (JaspellTernarySearchTrie) from Java Spelling
+Checking Package (jaspell): http://jaspell.sourceforge.net/ (BSD License).
+
+The snowball stemmers (in analysis/common/src/java/net/sf/snowball) were
+developed by Martin Porter and Richard Boulton.
+
+The KStem stemmer in analysis/common/src/org/apache/lucene/analysis/en was
+developed by Bob Krovetz and Sergio Guzman-Lara (CIIR-UMass Amherst) under
+the BSD license.
+
+Arabic, Persian, Romanian, Bulgarian, Hindi and Bengali analyzer stopword
+lists are BSD-licensed and were created by Jacques Savoy.
+
+The German, Spanish, Finnish, French, Hungarian, Italian, Portuguese,
+Russian and Swedish light stemmers are based on BSD-licensed reference
+implementations created by Jacques Savoy and Ljiljana Dolamic.
+
+The Stempel analyzer includes BSD-licensed software developed by the
+Egothor project (http://egothor.sf.net/), created by Leo Galambos,
+Martin Kvapil, and Edmond Nolan.
+
+The Polish analyzer stopword list is BSD-licensed and was created by the
+Carrot2 project.
+
+The SmartChineseAnalyzer source code (smartcn) was provided by
+Xiaoping Gao and copyright 2009 by www.imdict.net.
+
+WordBreakTestUnicode_*.java is derived from Unicode data such as the
+Unicode Character Database.
+
+--------------------------------------------------------------------------------
+Apache Pekko
+--------------------------------------------------------------------------------
+
+Apache Pekko
+Copyright 2022-2025 The Apache Software Foundation
+
+This product contains significant parts that were originally based on
+software from Lightbend (Akka ).
+Copyright (C) 2009-2022 Lightbend Inc.
+
+Apache Pekko is derived from Akka 2.6.x, the last version that was
+distributed under the Apache License, Version 2.0.
+
+pekko-actor contains MurmurHash.scala, scala-collection-compat, and code
+from scala-library, each modified by the Scala-Lang team under an Apache
+2.0 license.
+
+ Scala
+ Copyright (c) 2002-2023 EPFL
+ Copyright (c) 2011-2023 Lightbend, Inc.
+
+ Scala includes software developed at LAMP/EPFL (https://lamp.epfl.ch/)
+ and Lightbend, Inc. (https://www.lightbend.com/).
+
+pekko-actor contains code from Netty, released under an Apache 2.0
+license. Copyright 2014 The Netty Project (https://netty.io/).
+
+pekko-actor contains code from java-uuid-generator
+(https://github.com/cowtowncoder/java-uuid-generator) in
+`org.apache.pekko.util.UUIDComparator.scala`, released under an Apache 2.0
+license. Java UUID generator library has been written by Tatu Saloranta
+(tatu.saloranta@iki.fi).
+
+--------------------------------------------------------------------------------
+Apache Parquet
+--------------------------------------------------------------------------------
+
+Apache Parquet MR
+Copyright 2014-2024 The Apache Software Foundation
+
+This product includes code from Apache Avro.
+
+ Apache Avro
+ Copyright 2010-2024 The Apache Software Foundation
+
+--------------------------------------------------------------------------------
+Apache Iceberg
+--------------------------------------------------------------------------------
+
+
+Apache Iceberg
+Copyright 2017-2024 The Apache Software Foundation
+
+This product includes software developed at
+The Apache Software Foundation (http://www.apache.org/).
+
+--------------------------------------------------------------------------------
+
+This project includes code from Kite, developed at Cloudera, Inc. with
+the following copyright notice:
+
+| Copyright 2013 Cloudera 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.
+
+ Apache Arrow (arrow-format, arrow-memory-core, arrow-memory-netty,
+ arrow-vector, flight-core, flight-grpc)
+ Copyright 2016-2023 The Apache Software Foundation
+
+ Apache Avro
+ Copyright 2009-2024 The Apache Software Foundation
+
+ Apache Commons BeanUtils
+ Copyright 2000-2019 The Apache Software Foundation
+
+ Apache Commons CLI
+ Copyright 2001-2022 The Apache Software Foundation
+
+ Apache Commons Codec
+ Copyright 2002-2024 The Apache Software Foundation
+
+ Apache Commons Collections (3.x and 4.x)
+ Copyright 2001-2024 The Apache Software Foundation
+
+ Apache Commons Compress
+ Copyright 2002-2024 The Apache Software Foundation
+
+ Apache Commons Configuration
+ Copyright 2001-2024 The Apache Software Foundation
+
+ Apache Commons IO
+ Copyright 2002-2024 The Apache Software Foundation
+
+ Apache Commons JCS
+ Copyright 2002-2024 The Apache Software Foundation
+
+ Apache Commons Lang (2.x and 3.x)
+ Copyright 2001-2024 The Apache Software Foundation
+
+ Apache Commons Logging
+ Copyright 2003-2014 The Apache Software Foundation
+
+ Apache Commons Math
+ Copyright 2001-2016 The Apache Software Foundation
+
+ Apache Commons Net
+ Copyright 2001-2023 The Apache Software Foundation
+
+ Apache Commons Pool
+ Copyright 2001-2024 The Apache Software Foundation
+
+ Apache Commons Text
+ Copyright 2014-2024 The Apache Software Foundation
+
+ Apache Commons VFS
+ Copyright 2002-2024 The Apache Software Foundation
+
+ Apache Curator
+ Copyright 2011-2024 The Apache Software Foundation
+
+ Apache HttpComponents (httpclient, httpcore, httpasyncclient,
+ httpclient5, httpcore5, httpcore5-h2, httpmime)
+ Copyright 1999-2024 The Apache Software Foundation
+ Apache HTrace (Incubating)
+ Copyright 2016-2017 The Apache Software Foundation
+
+ Apache Iceberg
+ Copyright 2017-2024 The Apache Software Foundation
+
+ Apache Kerby (kerb-* and kerby-* subprojects)
+ Copyright 2014-2017 The Apache Software Foundation
+
+ Apache Maven (many subprojects including wagon-*)
+ Copyright 2001-2024 The Apache Software Foundation
+
+ Apache ORC
+ Copyright 2013-2024 The Apache Software Foundation
+
+ Apache Yetus
+ Copyright 2015-2023 The Apache Software Foundation
+
+ Apache ZooKeeper
+ Copyright 2008-2024 The Apache Software Foundation
+
+ Apache log4j 1.2 / reload4j
+ Copyright 2007 The Apache Software Foundation
+
+--------------------------------------------------------------------------------
+Netty
+--------------------------------------------------------------------------------
+
+The Netty Project
+Copyright 2011-2024 The Netty Project (https://netty.io/).
+
+This product contains the extensions to Java Collections Framework derived
+from the works by JSR-166 EG, Doug Lea, and Jason T. Greene (Public
+Domain).
+
+This product contains a modified version of Robert Harder's Public Domain
+Base64 Encoder and Decoder.
+
+This product contains a modified version of 'JZlib', a re-implementation
+of zlib in pure Java (BSD-style license,
+http://www.jcraft.com/jzlib/).
+
+This product contains a modified version of 'Webbit' (BSD License,
+https://github.com/joewalnes/webbit).
+
+This product optionally depends on 'Protocol Buffers' (New BSD License,
+http://code.google.com/p/protobuf/), 'Bouncy Castle Crypto APIs' (MIT
+License, http://www.bouncycastle.org/), 'SLF4J' (MIT License,
+http://www.slf4j.org/), 'Apache Commons Logging' (Apache License 2.0),
+'Apache Log4J' (Apache License 2.0), 'JBoss Logging' (GNU LGPL 2.1), and
+'Apache Felix' (Apache License 2.0).
+
+--------------------------------------------------------------------------------
+Eclipse Jetty
+--------------------------------------------------------------------------------
+
+Jetty Web Container
+Copyright 1995-2018 Mort Bay Consulting Pty Ltd.
+
+The Jetty Web Container is Copyright Mort Bay Consulting Pty Ltd unless
+otherwise noted. Jetty is dual licensed under both the Apache 2.0 License
+and the Eclipse Public 1.0 License; Texera redistributes it under the
+Apache 2.0 terms.
+
+Jetty bundles select artifacts under secondary licenses:
+ * Eclipse Public License: org.eclipse.jetty.orbit:org.eclipse.jdt.core,
+ javax.security.auth.message (EPL + ASL2),
+ javax.mail.glassfish (EPL + CDDL 1.0)
+ * CDDL + GPLv2 with classpath exception: javax.servlet:javax.servlet-api,
+ javax.annotation:javax.annotation-api,
+ javax.transaction:javax.transaction-api,
+ javax.websocket:javax.websocket-api
+ * OW2 license: org.ow2.asm:asm-commons, org.ow2.asm:asm
+ * MortBay ASL2: org.mortbay.jasper:apache-jsp, apache-el (based on
+ selected classes from Apache Tomcat)
+
+The UnixCrypt.java code implements one-way cryptography used by Unix
+systems for simple password protection. Copyright 1996 Aki Yoshida,
+modified April 2001 by Iris Van den Broeke, Daniel Deville.
+
+--------------------------------------------------------------------------------
+Jackson (FasterXML)
+--------------------------------------------------------------------------------
+
+Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and
+has been in development since 2007. It is currently developed by a
+community of developers.
+
+Copyright 2007- Tatu Saloranta (tatu.saloranta@iki.fi)
+
+Jackson 2.x core and extension components are licensed under Apache
+License 2.0. This attribution applies to jackson-core, jackson-databind,
+jackson-annotations, and every jackson-datatype-*, jackson-module-*,
+jackson-dataformat-*, and jackson-jaxrs-* artifact bundled in this
+distribution.
+
+Java ClassMate library (com.fasterxml:classmate) was originally written
+by Tatu Saloranta (tatu.saloranta@iki.fi), with contributions from
+Brian Langel.
+
+--------------------------------------------------------------------------------
+Google Guice
+--------------------------------------------------------------------------------
+
+Google Guice - Core Library (and guice-servlet extension)
+Copyright 2006-2015 Google, Inc.
+
+--------------------------------------------------------------------------------
+AWS SDK for Java 2.0
+--------------------------------------------------------------------------------
+
+AWS SDK for Java 2.0
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+This product includes software developed by Amazon Technologies, Inc
+(http://www.amazon.com/).
+
+The AWS SDK bundles the following third-party works:
+ * XML parsing and utility functions from JetS3t
+ Copyright 2006-2009 James Murty.
+ * PKCS#1 PEM encoded private key parsing and utility functions from
+ oauth.googlecode.com - Copyright 1998-2010 AOL Inc.
+ * Apache Commons Lang (https://github.com/apache/commons-lang)
+ * Netty Reactive Streams
+ (https://github.com/playframework/netty-reactive-streams)
+ * Jackson-core (https://github.com/FasterXML/jackson-core), shaded as
+ software.amazon.awssdk:third-party-jackson-core
+ * Jackson-dataformat-cbor
+ (https://github.com/FasterXML/jackson-dataformats-binary)
+
+Required Apache Commons Lang attribution:
+ Apache Commons Lang
+ Copyright 2001-2020 The Apache Software Foundation
+
+--------------------------------------------------------------------------------
+Kryo
+--------------------------------------------------------------------------------
+
+Kryo bundles Objenesis (see below).
+
+--------------------------------------------------------------------------------
+Objenesis
+--------------------------------------------------------------------------------
+
+Objenesis
+Copyright 2006-2024 Joe Walnes, Henri Tremblay, Leonardo Mesquita
+
+(NOTICE file corresponding to section 4d of the Apache License, Version
+2.0, in this case for Objenesis.)
+
+--------------------------------------------------------------------------------
+Jasypt
+--------------------------------------------------------------------------------
+
+Copyright (c) 2007-2010, The JASYPT team (http://www.jasypt.org)
+
+This distribution includes cryptographic software. The country in which
+you currently reside may have restrictions on the import, possession, use,
+and/or re-export to another country, of encryption software. BEFORE using
+any encryption software, please check your country's laws, regulations and
+policies concerning the import, possession, use, or re-export of
+encryption software.
+
+The U.S. Government Department of Commerce, Bureau of Industry and
+Security (BIS), has classified this software as Export Commodity Control
+Number (ECCN) 5D002.C.1. The PBE Encryption facilities require the Java
+Cryptography Extensions.
+
+Jasypt includes the ICU License (ICU 1.8.1 and later):
+
+ Copyright (c) 1995-2006 International Business Machines Corporation and
+ others. All rights reserved.
+
+ Permission is hereby granted, free of charge, to any person obtaining a
+ copy of this software and associated documentation files (the
+ "Software"), to deal in the Software without restriction, including
+ without limitation the rights to use, copy, modify, merge, publish,
+ distribute, and/or sell copies of the Software, and to permit persons
+ to whom the Software is furnished to do so, provided that the above
+ copyright notice(s) and this permission notice appear in all copies of
+ the Software and that both the above copyright notice(s) and this
+ permission notice appear in supporting documentation.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+ OR IMPLIED.
+
+ Except as contained in this notice, the name of a copyright holder
+ shall not be used in advertising or otherwise to promote the sale, use
+ or other dealings in this Software without prior written authorization
+ of the copyright holder.
+
+--------------------------------------------------------------------------------
+Joda-Time
+--------------------------------------------------------------------------------
+
+This product includes software developed by Joda.org (https://www.joda.org/).
+
+--------------------------------------------------------------------------------
+Jackson core (verbatim upstream NOTICE)
+--------------------------------------------------------------------------------
+
+# Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has
+been in development since 2007.
+It is currently developed by a community of developers.
+
+## Copyright
+
+Copyright 2007-, Tatu Saloranta (tatu.saloranta@iki.fi)
+
+## Licensing
+
+Jackson 2.x core and extension components are licensed under Apache License 2.0
+To find the details that apply to this artifact see the accompanying LICENSE file.
+
+## Credits
+
+A list of contributors may be found from CREDITS(-2.x) file, which is included
+in some artifacts (usually source distributions); but is always available
+from the source code management (SCM) system project uses.
+
+## FastDoubleParser
+
+jackson-core bundles a shaded copy of FastDoubleParser .
+That code is available under an MIT license
+under the following copyright.
+
+Copyright © 2023 Werner Randelshofer, Switzerland. MIT License.
+
+See FastDoubleParser-NOTICE for details of other source code included in FastDoubleParser
+and the licenses and copyrights that apply to that code.
+
+# FastDoubleParser
+
+This is a Java port of Daniel Lemire's fast_float project.
+This project provides parsers for double, float, BigDecimal and BigInteger values.
+
+## Copyright
+
+Copyright © 2024 Werner Randelshofer, Switzerland.
+
+## Licensing
+
+This code is licensed under MIT License.
+https://github.com/wrandelshofer/FastDoubleParser/blob/522be16e145f43308c43b23094e31d5efcaa580e/LICENSE
+(The file 'LICENSE' is included in the sources and classes Jar files that are released by this project
+- as is required by that license.)
+
+Some portions of the code have been derived from other projects.
+All these projects require that we include a copyright notice, and some require that we also include some text of their
+license file.
+
+fast_double_parser, Copyright (c) 2022 Daniel Lemire. BSL License.
+https://github.com/lemire/fast_double_parser
+https://github.com/lemire/fast_double_parser/blob/07d9189a8fb815fe800cb15ca022e7a07093236e/LICENSE.BSL
+(The file 'thirdparty-LICENSE' is included in the sources and classes Jar files that are released by this project
+- as is required by that license.)
+
+fast_float, Copyright (c) 2021 The fast_float authors. MIT License.
+https://github.com/fastfloat/fast_float
+https://github.com/fastfloat/fast_float/blob/cc1e01e9eee74128e48d51488a6b1df4a767a810/LICENSE-MIT
+(The file 'thirdparty-LICENSE' is included in the sources and classes Jar files that are released by this project
+- as is required by that license.)
+
+bigint, Copyright 2020 Tim Buktu. 2-clause BSD License.
+https://github.com/tbuktu/bigint/tree/floatfft
+https://github.com/tbuktu/bigint/blob/617c8cd8a7c5e4fb4d919c6a4d11e2586107f029/LICENSE
+https://github.com/wrandelshofer/FastDoubleParser/blob/39e123b15b71f29a38a087d16a0bc620fc879aa6/bigint-LICENSE
+(We only use those portions of the bigint project that can be licensed under 2-clause BSD License.)
+(The file 'thirdparty-LICENSE' is included in the sources and classes Jar files that are released by this project
+- as is required by that license.)
+
+--------------------------------------------------------------------------------
+Jackson modules and datatypes
+--------------------------------------------------------------------------------
+
+# Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has
+been in development since 2007.
+It is currently developed by a community of developers.
+
+## Copyright
+
+Copyright 2007-, Tatu Saloranta (tatu.saloranta@iki.fi)
+
+## Licensing
+
+Jackson 2.x core and extension components are licensed under Apache License 2.0
+To find the details that apply to this artifact see the accompanying LICENSE file.
+
+## Credits
+
+A list of contributors may be found from CREDITS(-2.x) file, which is included
+in some artifacts (usually source distributions); but is always available
+from the source code management (SCM) system project uses.
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has
+been in development since 2007.
+It is currently developed by a community of developers, as well as supported
+commercially by FasterXML.com.
+
+## Licensing
+
+Jackson core and extension components (as well their dependencies) may be licensed under
+different licenses.
+To find the details that apply to this artifact see the accompanying LICENSE file.
+For more information, including possible other licensing options, contact
+FasterXML.com (http://fasterxml.com).
+
+## Credits
+
+A list of contributors may be found from CREDITS file, which is included
+in some artifacts (usually source distributions); but is always available
+from the source code management (SCM) system project uses.
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has
+been in development since 2007.
+It is currently developed by a community of developers, as well as supported
+commercially by FasterXML.com.
+
+## Licensing
+
+Jackson core and extension components may be licensed under different licenses.
+To find the details that apply to this artifact see the accompanying LICENSE file.
+For more information, including possible other licensing options, contact
+FasterXML.com (http://fasterxml.com).
+
+## Credits
+
+A list of contributors may be found from CREDITS file, which is included
+in some artifacts (usually source distributions); but is always available
+from the source code management (SCM) system project uses.
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has
+been in development since 2007.
+It is currently developed by a community of developers.
+
+## Copyright
+
+Copyright 2007-, Tatu Saloranta (tatu.saloranta@iki.fi)
+
+## Licensing
+
+Jackson components are licensed under Apache (Software) License, version 2.0,
+as per accompanying LICENSE file.
+
+## Credits
+
+A list of contributors may be found from CREDITS file, which is included
+in some artifacts (usually source distributions); but is always available
+from the source code management (SCM) system project uses.
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has
+been in development since 2007.
+It is currently developed by a community of developers, as well as supported
+commercially by FasterXML.com.
+
+## Licensing
+
+Jackson core and extension components may licensed under different licenses.
+To find the details that apply to this artifact see the accompanying LICENSE file.
+For more information, including possible other licensing options, contact
+FasterXML.com (http://fasterxml.com).
+
+## Credits
+
+A list of contributors may be found from CREDITS file, which is included
+in some artifacts (usually source distributions); but is always available
+from the source code management (SCM) system project uses.
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has
+been in development since 2007.
+It is currently developed by a community of developers.
+
+## Licensing
+
+Jackson components are licensed under Apache (Software) License, version 2.0,
+as per accompanying LICENSE file.
+
+## Credits
+
+A list of contributors may be found from CREDITS file, which is included
+in some artifacts (usually source distributions); but is always available
+from the source code management (SCM) system project uses.
+
+--------------------------------------------------------------------------------
+Eclipse Jetty 11.0
+--------------------------------------------------------------------------------
+
+Notices for Eclipse Jetty
+=========================
+This content is produced and maintained by the Eclipse Jetty project.
+
+Project home: https://eclipse.dev/jetty/
+
+Trademarks
+----------
+Eclipse Jetty, and Jetty are trademarks of the Eclipse Foundation.
+
+Copyright
+---------
+All contributions are the property of the respective authors or of
+entities to which copyright has been assigned by the authors (eg. employer).
+
+Declared Project Licenses
+-------------------------
+This artifacts of this project are made available under the terms of:
+
+ * the Eclipse Public License v2.0
+ https://www.eclipse.org/legal/epl-2.0
+ SPDX-License-Identifier: EPL-2.0
+
+ or
+
+ * the Apache License, Version 2.0
+ https://www.apache.org/licenses/LICENSE-2.0
+ SPDX-License-Identifier: Apache-2.0
+
+The following dependencies are EPL.
+ * org.eclipse.jetty.orbit:org.eclipse.jdt.core
+
+The following dependencies are EPL and ASL2.
+ * org.eclipse.jetty.orbit:javax.security.auth.message
+
+The following dependencies are EPL and CDDL 1.0.
+ * org.eclipse.jetty.orbit:javax.mail.glassfish
+
+The following dependencies are CDDL + GPLv2 with classpath exception.
+https://glassfish.dev.java.net/nonav/public/CDDL+GPL.html
+
+ * jakarta.servlet:jakarta.servlet-api
+ * javax.annotation:javax.annotation-api
+ * javax.transaction:javax.transaction-api
+ * javax.websocket:javax.websocket-api
+
+The following dependencies are licensed by the OW2 Foundation according to the
+terms of http://asm.ow2.org/license.html
+
+ * org.ow2.asm:asm-commons
+ * org.ow2.asm:asm
+
+The following dependencies are ASL2 licensed.
+
+ * org.apache.taglibs:taglibs-standard-spec
+ * org.apache.taglibs:taglibs-standard-impl
+
+The following dependencies are ASL2 licensed. Based on selected classes from
+following Apache Tomcat jars, all ASL2 licensed.
+
+ * org.mortbay.jasper:apache-jsp
+ * org.apache.tomcat:tomcat-jasper
+ * org.apache.tomcat:tomcat-juli
+ * org.apache.tomcat:tomcat-jsp-api
+ * org.apache.tomcat:tomcat-el-api
+ * org.apache.tomcat:tomcat-jasper-el
+ * org.apache.tomcat:tomcat-api
+ * org.apache.tomcat:tomcat-util-scan
+ * org.apache.tomcat:tomcat-util
+ * org.mortbay.jasper:apache-el
+ * org.apache.tomcat:tomcat-jasper-el
+ * org.apache.tomcat:tomcat-el-api
+
+The following artifacts are CDDL + GPLv2 with classpath exception.
+https://glassfish.dev.java.net/nonav/public/CDDL+GPL.html
+
+ * org.eclipse.jetty.toolchain:jetty-schemas
+
+Cryptography
+------------
+Content may contain encryption software. The country in which you are currently
+may have restrictions on the import, possession, and use, and/or re-export to
+another country, of encryption software. BEFORE using any encryption software,
+please check the country's laws, regulations and policies concerning the import,
+possession, or use, and re-export of encryption software, to see if this is
+permitted.
+
+The UnixCrypt.java code implements the one way cryptography used by
+Unix systems for simple password protection. Copyright 1996 Aki Yoshida,
+modified April 2001 by Iris Van den Broeke, Daniel Deville.
+Permission to use, copy, modify and distribute UnixCrypt
+for non-commercial or commercial purposes and without fee is
+granted provided that the copyright notice appears in all copies.
+
+--------------------------------------------------------------------------------
+Apache Parquet (per-component supplementary notices)
+--------------------------------------------------------------------------------
+
+
+Apache Parquet MR (Incubating)
+Copyright 2014-2015 The Apache Software Foundation
+
+This product includes software developed at
+The Apache Software Foundation (http://www.apache.org/).
+
+--------------------------------------------------------------------------------
+
+This product includes code from Apache Avro, which includes the following in
+its NOTICE file:
+
+ Apache Avro
+ Copyright 2010-2015 The Apache Software Foundation
+
+ This product includes software developed at
+ The Apache Software Foundation (http://www.apache.org/).
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Jackson JSON processor
+
+Jackson is a high-performance, Free/Open Source JSON processing library.
+It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has
+been in development since 2007.
+It is currently developed by a community of developers.
+
+## Licensing
+
+Jackson 2.x core and extension components are licensed under Apache License 2.0
+To find the details that apply to this artifact see the accompanying LICENSE file.
+
+## Credits
+
+A list of contributors may be found from CREDITS(-2.x) file, which is included
+in some artifacts (usually source distributions); but is always available
+from the source code management (SCM) system project uses.
+
+--------------------------------------------------------------------------------
+R2DBC SPI
+--------------------------------------------------------------------------------
+
+Reactive Relational Database Connectivity
+
+Copyright 2017-2021 the original author or authors.
+
+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
+
+ https://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.
+
+--------------------------------------------------------------------------------
+Joda-Convert
+--------------------------------------------------------------------------------
+
+Joda Convert
+Copyright 2010-present Stephen Colebourne
+
+This product includes software developed by
+Joda.org (https://www.joda.org/).
+
+
+Joda-Convert includes code from Google Guava, which is licensed as follows:
+
+Copyright (C) 2011 The Guava Authors
+
+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
+
+https://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.
+
+--------------------------------------------------------------------------------
+Eclipse Jersey (jersey-container-servlet, jersey-container-servlet-core, jersey-client, jersey-hk2, jersey-media-jaxb)
+--------------------------------------------------------------------------------
+
+# Notice for Jersey
+This content is produced and maintained by the Eclipse Jersey project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jersey
+
+## Trademarks
+Eclipse Jersey is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jersey
+
+## Third-party Content
+
+Angular JS, v1.6.6
+* License MIT (http://www.opensource.org/licenses/mit-license.php)
+* Project: http://angularjs.org
+* Coyright: (c) 2010-2017 Google, Inc.
+
+aopalliance Version 1
+* License: all the source code provided by AOP Alliance is Public Domain.
+* Project: http://aopalliance.sourceforge.net
+* Copyright: Material in the public domain is not protected by copyright
+
+Bean Validation API 3.0.2
+* License: Apache License, 2.0
+* Project: http://beanvalidation.org/1.1/
+* Copyright: 2009, Red Hat, Inc. and/or its affiliates, and individual contributors
+* by the @authors tag.
+
+Hibernate Validator CDI, 7.0.5.Final
+* License: Apache License, 2.0
+* Project: https://beanvalidation.org/
+* Repackaged in org.glassfish.jersey.server.validation.internal.hibernate
+
+Bootstrap v3.3.7
+* License: MIT license (https://github.com/twbs/bootstrap/blob/master/LICENSE)
+* Project: http://getbootstrap.com
+* Copyright: 2011-2016 Twitter, Inc
+
+Google Guava Version 18.0
+* License: Apache License, 2.0
+* Copyright (C) 2009 The Guava Authors
+
+jakarta.inject Version: 1
+* License: Apache License, 2.0
+* Copyright (C) 2009 The JSR-330 Expert Group
+
+Javassist Version 3.29.2-GA
+* License: Apache License, 2.0
+* Project: http://www.javassist.org/
+* Copyright (C) 1999- Shigeru Chiba. All Rights Reserved.
+
+Jackson JAX-RS Providers Version 2.15.3
+* License: Apache License, 2.0
+* Project: https://github.com/FasterXML/jackson-jaxrs-providers
+* Copyright: (c) 2009-2023 FasterXML, LLC. All rights reserved unless otherwise indicated.
+
+jQuery v1.12.4
+* License: jquery.org/license
+* Project: jquery.org
+* Copyright: (c) jQuery Foundation
+
+jQuery Barcode plugin 0.3
+* License: MIT & GPL (http://www.opensource.org/licenses/mit-license.php & http://www.gnu.org/licenses/gpl.html)
+* Project: http://www.pasella.it/projects/jQuery/barcode
+* Copyright: (c) 2009 Antonello Pasella antonello.pasella@gmail.com
+
+JSR-166 Extension - JEP 266
+* License: CC0
+* No copyright
+* Written by Doug Lea with assistance from members of JCP JSR-166 Expert Group and released to the public domain, as explained at http://creativecommons.org/publicdomain/zero/1.0/
+
+KineticJS, v4.7.1
+* License: MIT license (http://www.opensource.org/licenses/mit-license.php)
+* Project: http://www.kineticjs.com, https://github.com/ericdrowell/KineticJS
+* Copyright: Eric Rowell
+
+org.objectweb.asm Version 9.6
+* License: Modified BSD (https://asm.ow2.io/license.html)
+* Copyright (c) 2000-2011 INRIA, France Telecom. All rights reserved.
+
+org.osgi.core version 6.0.0
+* License: Apache License, 2.0
+* Copyright (c) OSGi Alliance (2005, 2008). All Rights Reserved.
+
+org.glassfish.jersey.server.internal.monitoring.core
+* License: Apache License, 2.0
+* Copyright (c) 2015-2018 Oracle and/or its affiliates. All rights reserved.
+* Copyright 2010-2013 Coda Hale and Yammer, Inc.
+
+W3.org documents
+* License: W3C License
+* Copyright: Copyright (c) 1994-2001 World Wide Web Consortium, (Massachusetts Institute of Technology, Institut National de Recherche en Informatique et en Automatique, Keio University). All Rights Reserved. http://www.w3.org/Consortium/Legal/
+
+--------------------------------------------------------------------------------
+Eclipse Jersey Core Server
+--------------------------------------------------------------------------------
+
+# Notice for Jersey Core Server module
+This content is produced and maintained by the Eclipse Jersey project.
+
+* https://projects.eclipse.org/projects/ee4j.jersey
+
+## Trademarks
+Eclipse Jersey is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jersey
+
+## Third-party Content
+
+org.glassfish.jersey.server.internal.monitoring.core
+* License: Apache License, 2.0
+* Copyright (c) 2015-2018 Oracle and/or its affiliates. All rights reserved.
+* Copyright 2010-2013 Coda Hale and Yammer, Inc.
+
+org.objectweb.asm Version 9.6
+* License: Modified BSD (https://asm.ow2.io/license.html)
+* Copyright: (c) 2000-2011 INRIA, France Telecom. All rights reserved.
+
+W3.org documents
+* License: W3C License
+* Copyright: Copyright (c) 1994-2001 World Wide Web Consortium, (Massachusetts Institute of Technology, Institut National de Recherche en Informatique et en Automatique, Keio University). All Rights Reserved. http://www.w3.org/Consortium/Legal/
+
+# Notice for Jersey
+This content is produced and maintained by the Eclipse Jersey project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jersey
+
+## Trademarks
+Eclipse Jersey is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jersey
+
+## Third-party Content
+
+Angular JS, v1.6.6
+* License MIT (http://www.opensource.org/licenses/mit-license.php)
+* Project: http://angularjs.org
+* Coyright: (c) 2010-2017 Google, Inc.
+
+aopalliance Version 1
+* License: all the source code provided by AOP Alliance is Public Domain.
+* Project: http://aopalliance.sourceforge.net
+* Copyright: Material in the public domain is not protected by copyright
+
+Bean Validation API 3.0.2
+* License: Apache License, 2.0
+* Project: http://beanvalidation.org/1.1/
+* Copyright: 2009, Red Hat, Inc. and/or its affiliates, and individual contributors
+* by the @authors tag.
+
+Hibernate Validator CDI, 7.0.5.Final
+* License: Apache License, 2.0
+* Project: https://beanvalidation.org/
+* Repackaged in org.glassfish.jersey.server.validation.internal.hibernate
+
+Bootstrap v3.3.7
+* License: MIT license (https://github.com/twbs/bootstrap/blob/master/LICENSE)
+* Project: http://getbootstrap.com
+* Copyright: 2011-2016 Twitter, Inc
+
+Google Guava Version 18.0
+* License: Apache License, 2.0
+* Copyright (C) 2009 The Guava Authors
+
+jakarta.inject Version: 1
+* License: Apache License, 2.0
+* Copyright (C) 2009 The JSR-330 Expert Group
+
+Javassist Version 3.29.2-GA
+* License: Apache License, 2.0
+* Project: http://www.javassist.org/
+* Copyright (C) 1999- Shigeru Chiba. All Rights Reserved.
+
+Jackson JAX-RS Providers Version 2.15.3
+* License: Apache License, 2.0
+* Project: https://github.com/FasterXML/jackson-jaxrs-providers
+* Copyright: (c) 2009-2023 FasterXML, LLC. All rights reserved unless otherwise indicated.
+
+jQuery v1.12.4
+* License: jquery.org/license
+* Project: jquery.org
+* Copyright: (c) jQuery Foundation
+
+jQuery Barcode plugin 0.3
+* License: MIT & GPL (http://www.opensource.org/licenses/mit-license.php & http://www.gnu.org/licenses/gpl.html)
+* Project: http://www.pasella.it/projects/jQuery/barcode
+* Copyright: (c) 2009 Antonello Pasella antonello.pasella@gmail.com
+
+JSR-166 Extension - JEP 266
+* License: CC0
+* No copyright
+* Written by Doug Lea with assistance from members of JCP JSR-166 Expert Group and released to the public domain, as explained at http://creativecommons.org/publicdomain/zero/1.0/
+
+KineticJS, v4.7.1
+* License: MIT license (http://www.opensource.org/licenses/mit-license.php)
+* Project: http://www.kineticjs.com, https://github.com/ericdrowell/KineticJS
+* Copyright: Eric Rowell
+
+org.objectweb.asm Version 9.6
+* License: Modified BSD (https://asm.ow2.io/license.html)
+* Copyright (c) 2000-2011 INRIA, France Telecom. All rights reserved.
+
+org.osgi.core version 6.0.0
+* License: Apache License, 2.0
+* Copyright (c) OSGi Alliance (2005, 2008). All Rights Reserved.
+
+org.glassfish.jersey.server.internal.monitoring.core
+* License: Apache License, 2.0
+* Copyright (c) 2015-2018 Oracle and/or its affiliates. All rights reserved.
+* Copyright 2010-2013 Coda Hale and Yammer, Inc.
+
+W3.org documents
+* License: W3C License
+* Copyright: Copyright (c) 1994-2001 World Wide Web Consortium, (Massachusetts Institute of Technology, Institut National de Recherche en Informatique et en Automatique, Keio University). All Rights Reserved. http://www.w3.org/Consortium/Legal/
+
+--------------------------------------------------------------------------------
+Eclipse Jersey Core Common
+--------------------------------------------------------------------------------
+
+# Notice for Jersey Core Common module
+This content is produced and maintained by the Eclipse Jersey project.
+
+
+* https://projects.eclipse.org/projects/ee4j.jersey
+
+## Trademarks
+Eclipse Jersey is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jersey
+
+## Third-party Content
+
+Google Guava Version 18.0
+* License: Apache License, 2.0
+* Copyright: (C) 2009 The Guava Authors
+
+JSR-166 Extension - JEP 266
+* License: Creative Commons 1.0 (CC0)
+* No copyright
+* Written by Doug Lea with assistance from members of JCP JSR-166
+* Expert Group and released to the public domain, as explained at
+* http://creativecommons.org/publicdomain/zero/1.0/
+
+# Notice for Jersey
+This content is produced and maintained by the Eclipse Jersey project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jersey
+
+## Trademarks
+Eclipse Jersey is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jersey
+
+## Third-party Content
+
+Angular JS, v1.6.6
+* License MIT (http://www.opensource.org/licenses/mit-license.php)
+* Project: http://angularjs.org
+* Coyright: (c) 2010-2017 Google, Inc.
+
+aopalliance Version 1
+* License: all the source code provided by AOP Alliance is Public Domain.
+* Project: http://aopalliance.sourceforge.net
+* Copyright: Material in the public domain is not protected by copyright
+
+Bean Validation API 3.0.2
+* License: Apache License, 2.0
+* Project: http://beanvalidation.org/1.1/
+* Copyright: 2009, Red Hat, Inc. and/or its affiliates, and individual contributors
+* by the @authors tag.
+
+Hibernate Validator CDI, 7.0.5.Final
+* License: Apache License, 2.0
+* Project: https://beanvalidation.org/
+* Repackaged in org.glassfish.jersey.server.validation.internal.hibernate
+
+Bootstrap v3.3.7
+* License: MIT license (https://github.com/twbs/bootstrap/blob/master/LICENSE)
+* Project: http://getbootstrap.com
+* Copyright: 2011-2016 Twitter, Inc
+
+Google Guava Version 18.0
+* License: Apache License, 2.0
+* Copyright (C) 2009 The Guava Authors
+
+jakarta.inject Version: 1
+* License: Apache License, 2.0
+* Copyright (C) 2009 The JSR-330 Expert Group
+
+Javassist Version 3.29.2-GA
+* License: Apache License, 2.0
+* Project: http://www.javassist.org/
+* Copyright (C) 1999- Shigeru Chiba. All Rights Reserved.
+
+Jackson JAX-RS Providers Version 2.15.3
+* License: Apache License, 2.0
+* Project: https://github.com/FasterXML/jackson-jaxrs-providers
+* Copyright: (c) 2009-2023 FasterXML, LLC. All rights reserved unless otherwise indicated.
+
+jQuery v1.12.4
+* License: jquery.org/license
+* Project: jquery.org
+* Copyright: (c) jQuery Foundation
+
+jQuery Barcode plugin 0.3
+* License: MIT & GPL (http://www.opensource.org/licenses/mit-license.php & http://www.gnu.org/licenses/gpl.html)
+* Project: http://www.pasella.it/projects/jQuery/barcode
+* Copyright: (c) 2009 Antonello Pasella antonello.pasella@gmail.com
+
+JSR-166 Extension - JEP 266
+* License: CC0
+* No copyright
+* Written by Doug Lea with assistance from members of JCP JSR-166 Expert Group and released to the public domain, as explained at http://creativecommons.org/publicdomain/zero/1.0/
+
+KineticJS, v4.7.1
+* License: MIT license (http://www.opensource.org/licenses/mit-license.php)
+* Project: http://www.kineticjs.com, https://github.com/ericdrowell/KineticJS
+* Copyright: Eric Rowell
+
+org.objectweb.asm Version 9.6
+* License: Modified BSD (https://asm.ow2.io/license.html)
+* Copyright (c) 2000-2011 INRIA, France Telecom. All rights reserved.
+
+org.osgi.core version 6.0.0
+* License: Apache License, 2.0
+* Copyright (c) OSGi Alliance (2005, 2008). All Rights Reserved.
+
+org.glassfish.jersey.server.internal.monitoring.core
+* License: Apache License, 2.0
+* Copyright (c) 2015-2018 Oracle and/or its affiliates. All rights reserved.
+* Copyright 2010-2013 Coda Hale and Yammer, Inc.
+
+W3.org documents
+* License: W3C License
+* Copyright: Copyright (c) 1994-2001 World Wide Web Consortium, (Massachusetts Institute of Technology, Institut National de Recherche en Informatique et en Automatique, Keio University). All Rights Reserved. http://www.w3.org/Consortium/Legal/
+
+--------------------------------------------------------------------------------
+Eclipse Jersey Bean Validation
+--------------------------------------------------------------------------------
+
+# Notice for Jersey Bean Validation module
+This content is produced and maintained by the Eclipse Jersey project.
+
+* https://projects.eclipse.org/projects/ee4j.jersey
+
+## Trademarks
+Eclipse Jersey is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jersey
+
+## Third-party Content
+
+Hibernate Validator CDI, 7.0.5.Final
+* License: Apache License, 2.0
+* Project: https://beanvalidation.org/
+* Repackaged in org.glassfish.jersey.server.validation.internal.hibernate
+# Notice for Jersey
+This content is produced and maintained by the Eclipse Jersey project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jersey
+
+## Trademarks
+Eclipse Jersey is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jersey
+
+## Third-party Content
+
+Angular JS, v1.6.6
+* License MIT (http://www.opensource.org/licenses/mit-license.php)
+* Project: http://angularjs.org
+* Coyright: (c) 2010-2017 Google, Inc.
+
+aopalliance Version 1
+* License: all the source code provided by AOP Alliance is Public Domain.
+* Project: http://aopalliance.sourceforge.net
+* Copyright: Material in the public domain is not protected by copyright
+
+Bean Validation API 3.0.2
+* License: Apache License, 2.0
+* Project: http://beanvalidation.org/1.1/
+* Copyright: 2009, Red Hat, Inc. and/or its affiliates, and individual contributors
+* by the @authors tag.
+
+Hibernate Validator CDI, 7.0.5.Final
+* License: Apache License, 2.0
+* Project: https://beanvalidation.org/
+* Repackaged in org.glassfish.jersey.server.validation.internal.hibernate
+
+Bootstrap v3.3.7
+* License: MIT license (https://github.com/twbs/bootstrap/blob/master/LICENSE)
+* Project: http://getbootstrap.com
+* Copyright: 2011-2016 Twitter, Inc
+
+Google Guava Version 18.0
+* License: Apache License, 2.0
+* Copyright (C) 2009 The Guava Authors
+
+jakarta.inject Version: 1
+* License: Apache License, 2.0
+* Copyright (C) 2009 The JSR-330 Expert Group
+
+Javassist Version 3.29.2-GA
+* License: Apache License, 2.0
+* Project: http://www.javassist.org/
+* Copyright (C) 1999- Shigeru Chiba. All Rights Reserved.
+
+Jackson JAX-RS Providers Version 2.15.3
+* License: Apache License, 2.0
+* Project: https://github.com/FasterXML/jackson-jaxrs-providers
+* Copyright: (c) 2009-2023 FasterXML, LLC. All rights reserved unless otherwise indicated.
+
+jQuery v1.12.4
+* License: jquery.org/license
+* Project: jquery.org
+* Copyright: (c) jQuery Foundation
+
+jQuery Barcode plugin 0.3
+* License: MIT & GPL (http://www.opensource.org/licenses/mit-license.php & http://www.gnu.org/licenses/gpl.html)
+* Project: http://www.pasella.it/projects/jQuery/barcode
+* Copyright: (c) 2009 Antonello Pasella antonello.pasella@gmail.com
+
+JSR-166 Extension - JEP 266
+* License: CC0
+* No copyright
+* Written by Doug Lea with assistance from members of JCP JSR-166 Expert Group and released to the public domain, as explained at http://creativecommons.org/publicdomain/zero/1.0/
+
+KineticJS, v4.7.1
+* License: MIT license (http://www.opensource.org/licenses/mit-license.php)
+* Project: http://www.kineticjs.com, https://github.com/ericdrowell/KineticJS
+* Copyright: Eric Rowell
+
+org.objectweb.asm Version 9.6
+* License: Modified BSD (https://asm.ow2.io/license.html)
+* Copyright (c) 2000-2011 INRIA, France Telecom. All rights reserved.
+
+org.osgi.core version 6.0.0
+* License: Apache License, 2.0
+* Copyright (c) OSGi Alliance (2005, 2008). All Rights Reserved.
+
+org.glassfish.jersey.server.internal.monitoring.core
+* License: Apache License, 2.0
+* Copyright (c) 2015-2018 Oracle and/or its affiliates. All rights reserved.
+* Copyright 2010-2013 Coda Hale and Yammer, Inc.
+
+W3.org documents
+* License: W3C License
+* Copyright: Copyright (c) 1994-2001 World Wide Web Consortium, (Massachusetts Institute of Technology, Institut National de Recherche en Informatique et en Automatique, Keio University). All Rights Reserved. http://www.w3.org/Consortium/Legal/
+
+--------------------------------------------------------------------------------
+Eclipse GlassFish HK2 (aopalliance-repackaged, hk2-api, hk2-locator, hk2-utils)
+--------------------------------------------------------------------------------
+
+# Notices for Eclipse GlassFish
+
+This content is produced and maintained by the Eclipse GlassFish project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.glassfish
+
+## Trademarks
+
+Eclipse GlassFish, and GlassFish are trademarks of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/glassfish-ha-api
+* https://github.com/eclipse-ee4j/glassfish-logging-annotation-processor
+* https://github.com/eclipse-ee4j/glassfish-shoal
+* https://github.com/eclipse-ee4j/glassfish-cdi-porting-tck
+* https://github.com/eclipse-ee4j/glassfish-jsftemplating
+* https://github.com/eclipse-ee4j/glassfish-hk2-extra
+* https://github.com/eclipse-ee4j/glassfish-hk2
+* https://github.com/eclipse-ee4j/glassfish-fighterfish
+
+## Third-party Content
+
+This project leverages the following third party content.
+
+None
+
+## Cryptography
+
+Content may contain encryption software. The country in which you are currently
+may have restrictions on the import, possession, and use, and/or re-export to
+another country, of encryption software. BEFORE using any encryption software,
+please check the country's laws, regulations and policies concerning the import,
+possession, or use, and re-export of encryption software, to see if this is
+permitted.
+
+--------------------------------------------------------------------------------
+Eclipse Jetty Servlet API (jakarta-servlet-api 5.0.2)
+--------------------------------------------------------------------------------
+
+# Notices for Eclipse Project for Servlet
+
+This content is produced and maintained by the Eclipse Project for Servlet
+project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.servlet
+
+
+## Trademarks
+
+Eclipse Project for Servlet is a trademark of the Eclipse Foundation.
+
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+ * https://github.com/eclipse-ee4j/servlet-api
+ * https://github.com/eclipse/jetty.toolchain
+
+
+## Third-party Content
+
+## Jakarta
+
+The following artifacts are EPL 2.0 + GPLv2 with classpath exception.
+https://projects.eclipse.org/projects/ee4j.servlet
+
+ * jakarta.servlet:jakarta.servlet-api
+
+
+## GlassFish
+
+The following artifacts are CDDL + GPLv2 with classpath exception.
+https://glassfish.dev.java.net/nonav/public/CDDL+GPL.html
+
+ * org.eclipse.jetty.toolchain:jetty-schemas
+
+--------------------------------------------------------------------------------
+Jakarta XML Binding API (jakarta.xml.bind-api 3.0.x)
+--------------------------------------------------------------------------------
+
+[//]: # " Copyright (c) 2018, 2019 Oracle and/or its affiliates. All rights reserved. "
+[//]: # " "
+[//]: # " This program and the accompanying materials are made available under the "
+[//]: # " terms of the Eclipse Distribution License v. 1.0, which is available at "
+[//]: # " http://www.eclipse.org/org/documents/edl-v10.php. "
+[//]: # " "
+[//]: # " SPDX-License-Identifier: BSD-3-Clause "
+
+# Notices for Jakarta XML Binding
+
+This content is produced and maintained by the Jakarta XML Binding
+project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jaxb
+
+## Trademarks
+
+Jakarta XML Binding is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Distribution License v. 1.0 which is available at
+http://www.eclipse.org/org/documents/edl-v10.php.
+
+SPDX-License-Identifier: BSD-3-Clause
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jaxb-api
+* https://github.com/eclipse-ee4j/jaxb-tck
+
+## Third-party Content
+
+This project leverages the following third party content.
+
+Apache River (3.0.0)
+
+* License: Apache-2.0 AND BSD-3-Clause
+
+ASM 7 (n/a)
+
+* License: BSD-3-Clause
+* Project: https://asm.ow2.io/
+* Source:
+ https://repository.ow2.org/nexus/#nexus-search;gav~org.ow2.asm~asm-commons~~~~kw,versionexpand
+
+JTHarness (5.0)
+
+* License: (GPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0)
+* Project: https://wiki.openjdk.java.net/display/CodeTools/JT+Harness
+* Source: http://hg.openjdk.java.net/code-tools/jtharness/
+
+normalize.css (3.0.2)
+
+* License: MIT
+
+SigTest (n/a)
+
+* License: GPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Cryptography
+
+Content may contain encryption software. The country in which you are currently
+may have restrictions on the import, possession, and use, and/or re-export to
+another country, of encryption software. BEFORE using any encryption software,
+please check the country's laws, regulations and policies concerning the import,
+possession, or use, and re-export of encryption software, to see if this is
+permitted.
+
+--------------------------------------------------------------------------------
+Jakarta RESTful Web Services API (jakarta.ws.rs-api 3.0.x / 3.1.0)
+--------------------------------------------------------------------------------
+
+# Notices for Jakarta RESTful Web Services
+
+This content is produced and maintained by the **Jakarta RESTful Web Services**
+project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jaxrs
+
+## Trademarks
+
+**Jakarta RESTful Web Services** is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jaxrs-api
+
+## Third-party Content
+
+This project leverages the following third party content.
+
+javaee-api (7.0)
+
+* License: Apache-2.0 AND W3C
+
+JUnit (4.11)
+
+* License: Common Public License 1.0
+
+Mockito (2.16.0)
+
+* Project: http://site.mockito.org
+* Source: https://github.com/mockito/mockito/releases/tag/v2.16.0
+
+## Cryptography
+
+Content may contain encryption software. The country in which you are currently
+may have restrictions on the import, possession, and use, and/or re-export to
+another country, of encryption software. BEFORE using any encryption software,
+please check the country's laws, regulations and policies concerning the import,
+possession, or use, and re-export of encryption software, to see if this is
+permitted.
+
+--------------------------------------------------------------------------------
+Jakarta Expression Language API (jakarta.el-api 4.0.0, glassfish jakarta.el 4.0.2)
+--------------------------------------------------------------------------------
+
+# Notices for Jakarta Expression Language
+
+This content is produced and maintained by the Jakarta Expression Language project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.el
+
+## Trademarks
+
+Jakarta Expression Language is a trademark of the Eclipse
+Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/el-ri
+
+## Third-party Content
+
+## Cryptography
+
+Content may contain encryption software. The country in which you are currently
+may have restrictions on the import, possession, and use, and/or re-export to
+another country, of encryption software. BEFORE using any encryption software,
+please check the country's laws, regulations and policies concerning the import,
+possession, or use, and re-export of encryption software, to see if this is
+permitted.
+
+--------------------------------------------------------------------------------
+Jakarta Annotations API (jakarta.annotation-api 2.1.1 and 3.0.0)
+--------------------------------------------------------------------------------
+
+# Notices for Jakarta Annotations
+
+This content is produced and maintained by the Jakarta Annotations project.
+
+ * Project home: https://projects.eclipse.org/projects/ee4j.ca
+
+## Trademarks
+
+Jakarta Annotations is a trademark of the Eclipse Foundation.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+http://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied: GNU
+General Public License, version 2 with the GNU Classpath Exception which is
+available at https://www.gnu.org/software/classpath/license.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+ * https://github.com/eclipse-ee4j/common-annotations-api
+
+## Third-party Content
+
+## Cryptography
+
+Content may contain encryption software. The country in which you are currently
+may have restrictions on the import, possession, and use, and/or re-export to
+another country, of encryption software. BEFORE using any encryption software,
+please check the country's laws, regulations and policies concerning the import,
+possession, or use, and re-export of encryption software, to see if this is
+permitted.
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Notices for Jakarta Annotations
+
+This content is produced and maintained by the Jakarta Annotations project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.ca
+
+## Trademarks
+
+Jakarta Annotations™ is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Public License v. 2.0 which is available at
+https://www.eclipse.org/legal/epl-2.0. This Source Code may also be made
+available under the following Secondary Licenses when the conditions for such
+availability set forth in the Eclipse Public License v. 2.0 are satisfied:
+GPL-2.0 with Classpath-exception-2.0 which is available at
+https://openjdk.java.net/legal/gplv2+ce.html.
+
+SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only with Classpath-exception-2.0
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+* https://github.com/jakartaee/common-annotations-api
+
+## Cryptography
+
+Content may contain encryption software. The country in which you are currently
+may have restrictions on the import, possession, and use, and/or re-export to
+another country, of encryption software. BEFORE using any encryption software,
+please check the country's laws, regulations and policies concerning the import,
+possession, or use, and re-export of encryption software, to see if this is
+permitted.
+
+--------------------------------------------------------------------------------
+Jakarta Inject API (jakarta.inject-api 2.0.1)
+--------------------------------------------------------------------------------
+
+# Notices for Eclipse Jakarta Dependency Injection
+
+This content is produced and maintained by the Eclipse Jakarta Dependency Injection project.
+
+* Project home: https://projects.eclipse.org/projects/cdi.batch
+
+## Trademarks
+
+Jakarta Dependency Injection is a trademark of the Eclipse Foundation.
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Apache License, Version 2.0 which is available at
+https://www.apache.org/licenses/LICENSE-2.0.
+
+SPDX-License-Identifier: Apache-2.0
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+https://github.com/eclipse-ee4j/injection-api
+https://github.com/eclipse-ee4j/injection-spec
+https://github.com/eclipse-ee4j/injection-tck
+
+## Third-party Content
+
+This project leverages the following third party content.
+
+None
+
+## Cryptography
+
+None
+
+--------------------------------------------------------------------------------
+Jakarta Activation (jakarta.activation 2.0.0, 2.0.1, jakarta.activation-api 1.2.1, 2.1.0)
+--------------------------------------------------------------------------------
+
+# Notices for Eclipse Project for JAF
+
+This content is produced and maintained by the Eclipse Project for JAF project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jaf
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Distribution License v. 1.0,
+which is available at http://www.eclipse.org/org/documents/edl-v10.php.
+
+SPDX-License-Identifier: BSD-3-Clause
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jaf
+
+## Third-party Content
+
+This project leverages the following third party content.
+
+JUnit (4.12)
+
+* License: Eclipse Public License
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Notices for Jakarta Activation
+
+This content is produced and maintained by Jakarta Activation project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jaf
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Distribution License v. 1.0,
+which is available at http://www.eclipse.org/org/documents/edl-v10.php.
+
+SPDX-License-Identifier: BSD-3-Clause
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jaf
+
+## Third-party Content
+
+This project leverages the following third party content.
+
+JUnit (4.12)
+
+* License: Eclipse Public License
+
+- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+# Notices for Jakarta Activation
+
+This content is produced and maintained by Jakarta Activation project.
+
+* Project home: https://projects.eclipse.org/projects/ee4j.jaf
+
+## Copyright
+
+All content is the property of the respective authors or their employers. For
+more information regarding authorship of content, please consult the listed
+source code repository logs.
+
+## Declared Project Licenses
+
+This program and the accompanying materials are made available under the terms
+of the Eclipse Distribution License v. 1.0,
+which is available at http://www.eclipse.org/org/documents/edl-v10.php.
+
+SPDX-License-Identifier: BSD-3-Clause
+
+## Source Code
+
+The project maintains the following source code repositories:
+
+* https://github.com/eclipse-ee4j/jaf
+
+--------------------------------------------------------------------------------
+Bouncy Castle
+--------------------------------------------------------------------------------
+
+Bouncy Castle Cryptography (org.bouncycastle.*) is distributed under the
+Bouncy Castle License, whose text reproduces the MIT License with Bouncy
+Castle's copyright line:
+
+ Copyright (c) 2000-2024 The Legion of the Bouncy Castle Inc.
+ (https://www.bouncycastle.org)
+
+The full MIT-equivalent text is reproduced in
+licenses/LICENSE-MIT.txt; the Bouncy Castle copyright line above must be
+preserved alongside it.
+
+--------------------------------------------------------------------------------
+aiohttp (vendored MIT-licensed C parser code)
+--------------------------------------------------------------------------------
+
+The aiohttp wheel is primarily Apache-2.0 but bundles a small amount of
+MIT-licensed C code used for HTTP parsing. The MIT attribution required
+for redistribution of that code is:
+
+ Copyright (c) the aiohttp contributors and upstream authors of the
+ vendored parser source. Distributed under the MIT License.
+
+Consumers of the aiohttp wheel must preserve that MIT copyright notice in
+any downstream redistribution of the parser source or object code.
+
+--------------------------------------------------------------------------------
+Matplotlib
+--------------------------------------------------------------------------------
+
+Matplotlib is distributed under the "Matplotlib License", a modification of
+the PSF License. The attribution required for redistribution is:
+
+ Copyright (c) 2012- Matplotlib Development Team; All Rights Reserved.
+
+The full license text is covered by licenses/LICENSE-PSF-2.0.txt.
diff --git a/README.md b/README.md
index 7f70fdc5b22..54dfddb28e7 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,74 @@
-# Texera
-A system for Web-based text analytics
+Apache Texera - Human-AI Collaborative Data Science Using Visual Workflows
-travis-ci
-[](https://travis-ci.org/Texera/texera)
+
+
+
+ Apache Texera (Incubating) is an open-source platform for human-AI collaborative data science using visual workflows.
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+Apache Texera (Incubating) is an open-source platform for human-AI collaborative data science using visual workflows. It enables human analysts to construct, execute, and refine data analysis tasks through an intuitive GUI, assisted by AI agents that understand natural-language instructions. Texera is well suited for a wide range of applications, including “AI for Science,” by making advanced AI and data science capabilities accessible to a broader community. It can run on a laptop for local use or be deployed in the cloud to support scalable processing of large datasets.
+
+The platform has the following key features:
+
+* Natural-language data science through AI agents
+* Intuitive GUI-based workflows for data science
+* Real-time collaboration for workflow editing and execution
+* Runtime debugging and interactive workflow execution
+* Language-agnostic workflow runtime, native support for Python and Java
+* Parallel backend engine for scalable big-data processing
+* Separation of compute and storage for flexible cloud deployment
+
+
+
+
+
+# Citation
+Please cite Texera as
+```
+
+@article{DBLP:journals/pvldb/WangHNKALLDL24,
+ author = {Zuozhi Wang and
+ Yicong Huang and
+ Shengquan Ni and
+ Avinash Kumar and
+ Sadeem Alsudais and
+ Xiaozhen Liu and
+ Xinyuan Lin and
+ Yunyan Ding and
+ Chen Li},
+ title = {Texera: {A} System for Collaborative and Interactive Data Analytics
+ Using Workflows},
+ journal = {Proc. {VLDB} Endow.},
+ volume = {17},
+ number = {11},
+ pages = {3580--3588},
+ year = {2024},
+ url = {https://www.vldb.org/pvldb/vol17/p3580-wang.pdf},
+ timestamp = {Thu, 19 Sep 2024 13:09:37 +0200},
+ biburl = {https://dblp.org/rec/journals/pvldb/WangHNKALLDL24.bib},
+ bibsource = {dblp computer science bibliography, https://dblp.org}
+}
+```
diff --git a/SECURITY.md b/SECURITY.md
new file mode 100644
index 00000000000..0758ffce488
--- /dev/null
+++ b/SECURITY.md
@@ -0,0 +1,280 @@
+# Security Policy
+
+This document outlines Apache Texera (Incubating)'s security model, deployment considerations, and procedures for
+reporting security vulnerabilities.
+
+## Table of Contents
+
+- [Security Model Overview](#security-model-overview)
+- [Resources in Texera](#resources-in-texera)
+- [User Categories and Responsibilities](#user-categories-and-responsibilities)
+- [UI User Roles and Privileges](#ui-user-roles-and-privileges)
+- [Deployments and Computing Units](#deployments-and-computing-units)
+- [What is NOT a Security Issue](#what-is-not-a-security-issue)
+- [Reporting Security Vulnerabilities](#reporting-security-vulnerabilities)
+
+## Security Model Overview
+
+Texera's security architecture is built around:
+
+1. **Authentication**: JWT-based token authentication with configurable expiration
+2. **Authorization**: Role-based access control (RBAC) with four user roles
+3. **Resource Access Control**: Fine-grained privileges for datasets, workflows, and computing units
+4. **Deployment Isolation**: Separate security considerations for different deployment modes
+
+## Resources in Texera
+
+In Texera, a **resource** is any object within the system that can be created, accessed, modified, or shared by users
+via the web application. Understanding resource types and how access to them is managed is critical to following
+Texera’s security model.
+
+### Resource Types
+
+Texera supports the following resource types:
+
+- **Datasets**: Input data imported or uploaded for workflow processing
+- **Workflows**: Data analytics pipelines defined by users
+- **Computing Units**: Execution environments for running workflows (e.g., Kubernates PODs)
+- **Results**: Output from workflow executions, including but not limited to data, logs, metrics, and visualizations
+
+### Resource Ownership and Access Control
+
+Every resource is owned by a user. The owner controls the resource's visibility and can share it with other users by
+granting access permissions:
+
+- **READ**: View the resource and its contents
+- **WRITE**: Modify, execute, delete, and share the resource
+- **NONE**: No access to the resource
+
+Resources can be shared with specific users or made public. Public resources are visible to all users. Resource owners
+can modify access permissions at any time.
+
+### Resource Visibility
+
+- Users can only see resources for which they have at least READ access.
+- Access changes (e.g., revoking WRITE or READ) take effect immediately for affected users.
+
+## User Categories and Responsibilities
+
+Texera's security model distinguishes between two categories of users with distinct responsibilities:
+
+### Deployment Managers
+
+They have the highest level of access and control. They install and configure Texera, and make decisions about
+technologies, deployment modes, and permissions. They can potentially delete the entire installation and have access to
+all credentials, including database passwords, JWT secrets, and API keys. Deployment managers have full access to:
+
+- The underlying infrastructure (servers, Kubernetes clusters, cloud resources)
+- Database administration (e.g., PostgreSQL)
+- All configuration files, environment variables, and secrets
+- Network and security settings
+- Container orchestration and system logs
+
+Deployment managers can also decide to keep audits, backups, and copies of information outside of Texera, which are not
+covered by Texera's security model. They operate outside the Texera UI role system and may or may not have a UI user
+account.
+
+### UI Users
+
+**Who They Are**: Individuals who interact with Texera through the web interface.
+
+**Access Level**: UI users interact with Texera through the web interface and do not have direct access to:
+
+- The underlying infrastructure (servers, Kubernetes cluster)
+- Database administration
+- System configuration files
+- Network and firewall settings
+- Container orchestration
+
+**Important**: REGULAR and ADMIN users can execute arbitrary code through UDFs, which may access resources in the execution environment. Deployment managers are responsible for mitigating this risk. See [What is NOT a Security Issue](#what-is-not-a-security-issue) for details.
+
+**Roles**: UI users are assigned one of four roles (INACTIVE, RESTRICTED, REGULAR, ADMIN) that control their permissions
+within the Texera application.
+
+**Security Scope**: UI users are responsible for:
+
+- Protecting their login credentials
+- Managing access to their resources, e.g., datasets and workflows
+- Following organizational data security policies
+
+## UI User Roles and Privileges
+
+Texera implements four UI user roles with increasing levels of privilege. These roles control what users can do **within
+the Texera web application** and do not grant infrastructure-level access.
+
+### 1. INACTIVE
+
+Users with this role cannot log in to the system or access any resources. This is the default role for new registrations
+awaiting approval in controlled environments.
+
+### 2. RESTRICTED
+
+Users with this role cannot log in to the system or access any resources. Unlike INACTIVE users, RESTRICTED accounts
+typically represent users who previously used Texera but are now inactive and no longer use it. Any resources they
+created in the past remain in the system but are inaccessible to them. This role is used to preserve historical data
+while preventing further access.
+
+### 3. REGULAR
+
+Users with this role can create and manage their own resources (datasets, workflows, computing units). They have full
+READ and WRITE access to resources they own, and their access to other users' resources is determined by granted
+permissions (see Resources section above).
+
+They cannot:
+
+- Access other users' private resources without granted permissions
+- Manage user accounts or change user roles
+- Access system configuration, logs, or global settings
+
+This is the standard role for data scientists, analysts, and researchers.
+**Note**: REGULAR users can execute arbitrary code within workflows, so this role should only be granted to trusted
+individuals.
+
+### 4. ADMIN
+
+Users with this role are application administrators who manage users and resources through the web interface.
+
+They have all REGULAR privileges, plus:
+
+- Manage all UI user accounts (create, modify, and delete users)
+- Change user roles
+- View user login information.
+- Configure application settings available in the web interface
+
+They cannot:
+
+- Access the underlying servers or Kubernetes cluster
+- Modify JWT secrets or database passwords
+- Configure HTTPS/TLS or network settings
+- Access system-level logs or SSH into servers
+
+**Note**: ADMIN is an application-level role, not an infrastructure administrator. For infrastructure management,
+deployment manager access is required.
+
+## Deployments and Computing Units
+Texera can be deployed in several configurations, such as local development, single-node setups, or distributed Kubernetes
+clusters. For details on supported deployment options and their operational differences, see the deployment guides in
+our [wiki](https://github.com/apache/texera/wiki/How-to-run-Texera-on-local-Kubernetes).
+
+### Computing Unit Types
+
+Texera executes workflows on **computing units**. UI users (REGULAR and ADMIN) can execute arbitrary code (e.g., through
+UDFs written in Python, R, Java, Scala) within computing units as part of their workflows. See [What is NOT a Security Issue](#what-is-not-a-security-issue) for the security implications of UDF execution.
+
+Deployment managers configure which types of computing units are available:
+
+#### Local Computing Units
+
+Local computing units run as processes on the same machine as the Texera services (single-node deployment).
+
+**Security characteristics**:
+
+- Suitable for development, testing, and small team use
+- All computing units share the same host machine
+- No infrastructure-level isolation between users' workflows
+- Deployment managers control all computing resources
+
+**Security considerations**:
+
+- Users' workflow code executes on the host machine with limited isolation
+- UDF code executes with access to resources in the host environment — see [What is NOT a Security Issue](#what-is-not-a-security-issue)
+- Deployment managers must trust all REGULAR and ADMIN users
+- Resource exhaustion by one user can affect all users
+
+#### Kubernetes Computing Units
+
+Kubernetes computing units run as separate PODs in a Kubernetes cluster. Each computing unit is dynamically created when
+a user needs it.
+
+**Security characteristics**:
+
+- Suitable for production environments and multi-tenant deployments
+- Each computing unit runs in an isolated Kubernetes pod
+- UI users configure resource limits (CPU, memory, GPU) per pod
+- Pods can be scheduled across multiple nodes for better resource distribution
+
+**Security considerations**:
+
+- Better isolation between users compared to local computing units
+- Kubernetes provides namespace and pod-level isolation
+- Resource limits prevent individual users from consuming excessive resources
+- UDF code within a pod can still access resources available inside that pod's environment (e.g., environment variables, mounted secrets)
+- Container security and image scanning should be implemented
+- Deployment managers must secure the Kubernetes cluster infrastructure
+
+### What is NOT Guaranteed
+
+Texera's security model does NOT guarantee:
+
+- Protection against malicious code in user workflows (users can execute arbitrary code)
+- Isolation of application secrets from UDF code executing within the same process or pod
+- Strong isolation between workflows in local computing units
+- Complete isolation between workflows in Kubernetes computing units within the same namespace
+- Protection against infrastructure-level compromises
+- Protection against deployment manager misconfigurations
+- DDoS protection (requires external infrastructure)
+- Compliance with specific regulatory requirements without additional configuration
+
+## What are NOT Security Issues
+
+The following are **NOT considered security vulnerabilities** in Texera:
+
+### User Code Execution
+
+REGULAR and ADMIN users can execute arbitrary code (Python, R, Java, Scala) within computing units through UDFs. This is by design — custom code execution is a core feature of the platform.
+
+UDF code may access resources available in the execution environment, including but not limited to:
+
+- Texera's application configurations
+- Environment variables of the host
+
+### Resource Consumption
+
+Users can create workflows that consume significant CPU, memory, or storage. Texera is designed for data-intensive
+workloads. Deployment managers control this through computing unit resource limits, quotas, and monitoring.
+
+### Information Disclosure within Authorized Access
+
+Users with READ or WRITE access to a resource can view all its contents. Access control is at the resource level - once
+access is granted, full visibility is expected. Resource owners should grant access only to trusted users.
+
+### Public Resources
+
+Resources marked as public are visible to all users. Public sharing is a deliberate collaboration feature. Users should
+review resources before making them public and avoid including sensitive data or credentials.
+
+### Issues Requiring Deployment Manager Access
+
+Issues requiring physical access to servers, administrative access to infrastructure, database access, or access to
+configuration files are out of scope. These access levels are considered trusted.
+
+### Third-Party Dependencies
+
+Theoretical vulnerabilities in dependencies that have not been exploited in Texera's usage are not in scope.
+You are they are welcome to raise an issue or a PR.
+
+## Reporting Security Vulnerabilities
+
+The [Apache Software Foundation](https://apache.org/) takes a rigorous stance on eliminating security issues in its software projects. If you
+find a security bug, with that in mind, please **DO NOT** file public issues (e.g., GitHub issues). Before reporting a
+security issue, check the security model declared above. To report a new vulnerability you have discovered, please
+follow the ASF security [vulnerability reporting process](https://apache.org/security/#reporting-a-vulnerability).
+The Texera community follows the ASF
+security [vulnerability handling process](https://apache.org/security/#vulnerability-handling), and will fix it as soon
+as possible.
+
+## Changes to This Policy
+
+This security policy may be updated from time to time. Significant changes will be announced on the project mailing
+lists and website.
+
+---
+
+**Last Updated**: April 2026
+
+**Disclaimer**: This project is currently undergoing incubation at The Apache Software Foundation (ASF). Incubation is
+required of all newly accepted projects until a further review indicates that the infrastructure, communications, and
+decision-making process have stabilized in a manner consistent with other successful ASF projects. While incubation
+status is not necessarily a reflection of the completeness or stability of the code, it does indicate that the project
+has yet to be fully endorsed by the ASF.
+
diff --git a/access-control-service/build.sbt b/access-control-service/build.sbt
new file mode 100644
index 00000000000..4dcd74c7ac1
--- /dev/null
+++ b/access-control-service/build.sbt
@@ -0,0 +1,85 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import scala.collection.Seq
+
+name := "access-control-service"
+
+
+enablePlugins(JavaAppPackaging)
+
+// Ship LICENSE-binary, NOTICE-binary, DISCLAIMER-WIP, and the licenses/
+// directory at the top of the Universal dist zip.
+// See project/AddMetaInfLicenseFiles.scala.
+Universal / mappings := AddMetaInfLicenseFiles.distMappings(
+ (Universal / mappings).value,
+ (ThisBuild / baseDirectory).value
+)
+
+// Enable semanticdb for Scalafix
+ThisBuild / semanticdbEnabled := true
+ThisBuild / semanticdbVersion := scalafixSemanticdb.revision
+
+// Manage dependency conflicts by always using the latest revision
+ThisBuild / conflictManager := ConflictManager.latestRevision
+
+// Restrict parallel execution of tests to avoid conflicts
+Global / concurrentRestrictions += Tags.limit(Tags.Test, 1)
+
+/////////////////////////////////////////////////////////////////////////////
+// Compiler Options
+/////////////////////////////////////////////////////////////////////////////
+
+// Scala compiler options
+Compile / scalacOptions ++= Seq(
+ "-Xelide-below", "WARNING", // Turn on optimizations with "WARNING" as the threshold
+ "-feature", // Check feature warnings
+ "-deprecation", // Check deprecation warnings
+ "-Ywarn-unused:imports" // Check for unused imports
+)
+
+/////////////////////////////////////////////////////////////////////////////
+// Version Variables
+/////////////////////////////////////////////////////////////////////////////
+
+val dropwizardVersion = "4.0.7"
+val mockitoVersion = "5.4.0"
+val assertjVersion = "3.24.2"
+
+/////////////////////////////////////////////////////////////////////////////
+// Test-related Dependencies
+/////////////////////////////////////////////////////////////////////////////
+
+libraryDependencies ++= Seq(
+ "org.scalamock" %% "scalamock" % "5.2.0" % Test, // ScalaMock
+ "org.scalatest" %% "scalatest" % "3.2.17" % Test, // ScalaTest
+ "io.dropwizard" % "dropwizard-testing" % dropwizardVersion % Test, // Dropwizard Testing
+ "org.mockito" % "mockito-core" % mockitoVersion % Test, // Mockito for mocking
+ "org.assertj" % "assertj-core" % assertjVersion % Test, // AssertJ for assertions
+ "com.novocode" % "junit-interface" % "0.11" % Test // SBT interface for JUnit
+)
+
+/////////////////////////////////////////////////////////////////////////////
+// Dependencies
+/////////////////////////////////////////////////////////////////////////////
+
+// Core Dependencies
+libraryDependencies ++= Seq(
+ "io.dropwizard" % "dropwizard-core" % dropwizardVersion,
+ "io.dropwizard" % "dropwizard-auth" % dropwizardVersion, // Dropwizard Authentication module
+ "com.fasterxml.jackson.module" %% "jackson-module-scala" % "2.18.6"
+)
\ No newline at end of file
diff --git a/access-control-service/project/build.properties b/access-control-service/project/build.properties
new file mode 100644
index 00000000000..d0ba9f1421a
--- /dev/null
+++ b/access-control-service/project/build.properties
@@ -0,0 +1,18 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+sbt.version = 1.12.9
\ No newline at end of file
diff --git a/access-control-service/src/main/resources/access-control-service-web-config.yaml b/access-control-service/src/main/resources/access-control-service-web-config.yaml
new file mode 100644
index 00000000000..8c7895e9858
--- /dev/null
+++ b/access-control-service/src/main/resources/access-control-service-web-config.yaml
@@ -0,0 +1,36 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+server:
+ applicationConnectors:
+ - type: http
+ port: 9096
+ adminConnectors: []
+ requestLog:
+ type: classic
+ appenders: []
+
+logging:
+ level: ${TEXERA_SERVICE_LOG_LEVEL:-INFO}
+ appenders:
+ - type: console
+ threshold: ${TEXERA_SERVICE_LOG_LEVEL:-INFO}
+ - type: file
+ currentLogFilename: logs/access-control-service.log
+ archive: true
+ archivedLogFilenamePattern: logs/access-control-service-%d.log.gz
+ archivedFileCount: 5
\ No newline at end of file
diff --git a/access-control-service/src/main/resources/logback.xml b/access-control-service/src/main/resources/logback.xml
new file mode 100644
index 00000000000..99794f19684
--- /dev/null
+++ b/access-control-service/src/main/resources/logback.xml
@@ -0,0 +1,55 @@
+
+
+
+
+
+
+ [%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n
+
+
+
+
+
+
+ logs/access-control-service.log
+ true
+
+ logs/access-control-service-%d{yyyy-MM-dd}.log.gz
+
+
+ [%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n
+
+
+
+
+ 8192
+ true
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/access-control-service/src/main/scala/org/apache/texera/service/AccessControlService.scala b/access-control-service/src/main/scala/org/apache/texera/service/AccessControlService.scala
new file mode 100644
index 00000000000..0ab9f0fbfee
--- /dev/null
+++ b/access-control-service/src/main/scala/org/apache/texera/service/AccessControlService.scala
@@ -0,0 +1,99 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.texera.service
+
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import com.typesafe.scalalogging.LazyLogging
+import io.dropwizard.auth.AuthDynamicFeature
+import io.dropwizard.configuration.{EnvironmentVariableSubstitutor, SubstitutingSourceProvider}
+import io.dropwizard.core.Application
+import io.dropwizard.core.setup.{Bootstrap, Environment}
+import org.apache.texera.amber.config.StorageConfig
+import org.apache.texera.auth.{JwtAuthFilter, RequestLoggingFilter, SessionUser}
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.service.resource.{
+ AccessControlResource,
+ HealthCheckResource,
+ LiteLLMModelsResource,
+ LiteLLMProxyResource
+}
+import org.eclipse.jetty.server.session.SessionHandler
+import java.nio.file.Path
+
+class AccessControlService extends Application[AccessControlServiceConfiguration] with LazyLogging {
+ override def initialize(bootstrap: Bootstrap[AccessControlServiceConfiguration]): Unit = {
+ // enable environment variable substitution in YAML config
+ bootstrap.setConfigurationSourceProvider(
+ new SubstitutingSourceProvider(
+ bootstrap.getConfigurationSourceProvider,
+ new EnvironmentVariableSubstitutor(false)
+ )
+ )
+ // Register Scala module to Dropwizard default object mapper
+ bootstrap.getObjectMapper.registerModule(DefaultScalaModule)
+
+ SqlServer.initConnection(
+ StorageConfig.jdbcUrl,
+ StorageConfig.jdbcUsername,
+ StorageConfig.jdbcPassword
+ )
+ }
+
+ override def run(
+ configuration: AccessControlServiceConfiguration,
+ environment: Environment
+ ): Unit = {
+ // Serve backend at /api
+ environment.jersey.setUrlPattern("/api/*")
+
+ environment.jersey.register(classOf[SessionHandler])
+ environment.servlets.setSessionHandler(new SessionHandler)
+
+ environment.jersey.register(classOf[HealthCheckResource])
+ environment.jersey.register(classOf[AccessControlResource])
+ environment.jersey.register(classOf[LiteLLMProxyResource])
+ environment.jersey.register(classOf[LiteLLMModelsResource])
+
+ // Register JWT authentication filter
+ environment.jersey.register(new AuthDynamicFeature(classOf[JwtAuthFilter]))
+
+ // Enable @Auth annotation for injecting SessionUser
+ environment.jersey.register(
+ new io.dropwizard.auth.AuthValueFactoryProvider.Binder(classOf[SessionUser])
+ )
+
+ // Route request logs through SLF4J, controlled by TEXERA_SERVICE_LOG_LEVEL
+ RequestLoggingFilter.register(environment.getApplicationContext)
+ }
+}
+object AccessControlService {
+ def main(args: Array[String]): Unit = {
+ val accessControlPath = Path
+ .of(sys.env.getOrElse("TEXERA_HOME", "."))
+ .resolve("access-control-service")
+ .resolve("src")
+ .resolve("main")
+ .resolve("resources")
+ .resolve("access-control-service-web-config.yaml")
+ .toAbsolutePath
+ .toString
+
+ // Start the Dropwizard application
+ new AccessControlService().run("server", accessControlPath)
+ }
+}
diff --git a/access-control-service/src/main/scala/org/apache/texera/service/AccessControlServiceConfiguration.scala b/access-control-service/src/main/scala/org/apache/texera/service/AccessControlServiceConfiguration.scala
new file mode 100644
index 00000000000..a6e4764c16d
--- /dev/null
+++ b/access-control-service/src/main/scala/org/apache/texera/service/AccessControlServiceConfiguration.scala
@@ -0,0 +1,22 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.texera.service
+
+import io.dropwizard.core.Configuration
+
+class AccessControlServiceConfiguration extends Configuration {}
diff --git a/access-control-service/src/main/scala/org/apache/texera/service/resource/AccessControlResource.scala b/access-control-service/src/main/scala/org/apache/texera/service/resource/AccessControlResource.scala
new file mode 100644
index 00000000000..0cd52f49194
--- /dev/null
+++ b/access-control-service/src/main/scala/org/apache/texera/service/resource/AccessControlResource.scala
@@ -0,0 +1,340 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.texera.service.resource
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import com.typesafe.scalalogging.LazyLogging
+import jakarta.ws.rs.client.{Client, ClientBuilder, Entity}
+import jakarta.ws.rs.core._
+import jakarta.ws.rs.{Consumes, GET, POST, Path, Produces}
+import org.apache.texera.auth.JwtParser.parseToken
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.auth.util.{ComputingUnitAccess, HeaderField}
+import org.apache.texera.config.{GuiConfig, KubernetesConfig, LLMConfig}
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+
+import java.net.URLDecoder
+import java.nio.charset.StandardCharsets
+import java.util.Optional
+import scala.jdk.CollectionConverters.{CollectionHasAsScala, MapHasAsScala}
+import scala.util.matching.Regex
+
+object AccessControlResource extends LazyLogging {
+
+ private val mapper: ObjectMapper = new ObjectMapper().registerModule(DefaultScalaModule)
+
+ // Regex for the paths that require authorization
+ private val wsapiWorkflowWebsocket: Regex = """.*/wsapi/workflow-websocket.*""".r
+ private val apiExecutionsStats: Regex = """.*/api/executions/[0-9]+/stats/[0-9]+.*""".r
+ private val apiExecutionsResultExport: Regex = """.*/api/executions/result/export.*""".r
+
+ /**
+ * Authorize the request based on the path and headers.
+ * @param uriInfo URI sent by Envoy or API Gateway
+ * @param headers HTTP headers sent by Envoy or API Gateway which include
+ * headers sent by the client (browser)
+ * @return HTTP Response with appropriate status code and headers
+ */
+ def authorize(
+ uriInfo: UriInfo,
+ headers: HttpHeaders,
+ bodyOpt: Option[String] = None
+ ): Response = {
+ val path = uriInfo.getPath
+ logger.info(s"Authorizing request for path: $path")
+
+ path match {
+ case wsapiWorkflowWebsocket() | apiExecutionsStats() | apiExecutionsResultExport() =>
+ checkComputingUnitAccess(uriInfo, headers, bodyOpt)
+ case _ =>
+ logger.warn(s"No authorization logic for path: $path. Denying access.")
+ Response.status(Response.Status.FORBIDDEN).build()
+ }
+ }
+
+ private def checkComputingUnitAccess(
+ uriInfo: UriInfo,
+ headers: HttpHeaders,
+ bodyOpt: Option[String]
+ ): Response = {
+ val queryParams: Map[String, String] = uriInfo
+ .getQueryParameters()
+ .asScala
+ .view
+ .mapValues(values => values.asScala.headOption.getOrElse(""))
+ .toMap
+
+ logger.info(
+ s"Request URI: ${uriInfo.getRequestUri} and headers: ${headers.getRequestHeaders.asScala} and queryParams: $queryParams"
+ )
+
+ val token: String = {
+ val qToken = queryParams.get("access-token").filter(_.nonEmpty)
+ val hToken = Option(headers.getRequestHeader("Authorization"))
+ .flatMap(_.asScala.headOption)
+ .map(_.replaceFirst("(?i)^Bearer\\s+", "")) // case-insensitive "Bearer "
+ .map(_.trim)
+ .filter(_.nonEmpty)
+ val bToken = bodyOpt.flatMap(extractTokenFromBody)
+ qToken.orElse(hToken).orElse(bToken).getOrElse("")
+ }
+ logger.info(s"token extracted from request $token")
+ val cuid = queryParams.getOrElse("cuid", "")
+ val cuidInt =
+ try {
+ cuid.toInt
+ } catch {
+ case _: NumberFormatException =>
+ return Response.status(Response.Status.FORBIDDEN).build()
+ }
+
+ var cuAccess: PrivilegeEnum = PrivilegeEnum.NONE
+ var userSession: Optional[SessionUser] = Optional.empty()
+ try {
+ userSession = parseToken(token)
+ if (userSession.isEmpty)
+ return Response.status(Response.Status.FORBIDDEN).build()
+
+ val uid = userSession.get().getUid
+ cuAccess = ComputingUnitAccess.getComputingUnitAccess(cuidInt, uid)
+ if (cuAccess == PrivilegeEnum.NONE)
+ return Response.status(Response.Status.FORBIDDEN).build()
+ } catch {
+ case e: Exception =>
+ logger.error(s"Failed parsing token $e")
+ return Response.status(Response.Status.FORBIDDEN).build()
+ }
+
+ // Dynamic Routing Logic
+ val workflowComputingUnitPoolName = KubernetesConfig.computeUnitPoolName
+ val workflowComputingUnitPoolNamespace = KubernetesConfig.computeUnitPoolNamespace
+ val workflowComputingUnitPoolPort = KubernetesConfig.computeUnitPortNumber
+
+ val targetHost =
+ s"computing-unit-$cuidInt.$workflowComputingUnitPoolName-svc.$workflowComputingUnitPoolNamespace.svc.cluster.local:$workflowComputingUnitPoolPort"
+
+ Response
+ .ok()
+ .header(HeaderField.UserComputingUnitAccess, cuAccess.toString)
+ .header(HeaderField.UserId, userSession.get().getUid.toString)
+ .header(HeaderField.UserName, userSession.get().getName)
+ .header(HeaderField.UserEmail, userSession.get().getEmail)
+ .header("Host", targetHost) // Envoy ExtAuth: Rewrite Host
+ .build()
+ }
+
+ // Extracts a top-level "token" field from a JSON body
+ private def extractTokenFromBody(body: String): Option[String] = {
+ // 1) Try JSON
+ val jsonToken: Option[String] =
+ try {
+ val node = mapper.readTree(body)
+ if (node != null && node.has("token"))
+ Option(node.get("token").asText()).map(_.trim).filter(_.nonEmpty)
+ else None
+ } catch {
+ case _: Exception => None
+ }
+
+ // 2) Try application/x-www-form-urlencoded
+ def extractTokenFromUrlEncoded(s: String): Option[String] = {
+ // fast path: must contain '=' or '&'
+ if (!s.contains("=")) return None
+ val pairs = s.split("&").iterator
+ var found: Option[String] = None
+ while (pairs.hasNext && found.isEmpty) {
+ val p = pairs.next()
+ val idx = p.indexOf('=')
+ val key = if (idx >= 0) p.substring(0, idx) else p
+ if (key == "token") {
+ val raw = if (idx >= 0) p.substring(idx + 1) else ""
+ val decoded = URLDecoder.decode(raw, StandardCharsets.UTF_8.name())
+ val v = decoded.trim
+ if (v.nonEmpty) found = Some(v)
+ }
+ }
+ found
+ }
+
+ // 3) Try multipart/form-data (best-effort; parses raw body text)
+ def extractTokenFromMultipart(s: String): Option[String] = {
+ // Look for the part with name="token" and capture its content until the next boundary
+ val partWithBoundary = "(?s)name\\s*=\\s*\"token\"[^\\r\\n]*\\r?\\n\\r?\\n(.*?)\\r?\\n--".r
+ val partToEnd = "(?s)name\\s*=\\s*\"token\"[^\\r\\n]*\\r?\\n\\r?\\n(.*)".r
+
+ partWithBoundary
+ .findFirstMatchIn(s)
+ .map(_.group(1).trim)
+ .filter(_.nonEmpty)
+ .orElse(partToEnd.findFirstMatchIn(s).map(_.group(1).trim).filter(_.nonEmpty))
+ }
+
+ jsonToken
+ .orElse(extractTokenFromUrlEncoded(body))
+ .orElse(extractTokenFromMultipart(body))
+ }
+}
+@Produces(Array(MediaType.APPLICATION_JSON))
+@Path("/auth")
+class AccessControlResource extends LazyLogging {
+
+ @GET
+ @Path("/{path:.*}")
+ def authorizeGet(
+ @Context uriInfo: UriInfo,
+ @Context headers: HttpHeaders
+ ): Response = {
+ AccessControlResource.authorize(uriInfo, headers)
+ }
+
+ @POST
+ @Path("/{path:.*}")
+ def authorizePost(
+ @Context uriInfo: UriInfo,
+ @Context headers: HttpHeaders,
+ body: String
+ ): Response = {
+ logger.info("Request body: " + body)
+ AccessControlResource.authorize(uriInfo, headers, Option(body).map(_.trim).filter(_.nonEmpty))
+ }
+}
+
+@Path("/chat")
+@Produces(Array(MediaType.APPLICATION_JSON))
+@Consumes(Array(MediaType.APPLICATION_JSON))
+class LiteLLMProxyResource extends LazyLogging {
+
+ private val client: Client = ClientBuilder.newClient()
+ private val litellmBaseUrl: String = LLMConfig.baseUrl
+ private val litellmApiKey: String = LLMConfig.masterKey
+
+ @POST
+ @Path("/{path:.*}")
+ def proxyPost(
+ @Context uriInfo: UriInfo,
+ @Context headers: HttpHeaders,
+ body: String
+ ): Response = {
+ if (!GuiConfig.guiWorkflowWorkspaceCopilotEnabled) {
+ return Response
+ .status(Response.Status.FORBIDDEN)
+ .entity("""{"error": "Copilot feature is disabled"}""")
+ .build()
+ }
+
+ // uriInfo.getPath returns "chat/completions" for /api/chat/completions
+ // We want to forward as "/chat/completions" to LiteLLM
+ val fullPath = uriInfo.getPath
+ val targetUrl = s"$litellmBaseUrl/$fullPath"
+
+ logger.info(s"Proxying POST request to LiteLLM: $targetUrl")
+
+ try {
+ val requestBuilder = client
+ .target(targetUrl)
+ .request(MediaType.APPLICATION_JSON)
+ .header("Authorization", s"Bearer $litellmApiKey")
+
+ // Forward other relevant headers from the original request
+ headers.getRequestHeaders.asScala.foreach {
+ case (key, values)
+ if !key.equalsIgnoreCase("Authorization") &&
+ !key.equalsIgnoreCase("Host") &&
+ !key.equalsIgnoreCase("Content-Length") =>
+ values.asScala.foreach(value => requestBuilder.header(key, value))
+ case _ => // Skip Authorization, Host, and Content-Length headers
+ }
+
+ val response = requestBuilder.post(Entity.json(body))
+
+ // Build response with same status and body from LiteLLM
+ val responseBody = response.readEntity(classOf[String])
+ val responseBuilder = Response
+ .status(response.getStatus)
+ .entity(responseBody)
+
+ // Forward response headers
+ response.getHeaders.asScala.foreach {
+ case (key, values) =>
+ values.asScala.foreach(value => responseBuilder.header(key, value))
+ }
+
+ responseBuilder.build()
+ } catch {
+ case e: Exception =>
+ logger.error(s"Error proxying request to LiteLLM: ${e.getMessage}", e)
+ Response
+ .status(Response.Status.BAD_GATEWAY)
+ .entity(s"""{"error": "Failed to proxy request to LiteLLM: ${e.getMessage}"}""")
+ .build()
+ }
+ }
+}
+
+@Path("/models")
+@Produces(Array(MediaType.APPLICATION_JSON))
+class LiteLLMModelsResource extends LazyLogging {
+
+ private val client: Client = ClientBuilder.newClient()
+ private val litellmBaseUrl: String = LLMConfig.baseUrl
+ private val litellmApiKey: String = LLMConfig.masterKey
+
+ @GET
+ def getModels: Response = {
+ if (!GuiConfig.guiWorkflowWorkspaceCopilotEnabled) {
+ return Response
+ .status(Response.Status.FORBIDDEN)
+ .entity("""{"error": "Copilot feature is disabled"}""")
+ .build()
+ }
+
+ val targetUrl = s"$litellmBaseUrl/models"
+
+ logger.info(s"Fetching models from LiteLLM: $targetUrl")
+
+ try {
+ val response = client
+ .target(targetUrl)
+ .request(MediaType.APPLICATION_JSON)
+ .header("Authorization", s"Bearer $litellmApiKey")
+ .get()
+
+ // Build response with same status and body from LiteLLM
+ val responseBody = response.readEntity(classOf[String])
+ val responseBuilder = Response
+ .status(response.getStatus)
+ .entity(responseBody)
+
+ // Forward response headers
+ response.getHeaders.asScala.foreach {
+ case (key, values) =>
+ values.asScala.foreach(value => responseBuilder.header(key, value))
+ }
+
+ responseBuilder.build()
+ } catch {
+ case e: Exception =>
+ logger.error(s"Error fetching models from LiteLLM: ${e.getMessage}", e)
+ Response
+ .status(Response.Status.BAD_GATEWAY)
+ .entity(s"""{"error": "Failed to fetch models from LiteLLM: ${e.getMessage}"}""")
+ .build()
+ }
+ }
+}
diff --git a/access-control-service/src/main/scala/org/apache/texera/service/resource/HealthCheckResource.scala b/access-control-service/src/main/scala/org/apache/texera/service/resource/HealthCheckResource.scala
new file mode 100644
index 00000000000..08ecd1cd0cb
--- /dev/null
+++ b/access-control-service/src/main/scala/org/apache/texera/service/resource/HealthCheckResource.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.service.resource
+
+import jakarta.ws.rs.core.MediaType
+import jakarta.ws.rs.{GET, Path, Produces}
+
+@Path("/healthcheck")
+@Produces(Array(MediaType.APPLICATION_JSON))
+class HealthCheckResource {
+ @GET
+ def healthCheck: Map[String, String] = Map("status" -> "ok")
+}
diff --git a/access-control-service/src/test/scala/org/apache/texera/AccessControlResourceSpec.scala b/access-control-service/src/test/scala/org/apache/texera/AccessControlResourceSpec.scala
new file mode 100644
index 00000000000..751b51022f4
--- /dev/null
+++ b/access-control-service/src/test/scala/org/apache/texera/AccessControlResourceSpec.scala
@@ -0,0 +1,236 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.texera
+
+import jakarta.ws.rs.core.{HttpHeaders, MultivaluedHashMap, Response, UriInfo}
+import org.apache.texera.auth.JwtAuth
+import org.apache.texera.auth.util.HeaderField
+import org.apache.texera.dao.MockTexeraDB
+import org.apache.texera.dao.jooq.generated.enums.{
+ PrivilegeEnum,
+ UserRoleEnum,
+ WorkflowComputingUnitTypeEnum
+}
+import org.apache.texera.dao.jooq.generated.tables.daos.{
+ ComputingUnitUserAccessDao,
+ UserDao,
+ WorkflowComputingUnitDao
+}
+import org.apache.texera.dao.jooq.generated.tables.pojos.{
+ ComputingUnitUserAccess,
+ User,
+ WorkflowComputingUnit
+}
+import org.apache.texera.service.resource.AccessControlResource
+import org.mockito.Mockito._
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
+
+import java.net.URI
+import java.util
+
+class AccessControlResourceSpec
+ extends AnyFlatSpec
+ with Matchers
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with MockTexeraDB {
+
+ private val testURI: String = "http://localhost:8080/"
+ private val testPath: String = "/api/executions/1/stats/1"
+
+ private val testUser1: User = {
+ val user = new User()
+ user.setUid(1)
+ user.setName("testuser")
+ user.setEmail("test@example.com")
+ user.setRole(UserRoleEnum.REGULAR)
+ user.setPassword("password")
+ user
+ }
+
+ private val testUser2: User = {
+ val user = new User()
+ user.setUid(2)
+ user.setName("testuser2")
+ user.setEmail("test2@example.com")
+ user.setRole(UserRoleEnum.REGULAR)
+ user.setPassword("password")
+ user
+ }
+
+ private val testCU: WorkflowComputingUnit = {
+ val cu = new WorkflowComputingUnit()
+ cu.setUid(2)
+ cu.setType(WorkflowComputingUnitTypeEnum.kubernetes)
+ cu.setCuid(2)
+ cu.setName("test-cu")
+ cu
+ }
+
+ private var token: String = _
+
+ override protected def beforeAll(): Unit = {
+ initializeDBAndReplaceDSLContext()
+ val userDao = new UserDao(getDSLContext.configuration())
+ val computingUnitDao = new WorkflowComputingUnitDao(getDSLContext.configuration())
+ val computingUnitOfUserDao = new ComputingUnitUserAccessDao(getDSLContext.configuration())
+
+ // insert user, computing unit, and access privilege into the mock database
+ userDao.insert(testUser1)
+ userDao.insert(testUser2)
+ computingUnitDao.insert(testCU)
+
+ val cuAccess = new ComputingUnitUserAccess()
+ cuAccess.setUid(testUser1.getUid)
+ cuAccess.setCuid(testCU.getCuid)
+ cuAccess.setPrivilege(PrivilegeEnum.WRITE)
+ computingUnitOfUserDao.insert(cuAccess)
+
+ val claims = JwtAuth.jwtClaims(testUser1, 1)
+ token = JwtAuth.jwtToken(claims)
+ }
+
+ override protected def afterAll(): Unit = {
+ shutdownDB()
+ }
+
+ "AccessControlResource" should "return FORBIDDEN for a GET request without a token" in {
+ val mockUriInfo = mock(classOf[UriInfo])
+ val mockHttpHeaders = mock(classOf[HttpHeaders])
+ val queryParams = new MultivaluedHashMap[String, String]()
+ queryParams.add("cuid", "1")
+ val requestHeaders = new MultivaluedHashMap[String, String]()
+
+ when(mockUriInfo.getQueryParameters).thenReturn(queryParams)
+ when(mockUriInfo.getRequestUri).thenReturn(new URI(testURI))
+ when(mockUriInfo.getPath).thenReturn(testPath)
+ when(mockHttpHeaders.getRequestHeaders).thenReturn(requestHeaders)
+ when(mockHttpHeaders.getRequestHeader("Authorization")).thenReturn(new util.ArrayList[String]())
+
+ val accessControlResource = new AccessControlResource()
+ val response = accessControlResource.authorizeGet(mockUriInfo, mockHttpHeaders)
+
+ response.getStatus shouldBe Response.Status.FORBIDDEN.getStatusCode
+ }
+
+ it should "return FORBIDDEN for a GET request with a non-integer cuid" in {
+ val mockUriInfo = mock(classOf[UriInfo])
+ val mockHttpHeaders = mock(classOf[HttpHeaders])
+ val queryParams = new MultivaluedHashMap[String, String]()
+ queryParams.add("cuid", "abc")
+ val requestHeaders = new MultivaluedHashMap[String, String]()
+ requestHeaders.add("Authorization", "Bearer dummy-token")
+
+ when(mockUriInfo.getQueryParameters).thenReturn(queryParams)
+ when(mockUriInfo.getRequestUri).thenReturn(new URI(testURI))
+ when(mockUriInfo.getPath).thenReturn(testPath)
+ when(mockHttpHeaders.getRequestHeaders).thenReturn(requestHeaders)
+ when(mockHttpHeaders.getRequestHeader("Authorization"))
+ .thenReturn(util.Arrays.asList("Bearer dummy-token"))
+
+ val accessControlResource = new AccessControlResource()
+ val response = accessControlResource.authorizeGet(mockUriInfo, mockHttpHeaders)
+
+ response.getStatus shouldBe Response.Status.FORBIDDEN.getStatusCode
+ }
+
+ it should "return FORBIDDEN for a POST request without a token" in {
+ val mockUriInfo = mock(classOf[UriInfo])
+ val mockHttpHeaders = mock(classOf[HttpHeaders])
+ val queryParams = new MultivaluedHashMap[String, String]()
+ queryParams.add("cuid", "1")
+ val requestHeaders = new MultivaluedHashMap[String, String]()
+
+ when(mockUriInfo.getQueryParameters).thenReturn(queryParams)
+ when(mockUriInfo.getRequestUri).thenReturn(new URI(testURI))
+ when(mockUriInfo.getPath).thenReturn(testPath)
+ when(mockHttpHeaders.getRequestHeaders).thenReturn(requestHeaders)
+ when(mockHttpHeaders.getRequestHeader("Authorization")).thenReturn(new util.ArrayList[String]())
+
+ val accessControlResource = new AccessControlResource()
+ val response = accessControlResource.authorizePost(mockUriInfo, mockHttpHeaders, null)
+
+ response.getStatus shouldBe Response.Status.FORBIDDEN.getStatusCode
+ }
+
+ "AccessControlResource" should "return FORBIDDEN when user does not have access to the computing unit" in {
+ // Mock the request context
+ val mockUriInfo = mock(classOf[UriInfo])
+ val mockHttpHeaders = mock(classOf[HttpHeaders])
+
+ // Prepare query parameters with a computing unit ID (cuid)
+ val queryParams = new MultivaluedHashMap[String, String]()
+ queryParams.add("cuid", "1") // Assuming user 1 does not have access to cuid 1
+
+ // Prepare request headers with the generated JWT
+ val requestHeaders = new MultivaluedHashMap[String, String]()
+ requestHeaders.add("Authorization", "Bearer " + token)
+
+ // Stub the mock objects to return the prepared data
+ when(mockUriInfo.getQueryParameters).thenReturn(queryParams)
+ when(mockUriInfo.getRequestUri).thenReturn(new URI(testURI))
+ when(mockUriInfo.getPath).thenReturn(testPath)
+ when(mockHttpHeaders.getRequestHeaders).thenReturn(requestHeaders)
+ when(mockHttpHeaders.getRequestHeader("Authorization"))
+ .thenReturn(util.Arrays.asList("Bearer " + token))
+
+ // Instantiate the resource and call the method under test
+ val accessControlResource = new AccessControlResource()
+ val response = accessControlResource.authorizeGet(mockUriInfo, mockHttpHeaders)
+
+ // Assert that the response status is FORBIDDEN
+ response.getStatus shouldBe Response.Status.FORBIDDEN.getStatusCode
+ }
+
+ it should "return OK and correct headers when user has access" in {
+ // Mock the request context
+ val mockUriInfo = mock(classOf[UriInfo])
+ val mockHttpHeaders = mock(classOf[HttpHeaders])
+
+ // Prepare query parameters with a computing unit ID the user HAS access to
+ val queryParams = new MultivaluedHashMap[String, String]()
+ queryParams.add("cuid", testCU.getCuid.toString)
+
+ // Prepare request headers with the generated JWT
+ val requestHeaders = new MultivaluedHashMap[String, String]()
+ requestHeaders.add("Authorization", "Bearer " + token)
+
+ // Stub the mock objects to return the prepared data
+ when(mockUriInfo.getQueryParameters).thenReturn(queryParams)
+ when(mockUriInfo.getRequestUri).thenReturn(new URI(testURI))
+ when(mockUriInfo.getPath).thenReturn(testPath)
+ when(mockHttpHeaders.getRequestHeaders).thenReturn(requestHeaders)
+ when(mockHttpHeaders.getRequestHeader("Authorization"))
+ .thenReturn(util.Arrays.asList("Bearer " + token))
+
+ // Instantiate the resource and call the method under test
+ val accessControlResource = new AccessControlResource()
+ val response = accessControlResource.authorizeGet(mockUriInfo, mockHttpHeaders)
+
+ // Assert that the response status is OK and headers are correct
+ response.getStatus shouldBe Response.Status.OK.getStatusCode
+ response.getHeaderString(
+ HeaderField.UserComputingUnitAccess
+ ) shouldBe PrivilegeEnum.WRITE.toString
+ response.getHeaderString(HeaderField.UserId) shouldBe testUser1.getUid.toString
+ response.getHeaderString(HeaderField.UserName) shouldBe testUser1.getName
+ response.getHeaderString(HeaderField.UserEmail) shouldBe testUser1.getEmail
+ }
+}
diff --git a/agent-service/.dockerignore b/agent-service/.dockerignore
new file mode 100644
index 00000000000..44b274308ab
--- /dev/null
+++ b/agent-service/.dockerignore
@@ -0,0 +1,2 @@
+node_modules
+.yarn
diff --git a/agent-service/.env.example b/agent-service/.env.example
new file mode 100644
index 00000000000..605f157ca0c
--- /dev/null
+++ b/agent-service/.env.example
@@ -0,0 +1,21 @@
+# agent-service's webserver configuration
+PORT=3001
+API_PREFIX=/api
+
+# ERROR | WARN | INFO | DEBUG
+TEXERA_SERVICE_LOG_LEVEL=INFO
+
+# Human-readable dev logs via pino-pretty.
+LOG_PRETTY=true
+
+# LLM_API_KEY authenticates this service to the gateway — NOT the
+# upstream provider key (OpenAI / Anthropic / etc.), which is
+# configured inside the gateway. "dummy" works when the gateway
+# does not enforce a key (e.g. LiteLLM with no master_key).
+LLM_API_KEY=dummy
+LLM_ENDPOINT=http://localhost:9096
+
+# Texera backend services
+TEXERA_DASHBOARD_SERVICE_ENDPOINT=http://localhost:8080
+WORKFLOW_COMPILING_SERVICE_ENDPOINT=http://localhost:9090
+WORKFLOW_EXECUTION_SERVICE_ENDPOINT=http://localhost:8085
\ No newline at end of file
diff --git a/agent-service/.prettierrc b/agent-service/.prettierrc
new file mode 100644
index 00000000000..a4b3b08e4e5
--- /dev/null
+++ b/agent-service/.prettierrc
@@ -0,0 +1,13 @@
+{
+ "printWidth": 120,
+ "tabWidth": 2,
+ "useTabs": false,
+ "semi": true,
+ "singleQuote": false,
+ "quoteProps": "as-needed",
+ "trailingComma": "es5",
+ "bracketSameLine": true,
+ "bracketSpacing": true,
+ "arrowParens": "avoid",
+ "endOfLine": "lf"
+}
diff --git a/agent-service/bin/collect-licenses.ts b/agent-service/bin/collect-licenses.ts
new file mode 100755
index 00000000000..9451091d0f8
--- /dev/null
+++ b/agent-service/bin/collect-licenses.ts
@@ -0,0 +1,87 @@
+#!/usr/bin/env bun
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Walk node_modules and emit a {name, version, license}[] manifest in
+// the same shape license-webpack-plugin produces for the frontend.
+// Run after `bun install --production --frozen-lockfile`. Output goes
+// to stdout so a CI step can redirect it to dist/3rdpartylicenses.json.
+
+import { readdir, readFile, stat } from "node:fs/promises";
+import { join } from "node:path";
+
+type Entry = { name: string; version: string; license: string };
+
+function normalizeLicense(license: unknown): string {
+ if (typeof license === "string") return license;
+ if (license && typeof license === "object") {
+ // legacy { type, url } form, or { license: "X", licenses: [...] }
+ const obj = license as Record;
+ if (typeof obj.type === "string") return obj.type;
+ if (Array.isArray(obj)) {
+ return obj.map((l) => normalizeLicense(l)).filter(Boolean).join(" OR ");
+ }
+ }
+ return "UNKNOWN";
+}
+
+async function readPackageJson(dir: string): Promise {
+ try {
+ const raw = await readFile(join(dir, "package.json"), "utf8");
+ const pkg = JSON.parse(raw);
+ if (!pkg.name || !pkg.version) return null;
+ const license = pkg.license ?? pkg.licenses ?? "UNKNOWN";
+ return {
+ name: pkg.name,
+ version: pkg.version,
+ license: normalizeLicense(license),
+ };
+ } catch {
+ return null;
+ }
+}
+
+async function walk(nm: string): Promise {
+ const entries: Entry[] = [];
+ const top = await readdir(nm);
+ for (const name of top) {
+ if (name.startsWith(".")) continue;
+ const path = join(nm, name);
+ const st = await stat(path);
+ if (!st.isDirectory()) continue;
+ if (name.startsWith("@")) {
+ // scoped: walk one more level
+ const inner = await readdir(path);
+ for (const sub of inner) {
+ if (sub.startsWith(".")) continue;
+ const e = await readPackageJson(join(path, sub));
+ if (e) entries.push(e);
+ }
+ } else {
+ const e = await readPackageJson(path);
+ if (e) entries.push(e);
+ }
+ }
+ return entries;
+}
+
+const nm = join(import.meta.dir, "..", "node_modules");
+const entries = await walk(nm);
+entries.sort((a, b) =>
+ a.name === b.name ? a.version.localeCompare(b.version) : a.name.localeCompare(b.name),
+);
+process.stdout.write(JSON.stringify(entries, null, 2) + "\n");
diff --git a/agent-service/bun.lock b/agent-service/bun.lock
new file mode 100644
index 00000000000..8a1953f46b2
--- /dev/null
+++ b/agent-service/bun.lock
@@ -0,0 +1,237 @@
+{
+ "lockfileVersion": 1,
+ "configVersion": 1,
+ "workspaces": {
+ "": {
+ "name": "texera-agent-service",
+ "dependencies": {
+ "@ai-sdk/openai": "2.0.79",
+ "@elysiajs/cors": "1.4.0",
+ "ai": "5.0.108",
+ "ajv": "8.10.0",
+ "dagre": "0.8.5",
+ "elysia": "1.4.18",
+ "pino": "10.3.1",
+ "rxjs": "7.8.2",
+ "zod": "3.25.76",
+ },
+ "devDependencies": {
+ "@types/bun": "1.3.3",
+ "@types/dagre": "0.7.54",
+ "pino-pretty": "13.1.3",
+ "prettier": "3.4.2",
+ "tsx": "4.21.0",
+ "typescript": "5.9.3",
+ },
+ },
+ },
+ "packages": {
+ "@ai-sdk/gateway": ["@ai-sdk/gateway@2.0.18", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.18", "@vercel/oidc": "3.0.5" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-sDQcW+6ck2m0pTIHW6BPHD7S125WD3qNkx/B8sEzJp/hurocmJ5Cni0ybExg6sQMGo+fr/GWOwpHF1cmCdg5rQ=="],
+
+ "@ai-sdk/openai": ["@ai-sdk/openai@2.0.79", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.18" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-RIUbwD2VGpawSKUuZ0HmsDbyjPniQIa9wYyE3xQ5fIWnI+RZH8MfyRwnUFom1pko5YOGlhosZhJmvolG8lNr7Q=="],
+
+ "@ai-sdk/provider": ["@ai-sdk/provider@2.0.0", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="],
+
+ "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@3.0.18", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-ypv1xXMsgGcNKUP+hglKqtdDuMg68nWHucPPAhIENrbFAI+xCHiqPVN8Zllxyv1TNZwGWUghPxJXU+Mqps0YRQ=="],
+
+ "@borewit/text-codec": ["@borewit/text-codec@0.1.1", "", {}, "sha512-5L/uBxmjaCIX5h8Z+uu+kA9BQLkc/Wl06UGR5ajNRxu+/XjonB5i8JpgFMrPj3LXTCPA0pv8yxUvbUi+QthGGA=="],
+
+ "@elysiajs/cors": ["@elysiajs/cors@1.4.0", "", { "peerDependencies": { "elysia": ">= 1.4.0" } }, "sha512-pb0SCzBfFbFSYA/U40HHO7R+YrcXBJXOWgL20eSViK33ol1e20ru2/KUaZYo5IMUn63yaTJI/bQERuQ+77ND8g=="],
+
+ "@esbuild/aix-ppc64": ["@esbuild/aix-ppc64@0.27.1", "", { "os": "aix", "cpu": "ppc64" }, "sha512-HHB50pdsBX6k47S4u5g/CaLjqS3qwaOVE5ILsq64jyzgMhLuCuZ8rGzM9yhsAjfjkbgUPMzZEPa7DAp7yz6vuA=="],
+
+ "@esbuild/android-arm": ["@esbuild/android-arm@0.27.1", "", { "os": "android", "cpu": "arm" }, "sha512-kFqa6/UcaTbGm/NncN9kzVOODjhZW8e+FRdSeypWe6j33gzclHtwlANs26JrupOntlcWmB0u8+8HZo8s7thHvg=="],
+
+ "@esbuild/android-arm64": ["@esbuild/android-arm64@0.27.1", "", { "os": "android", "cpu": "arm64" }, "sha512-45fuKmAJpxnQWixOGCrS+ro4Uvb4Re9+UTieUY2f8AEc+t7d4AaZ6eUJ3Hva7dtrxAAWHtlEFsXFMAgNnGU9uQ=="],
+
+ "@esbuild/android-x64": ["@esbuild/android-x64@0.27.1", "", { "os": "android", "cpu": "x64" }, "sha512-LBEpOz0BsgMEeHgenf5aqmn/lLNTFXVfoWMUox8CtWWYK9X4jmQzWjoGoNb8lmAYml/tQ/Ysvm8q7szu7BoxRQ=="],
+
+ "@esbuild/darwin-arm64": ["@esbuild/darwin-arm64@0.27.1", "", { "os": "darwin", "cpu": "arm64" }, "sha512-veg7fL8eMSCVKL7IW4pxb54QERtedFDfY/ASrumK/SbFsXnRazxY4YykN/THYqFnFwJ0aVjiUrVG2PwcdAEqQQ=="],
+
+ "@esbuild/darwin-x64": ["@esbuild/darwin-x64@0.27.1", "", { "os": "darwin", "cpu": "x64" }, "sha512-+3ELd+nTzhfWb07Vol7EZ+5PTbJ/u74nC6iv4/lwIU99Ip5uuY6QoIf0Hn4m2HoV0qcnRivN3KSqc+FyCHjoVQ=="],
+
+ "@esbuild/freebsd-arm64": ["@esbuild/freebsd-arm64@0.27.1", "", { "os": "freebsd", "cpu": "arm64" }, "sha512-/8Rfgns4XD9XOSXlzUDepG8PX+AVWHliYlUkFI3K3GB6tqbdjYqdhcb4BKRd7C0BhZSoaCxhv8kTcBrcZWP+xg=="],
+
+ "@esbuild/freebsd-x64": ["@esbuild/freebsd-x64@0.27.1", "", { "os": "freebsd", "cpu": "x64" }, "sha512-GITpD8dK9C+r+5yRT/UKVT36h/DQLOHdwGVwwoHidlnA168oD3uxA878XloXebK4Ul3gDBBIvEdL7go9gCUFzQ=="],
+
+ "@esbuild/linux-arm": ["@esbuild/linux-arm@0.27.1", "", { "os": "linux", "cpu": "arm" }, "sha512-ieMID0JRZY/ZeCrsFQ3Y3NlHNCqIhTprJfDgSB3/lv5jJZ8FX3hqPyXWhe+gvS5ARMBJ242PM+VNz/ctNj//eA=="],
+
+ "@esbuild/linux-arm64": ["@esbuild/linux-arm64@0.27.1", "", { "os": "linux", "cpu": "arm64" }, "sha512-W9//kCrh/6in9rWIBdKaMtuTTzNj6jSeG/haWBADqLLa9P8O5YSRDzgD5y9QBok4AYlzS6ARHifAb75V6G670Q=="],
+
+ "@esbuild/linux-ia32": ["@esbuild/linux-ia32@0.27.1", "", { "os": "linux", "cpu": "ia32" }, "sha512-VIUV4z8GD8rtSVMfAj1aXFahsi/+tcoXXNYmXgzISL+KB381vbSTNdeZHHHIYqFyXcoEhu9n5cT+05tRv13rlw=="],
+
+ "@esbuild/linux-loong64": ["@esbuild/linux-loong64@0.27.1", "", { "os": "linux", "cpu": "none" }, "sha512-l4rfiiJRN7sTNI//ff65zJ9z8U+k6zcCg0LALU5iEWzY+a1mVZ8iWC1k5EsNKThZ7XCQ6YWtsZ8EWYm7r1UEsg=="],
+
+ "@esbuild/linux-mips64el": ["@esbuild/linux-mips64el@0.27.1", "", { "os": "linux", "cpu": "none" }, "sha512-U0bEuAOLvO/DWFdygTHWY8C067FXz+UbzKgxYhXC0fDieFa0kDIra1FAhsAARRJbvEyso8aAqvPdNxzWuStBnA=="],
+
+ "@esbuild/linux-ppc64": ["@esbuild/linux-ppc64@0.27.1", "", { "os": "linux", "cpu": "ppc64" }, "sha512-NzdQ/Xwu6vPSf/GkdmRNsOfIeSGnh7muundsWItmBsVpMoNPVpM61qNzAVY3pZ1glzzAxLR40UyYM23eaDDbYQ=="],
+
+ "@esbuild/linux-riscv64": ["@esbuild/linux-riscv64@0.27.1", "", { "os": "linux", "cpu": "none" }, "sha512-7zlw8p3IApcsN7mFw0O1Z1PyEk6PlKMu18roImfl3iQHTnr/yAfYv6s4hXPidbDoI2Q0pW+5xeoM4eTCC0UdrQ=="],
+
+ "@esbuild/linux-s390x": ["@esbuild/linux-s390x@0.27.1", "", { "os": "linux", "cpu": "s390x" }, "sha512-cGj5wli+G+nkVQdZo3+7FDKC25Uh4ZVwOAK6A06Hsvgr8WqBBuOy/1s+PUEd/6Je+vjfm6stX0kmib5b/O2Ykw=="],
+
+ "@esbuild/linux-x64": ["@esbuild/linux-x64@0.27.1", "", { "os": "linux", "cpu": "x64" }, "sha512-z3H/HYI9MM0HTv3hQZ81f+AKb+yEoCRlUby1F80vbQ5XdzEMyY/9iNlAmhqiBKw4MJXwfgsh7ERGEOhrM1niMA=="],
+
+ "@esbuild/netbsd-arm64": ["@esbuild/netbsd-arm64@0.27.1", "", { "os": "none", "cpu": "arm64" }, "sha512-wzC24DxAvk8Em01YmVXyjl96Mr+ecTPyOuADAvjGg+fyBpGmxmcr2E5ttf7Im8D0sXZihpxzO1isus8MdjMCXQ=="],
+
+ "@esbuild/netbsd-x64": ["@esbuild/netbsd-x64@0.27.1", "", { "os": "none", "cpu": "x64" }, "sha512-1YQ8ybGi2yIXswu6eNzJsrYIGFpnlzEWRl6iR5gMgmsrR0FcNoV1m9k9sc3PuP5rUBLshOZylc9nqSgymI+TYg=="],
+
+ "@esbuild/openbsd-arm64": ["@esbuild/openbsd-arm64@0.27.1", "", { "os": "openbsd", "cpu": "arm64" }, "sha512-5Z+DzLCrq5wmU7RDaMDe2DVXMRm2tTDvX2KU14JJVBN2CT/qov7XVix85QoJqHltpvAOZUAc3ndU56HSMWrv8g=="],
+
+ "@esbuild/openbsd-x64": ["@esbuild/openbsd-x64@0.27.1", "", { "os": "openbsd", "cpu": "x64" }, "sha512-Q73ENzIdPF5jap4wqLtsfh8YbYSZ8Q0wnxplOlZUOyZy7B4ZKW8DXGWgTCZmF8VWD7Tciwv5F4NsRf6vYlZtqg=="],
+
+ "@esbuild/openharmony-arm64": ["@esbuild/openharmony-arm64@0.27.1", "", { "os": "none", "cpu": "arm64" }, "sha512-ajbHrGM/XiK+sXM0JzEbJAen+0E+JMQZ2l4RR4VFwvV9JEERx+oxtgkpoKv1SevhjavK2z2ReHk32pjzktWbGg=="],
+
+ "@esbuild/sunos-x64": ["@esbuild/sunos-x64@0.27.1", "", { "os": "sunos", "cpu": "x64" }, "sha512-IPUW+y4VIjuDVn+OMzHc5FV4GubIwPnsz6ubkvN8cuhEqH81NovB53IUlrlBkPMEPxvNnf79MGBoz8rZ2iW8HA=="],
+
+ "@esbuild/win32-arm64": ["@esbuild/win32-arm64@0.27.1", "", { "os": "win32", "cpu": "arm64" }, "sha512-RIVRWiljWA6CdVu8zkWcRmGP7iRRIIwvhDKem8UMBjPql2TXM5PkDVvvrzMtj1V+WFPB4K7zkIGM7VzRtFkjdg=="],
+
+ "@esbuild/win32-ia32": ["@esbuild/win32-ia32@0.27.1", "", { "os": "win32", "cpu": "ia32" }, "sha512-2BR5M8CPbptC1AK5JbJT1fWrHLvejwZidKx3UMSF0ecHMa+smhi16drIrCEggkgviBwLYd5nwrFLSl5Kho96RQ=="],
+
+ "@esbuild/win32-x64": ["@esbuild/win32-x64@0.27.1", "", { "os": "win32", "cpu": "x64" }, "sha512-d5X6RMYv6taIymSk8JBP+nxv8DQAMY6A51GPgusqLdK9wBz5wWIXy1KjTck6HnjE9hqJzJRdk+1p/t5soSbCtw=="],
+
+ "@opentelemetry/api": ["@opentelemetry/api@1.9.0", "", {}, "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg=="],
+
+ "@pinojs/redact": ["@pinojs/redact@0.4.0", "", {}, "sha512-k2ENnmBugE/rzQfEcdWHcCY+/FM3VLzH9cYEsbdsoqrvzAKRhUZeRNhAZvB8OitQJ1TBed3yqWtdjzS6wJKBwg=="],
+
+ "@sinclair/typebox": ["@sinclair/typebox@0.34.41", "", {}, "sha512-6gS8pZzSXdyRHTIqoqSVknxolr1kzfy4/CeDnrzsVz8TTIWUbOBr6gnzOmTYJ3eXQNh4IYHIGi5aIL7sOZ2G/g=="],
+
+ "@standard-schema/spec": ["@standard-schema/spec@1.0.0", "", {}, "sha512-m2bOd0f2RT9k8QJx1JN85cZYyH1RqFBdlwtkSlf4tBDYLCiiZnv1fIIwacK6cqwXavOydf0NPToMQgpKq+dVlA=="],
+
+ "@tokenizer/inflate": ["@tokenizer/inflate@0.4.1", "", { "dependencies": { "debug": "^4.4.3", "token-types": "^6.1.1" } }, "sha512-2mAv+8pkG6GIZiF1kNg1jAjh27IDxEPKwdGul3snfztFerfPGI1LjDezZp3i7BElXompqEtPmoPx6c2wgtWsOA=="],
+
+ "@tokenizer/token": ["@tokenizer/token@0.3.0", "", {}, "sha512-OvjF+z51L3ov0OyAU0duzsYuvO01PH7x4t6DJx+guahgTnBHkhJdG7soQeTSFLWN3efnHyibZ4Z8l2EuWwJN3A=="],
+
+ "@types/bun": ["@types/bun@1.3.3", "", { "dependencies": { "bun-types": "1.3.3" } }, "sha512-ogrKbJ2X5N0kWLLFKeytG0eHDleBYtngtlbu9cyBKFtNL3cnpDZkNdQj8flVf6WTZUX5ulI9AY1oa7ljhSrp+g=="],
+
+ "@types/dagre": ["@types/dagre@0.7.54", "", {}, "sha512-QjcRY+adGbYvBFS7cwv5txhVIwX1XXIUswWl+kSQTbI6NjgZydrZkEKX/etzVd7i+bCsCb40Z/xlBY5eoFuvWQ=="],
+
+ "@types/node": ["@types/node@24.10.1", "", { "dependencies": { "undici-types": "~7.16.0" } }, "sha512-GNWcUTRBgIRJD5zj+Tq0fKOJ5XZajIiBroOF0yvj2bSU1WvNdYS/dn9UxwsujGW4JX06dnHyjV2y9rRaybH0iQ=="],
+
+ "@vercel/oidc": ["@vercel/oidc@3.0.5", "", {}, "sha512-fnYhv671l+eTTp48gB4zEsTW/YtRgRPnkI2nT7x6qw5rkI1Lq2hTmQIpHPgyThI0znLK+vX2n9XxKdXZ7BUbbw=="],
+
+ "ai": ["ai@5.0.108", "", { "dependencies": { "@ai-sdk/gateway": "2.0.18", "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.18", "@opentelemetry/api": "1.9.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-Jex3Lb7V41NNpuqJHKgrwoU6BCLHdI1Pg4qb4GJH4jRIDRXUBySJErHjyN4oTCwbiYCeb/8II9EnqSRPq9EifA=="],
+
+ "ajv": ["ajv@8.17.1", "", { "dependencies": { "fast-deep-equal": "^3.1.3", "fast-uri": "^3.0.1", "json-schema-traverse": "^1.0.0", "require-from-string": "^2.0.2" } }, "sha512-B/gBuNg5SiMTrPkC+A2+cW0RszwxYmn6VYxB/inlBStS5nx6xHIt/ehKRhIMhqusl7a8LjQoZnjCs5vhwxOQ1g=="],
+
+ "atomic-sleep": ["atomic-sleep@1.0.0", "", {}, "sha512-kNOjDqAh7px0XWNI+4QbzoiR/nTkHAWNud2uvnJquD1/x5a7EQZMJT0AczqK0Qn67oY/TTQ1LbUKajZpp3I9tQ=="],
+
+ "bun-types": ["bun-types@1.3.3", "", { "dependencies": { "@types/node": "*" } }, "sha512-z3Xwlg7j2l9JY27x5Qn3Wlyos8YAp0kKRlrePAOjgjMGS5IG6E7Jnlx736vH9UVI4wUICwwhC9anYL++XeOgTQ=="],
+
+ "colorette": ["colorette@2.0.20", "", {}, "sha512-IfEDxwoWIjkeXL1eXcDiow4UbKjhLdq6/EuSVR9GMN7KVH3r9gQ83e73hsz1Nd1T3ijd5xv1wcWRYO+D6kCI2w=="],
+
+ "cookie": ["cookie@1.1.1", "", {}, "sha512-ei8Aos7ja0weRpFzJnEA9UHJ/7XQmqglbRwnf2ATjcB9Wq874VKH9kfjjirM6UhU2/E5fFYadylyhFldcqSidQ=="],
+
+ "dagre": ["dagre@0.8.5", "", { "dependencies": { "graphlib": "^2.1.8", "lodash": "^4.17.15" } }, "sha512-/aTqmnRta7x7MCCpExk7HQL2O4owCT2h8NT//9I1OQ9vt29Pa0BzSAkR5lwFUcQ7491yVi/3CXU9jQ5o0Mn2Sw=="],
+
+ "dateformat": ["dateformat@4.6.3", "", {}, "sha512-2P0p0pFGzHS5EMnhdxQi7aJN+iMheud0UhG4dlE1DLAlvL8JHjJJTX/CSm4JXwV0Ka5nGk3zC5mcb5bUQUxxMA=="],
+
+ "debug": ["debug@4.4.3", "", { "dependencies": { "ms": "^2.1.3" } }, "sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA=="],
+
+ "elysia": ["elysia@1.4.18", "", { "dependencies": { "cookie": "^1.1.1", "exact-mirror": "0.2.5", "fast-decode-uri-component": "^1.0.1", "memoirist": "^0.4.0" }, "peerDependencies": { "@sinclair/typebox": ">= 0.34.0 < 1", "@types/bun": ">= 1.2.0", "file-type": ">= 20.0.0", "openapi-types": ">= 12.0.0", "typescript": ">= 5.0.0" }, "optionalPeers": ["@types/bun", "typescript"] }, "sha512-A6BhlipmSvgCy69SBgWADYZSdDIj3fT2gk8/9iMAC8iD+aGcnCr0fitziX0xr36MFDs/fsvVp8dWqxeq1VCgKg=="],
+
+ "end-of-stream": ["end-of-stream@1.4.5", "", { "dependencies": { "once": "^1.4.0" } }, "sha512-ooEGc6HP26xXq/N+GCGOT0JKCLDGrq2bQUZrQ7gyrJiZANJ/8YDTxTpQBXGMn+WbIQXNVpyWymm7KYVICQnyOg=="],
+
+ "esbuild": ["esbuild@0.27.1", "", { "optionalDependencies": { "@esbuild/aix-ppc64": "0.27.1", "@esbuild/android-arm": "0.27.1", "@esbuild/android-arm64": "0.27.1", "@esbuild/android-x64": "0.27.1", "@esbuild/darwin-arm64": "0.27.1", "@esbuild/darwin-x64": "0.27.1", "@esbuild/freebsd-arm64": "0.27.1", "@esbuild/freebsd-x64": "0.27.1", "@esbuild/linux-arm": "0.27.1", "@esbuild/linux-arm64": "0.27.1", "@esbuild/linux-ia32": "0.27.1", "@esbuild/linux-loong64": "0.27.1", "@esbuild/linux-mips64el": "0.27.1", "@esbuild/linux-ppc64": "0.27.1", "@esbuild/linux-riscv64": "0.27.1", "@esbuild/linux-s390x": "0.27.1", "@esbuild/linux-x64": "0.27.1", "@esbuild/netbsd-arm64": "0.27.1", "@esbuild/netbsd-x64": "0.27.1", "@esbuild/openbsd-arm64": "0.27.1", "@esbuild/openbsd-x64": "0.27.1", "@esbuild/openharmony-arm64": "0.27.1", "@esbuild/sunos-x64": "0.27.1", "@esbuild/win32-arm64": "0.27.1", "@esbuild/win32-ia32": "0.27.1", "@esbuild/win32-x64": "0.27.1" }, "bin": { "esbuild": "bin/esbuild" } }, "sha512-yY35KZckJJuVVPXpvjgxiCuVEJT67F6zDeVTv4rizyPrfGBUpZQsvmxnN+C371c2esD/hNMjj4tpBhuueLN7aA=="],
+
+ "eventsource-parser": ["eventsource-parser@3.0.6", "", {}, "sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg=="],
+
+ "exact-mirror": ["exact-mirror@0.2.5", "", { "peerDependencies": { "@sinclair/typebox": "^0.34.15" }, "optionalPeers": ["@sinclair/typebox"] }, "sha512-u8Wu2lO8nio5lKSJubOydsdNtQmH8ENba5m0nbQYmTvsjksXKYIS1nSShdDlO8Uem+kbo+N6eD5I03cpZ+QsRQ=="],
+
+ "fast-copy": ["fast-copy@4.0.3", "", {}, "sha512-58apWr0GUiDFM8+3afrO6eYwJBn9ZAhDOzG3L+/9llab/haCARS2UIfffmOurYLwbgDRs8n0rfr6qAAPEAuAQw=="],
+
+ "fast-decode-uri-component": ["fast-decode-uri-component@1.0.1", "", {}, "sha512-WKgKWg5eUxvRZGwW8FvfbaH7AXSh2cL+3j5fMGzUMCxWBJ3dV3a7Wz8y2f/uQ0e3B6WmodD3oS54jTQ9HVTIIg=="],
+
+ "fast-deep-equal": ["fast-deep-equal@3.1.3", "", {}, "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q=="],
+
+ "fast-safe-stringify": ["fast-safe-stringify@2.1.1", "", {}, "sha512-W+KJc2dmILlPplD/H4K9l9LcAHAfPtP6BY84uVLXQ6Evcz9Lcg33Y2z1IVblT6xdY54PXYVHEv+0Wpq8Io6zkA=="],
+
+ "fast-uri": ["fast-uri@3.1.0", "", {}, "sha512-iPeeDKJSWf4IEOasVVrknXpaBV0IApz/gp7S2bb7Z4Lljbl2MGJRqInZiUrQwV16cpzw/D3S5j5Julj/gT52AA=="],
+
+ "file-type": ["file-type@21.1.1", "", { "dependencies": { "@tokenizer/inflate": "^0.4.1", "strtok3": "^10.3.4", "token-types": "^6.1.1", "uint8array-extras": "^1.4.0" } }, "sha512-ifJXo8zUqbQ/bLbl9sFoqHNTNWbnPY1COImFfM6CCy7z+E+jC1eY9YfOKkx0fckIg+VljAy2/87T61fp0+eEkg=="],
+
+ "fsevents": ["fsevents@2.3.3", "", { "os": "darwin" }, "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw=="],
+
+ "get-tsconfig": ["get-tsconfig@4.13.0", "", { "dependencies": { "resolve-pkg-maps": "^1.0.0" } }, "sha512-1VKTZJCwBrvbd+Wn3AOgQP/2Av+TfTCOlE4AcRJE72W1ksZXbAx8PPBR9RzgTeSPzlPMHrbANMH3LbltH73wxQ=="],
+
+ "graphlib": ["graphlib@2.1.8", "", { "dependencies": { "lodash": "^4.17.15" } }, "sha512-jcLLfkpoVGmH7/InMC/1hIvOPSUh38oJtGhvrOFGzioE1DZ+0YW16RgmOJhHiuWTvGiJQ9Z1Ik43JvkRPRvE+A=="],
+
+ "help-me": ["help-me@5.0.0", "", {}, "sha512-7xgomUX6ADmcYzFik0HzAxh/73YlKR9bmFzf51CZwR+b6YtzU2m0u49hQCqV6SvlqIqsaxovfwdvbnsw3b/zpg=="],
+
+ "ieee754": ["ieee754@1.2.1", "", {}, "sha512-dcyqhDvX1C46lXZcVqCpK+FtMRQVdIMN6/Df5js2zouUsqG7I6sFxitIC+7KYK29KdXOLHdu9zL4sFnoVQnqaA=="],
+
+ "joycon": ["joycon@3.1.1", "", {}, "sha512-34wB/Y7MW7bzjKRjUKTa46I2Z7eV62Rkhva+KkopW7Qvv/OSWBqvkSY7vusOPrNuZcUG3tApvdVgNB8POj3SPw=="],
+
+ "json-schema": ["json-schema@0.4.0", "", {}, "sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA=="],
+
+ "json-schema-traverse": ["json-schema-traverse@1.0.0", "", {}, "sha512-NM8/P9n3XjXhIZn1lLhkFaACTOURQXjWhV4BA/RnOv8xvgqtqpAX9IO4mRQxSx1Rlo4tqzeqb0sOlruaOy3dug=="],
+
+ "lodash": ["lodash@4.18.1", "", {}, "sha512-dMInicTPVE8d1e5otfwmmjlxkZoUpiVLwyeTdUsi/Caj/gfzzblBcCE5sRHV/AsjuCmxWrte2TNGSYuCeCq+0Q=="],
+
+ "memoirist": ["memoirist@0.4.0", "", {}, "sha512-zxTgA0mSYELa66DimuNQDvyLq36AwDlTuVRbnQtB+VuTcKWm5Qc4z3WkSpgsFWHNhexqkIooqpv4hdcqrX5Nmg=="],
+
+ "minimist": ["minimist@1.2.8", "", {}, "sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA=="],
+
+ "ms": ["ms@2.1.3", "", {}, "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA=="],
+
+ "on-exit-leak-free": ["on-exit-leak-free@2.1.2", "", {}, "sha512-0eJJY6hXLGf1udHwfNftBqH+g73EU4B504nZeKpz1sYRKafAghwxEJunB2O7rDZkL4PGfsMVnTXZ2EjibbqcsA=="],
+
+ "once": ["once@1.4.0", "", { "dependencies": { "wrappy": "1" } }, "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w=="],
+
+ "openapi-types": ["openapi-types@12.1.3", "", {}, "sha512-N4YtSYJqghVu4iek2ZUvcN/0aqH1kRDuNqzcycDxhOUpg7GdvLa2F3DgS6yBNhInhv2r/6I0Flkn7CqL8+nIcw=="],
+
+ "pino": ["pino@10.3.1", "", { "dependencies": { "@pinojs/redact": "^0.4.0", "atomic-sleep": "^1.0.0", "on-exit-leak-free": "^2.1.0", "pino-abstract-transport": "^3.0.0", "pino-std-serializers": "^7.0.0", "process-warning": "^5.0.0", "quick-format-unescaped": "^4.0.3", "real-require": "^0.2.0", "safe-stable-stringify": "^2.3.1", "sonic-boom": "^4.0.1", "thread-stream": "^4.0.0" }, "bin": { "pino": "bin.js" } }, "sha512-r34yH/GlQpKZbU1BvFFqOjhISRo1MNx1tWYsYvmj6KIRHSPMT2+yHOEb1SG6NMvRoHRF0a07kCOox/9yakl1vg=="],
+
+ "pino-abstract-transport": ["pino-abstract-transport@3.0.0", "", { "dependencies": { "split2": "^4.0.0" } }, "sha512-wlfUczU+n7Hy/Ha5j9a/gZNy7We5+cXp8YL+X+PG8S0KXxw7n/JXA3c46Y0zQznIJ83URJiwy7Lh56WLokNuxg=="],
+
+ "pino-pretty": ["pino-pretty@13.1.3", "", { "dependencies": { "colorette": "^2.0.7", "dateformat": "^4.6.3", "fast-copy": "^4.0.0", "fast-safe-stringify": "^2.1.1", "help-me": "^5.0.0", "joycon": "^3.1.1", "minimist": "^1.2.6", "on-exit-leak-free": "^2.1.0", "pino-abstract-transport": "^3.0.0", "pump": "^3.0.0", "secure-json-parse": "^4.0.0", "sonic-boom": "^4.0.1", "strip-json-comments": "^5.0.2" }, "bin": { "pino-pretty": "bin.js" } }, "sha512-ttXRkkOz6WWC95KeY9+xxWL6AtImwbyMHrL1mSwqwW9u+vLp/WIElvHvCSDg0xO/Dzrggz1zv3rN5ovTRVowKg=="],
+
+ "pino-std-serializers": ["pino-std-serializers@7.1.0", "", {}, "sha512-BndPH67/JxGExRgiX1dX0w1FvZck5Wa4aal9198SrRhZjH3GxKQUKIBnYJTdj2HDN3UQAS06HlfcSbQj2OHmaw=="],
+
+ "prettier": ["prettier@3.4.2", "", { "bin": { "prettier": "bin/prettier.cjs" } }, "sha512-e9MewbtFo+Fevyuxn/4rrcDAaq0IYxPGLvObpQjiZBMAzB9IGmzlnG9RZy3FFas+eBMu2vA0CszMeduow5dIuQ=="],
+
+ "process-warning": ["process-warning@5.0.0", "", {}, "sha512-a39t9ApHNx2L4+HBnQKqxxHNs1r7KF+Intd8Q/g1bUh6q0WIp9voPXJ/x0j+ZL45KF1pJd9+q2jLIRMfvEshkA=="],
+
+ "pump": ["pump@3.0.4", "", { "dependencies": { "end-of-stream": "^1.1.0", "once": "^1.3.1" } }, "sha512-VS7sjc6KR7e1ukRFhQSY5LM2uBWAUPiOPa/A3mkKmiMwSmRFUITt0xuj+/lesgnCv+dPIEYlkzrcyXgquIHMcA=="],
+
+ "quick-format-unescaped": ["quick-format-unescaped@4.0.4", "", {}, "sha512-tYC1Q1hgyRuHgloV/YXs2w15unPVh8qfu/qCTfhTYamaw7fyhumKa2yGpdSo87vY32rIclj+4fWYQXUMs9EHvg=="],
+
+ "real-require": ["real-require@0.2.0", "", {}, "sha512-57frrGM/OCTLqLOAh0mhVA9VBMHd+9U7Zb2THMGdBUoZVOtGbJzjxsYGDJ3A9AYYCP4hn6y1TVbaOfzWtm5GFg=="],
+
+ "require-from-string": ["require-from-string@2.0.2", "", {}, "sha512-Xf0nWe6RseziFMu+Ap9biiUbmplq6S9/p+7w7YXP/JBHhrUDDUhwa+vANyubuqfZWTveU//DYVGsDG7RKL/vEw=="],
+
+ "resolve-pkg-maps": ["resolve-pkg-maps@1.0.0", "", {}, "sha512-seS2Tj26TBVOC2NIc2rOe2y2ZO7efxITtLZcGSOnHHNOQ7CkiUBfw0Iw2ck6xkIhPwLhKNLS8BO+hEpngQlqzw=="],
+
+ "rxjs": ["rxjs@7.8.2", "", { "dependencies": { "tslib": "^2.1.0" } }, "sha512-dhKf903U/PQZY6boNNtAGdWbG85WAbjT/1xYoZIC7FAY0yWapOBQVsVrDl58W86//e1VpMNBtRV4MaXfdMySFA=="],
+
+ "safe-stable-stringify": ["safe-stable-stringify@2.5.0", "", {}, "sha512-b3rppTKm9T+PsVCBEOUR46GWI7fdOs00VKZ1+9c1EWDaDMvjQc6tUwuFyIprgGgTcWoVHSKrU8H31ZHA2e0RHA=="],
+
+ "secure-json-parse": ["secure-json-parse@4.1.0", "", {}, "sha512-l4KnYfEyqYJxDwlNVyRfO2E4NTHfMKAWdUuA8J0yve2Dz/E/PdBepY03RvyJpssIpRFwJoCD55wA+mEDs6ByWA=="],
+
+ "sonic-boom": ["sonic-boom@4.2.1", "", { "dependencies": { "atomic-sleep": "^1.0.0" } }, "sha512-w6AxtubXa2wTXAUsZMMWERrsIRAdrK0Sc+FUytWvYAhBJLyuI4llrMIC1DtlNSdI99EI86KZum2MMq3EAZlF9Q=="],
+
+ "split2": ["split2@4.2.0", "", {}, "sha512-UcjcJOWknrNkF6PLX83qcHM6KHgVKNkV62Y8a5uYDVv9ydGQVwAHMKqHdJje1VTWpljG0WYpCDhrCdAOYH4TWg=="],
+
+ "strip-json-comments": ["strip-json-comments@5.0.3", "", {}, "sha512-1tB5mhVo7U+ETBKNf92xT4hrQa3pm0MZ0PQvuDnWgAAGHDsfp4lPSpiS6psrSiet87wyGPh9ft6wmhOMQ0hDiw=="],
+
+ "strtok3": ["strtok3@10.3.4", "", { "dependencies": { "@tokenizer/token": "^0.3.0" } }, "sha512-KIy5nylvC5le1OdaaoCJ07L+8iQzJHGH6pWDuzS+d07Cu7n1MZ2x26P8ZKIWfbK02+XIL8Mp4RkWeqdUCrDMfg=="],
+
+ "thread-stream": ["thread-stream@4.0.0", "", { "dependencies": { "real-require": "^0.2.0" } }, "sha512-4iMVL6HAINXWf1ZKZjIPcz5wYaOdPhtO8ATvZ+Xqp3BTdaqtAwQkNmKORqcIo5YkQqGXq5cwfswDwMqqQNrpJA=="],
+
+ "token-types": ["token-types@6.1.1", "", { "dependencies": { "@borewit/text-codec": "^0.1.0", "@tokenizer/token": "^0.3.0", "ieee754": "^1.2.1" } }, "sha512-kh9LVIWH5CnL63Ipf0jhlBIy0UsrMj/NJDfpsy1SqOXlLKEVyXXYrnFxFT1yOOYVGBSApeVnjPw/sBz5BfEjAQ=="],
+
+ "tslib": ["tslib@2.8.1", "", {}, "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w=="],
+
+ "tsx": ["tsx@4.21.0", "", { "dependencies": { "esbuild": "~0.27.0", "get-tsconfig": "^4.7.5" }, "optionalDependencies": { "fsevents": "~2.3.3" }, "bin": { "tsx": "dist/cli.mjs" } }, "sha512-5C1sg4USs1lfG0GFb2RLXsdpXqBSEhAaA/0kPL01wxzpMqLILNxIxIOKiILz+cdg/pLnOUxFYOR5yhHU666wbw=="],
+
+ "typescript": ["typescript@5.9.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw=="],
+
+ "uint8array-extras": ["uint8array-extras@1.5.0", "", {}, "sha512-rvKSBiC5zqCCiDZ9kAOszZcDvdAHwwIKJG33Ykj43OKcWsnmcBRL09YTU4nOeHZ8Y2a7l1MgTd08SBe9A8Qj6A=="],
+
+ "undici-types": ["undici-types@7.16.0", "", {}, "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw=="],
+
+ "wrappy": ["wrappy@1.0.2", "", {}, "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ=="],
+
+ "zod": ["zod@3.25.76", "", {}, "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ=="],
+ }
+}
diff --git a/agent-service/package.json b/agent-service/package.json
new file mode 100644
index 00000000000..608981e0a45
--- /dev/null
+++ b/agent-service/package.json
@@ -0,0 +1,37 @@
+{
+ "name": "texera-agent-service",
+ "version": "0.1.0",
+ "description": "Texera Agent Service - AI agents for workflow manipulation",
+ "type": "module",
+ "main": "src/index.ts",
+ "scripts": {
+ "dev": "bun run --watch src/server.ts",
+ "dev:node": "npx tsx --watch src/server.ts",
+ "start": "bun run src/server.ts",
+ "start:node": "npx tsx src/server.ts",
+ "test": "bun test",
+ "typecheck": "tsc --noEmit",
+ "format": "prettier --write \"src/**/*.{ts,tsx,json}\"",
+ "format:check": "prettier --check \"src/**/*.{ts,tsx,json}\""
+ },
+ "dependencies": {
+ "@ai-sdk/openai": "2.0.79",
+ "@elysiajs/cors": "1.4.0",
+ "ai": "5.0.108",
+ "ajv": "8.10.0",
+ "dagre": "0.8.5",
+ "elysia": "1.4.18",
+ "pino": "10.3.1",
+ "rxjs": "7.8.2",
+ "zod": "3.25.76"
+ },
+ "devDependencies": {
+ "@types/bun": "1.3.3",
+ "@types/dagre": "0.7.54",
+ "pino-pretty": "13.1.3",
+ "prettier": "3.4.2",
+ "tsx": "4.21.0",
+ "typescript": "5.9.3"
+ },
+ "packageManager": "yarn@4.5.1+sha512.341db9396b6e289fecc30cd7ab3af65060e05ebff4b3b47547b278b9e67b08f485ecd8c79006b405446262142c7a38154445ef7f17c1d5d1de7d90bf9ce7054d"
+}
diff --git a/agent-service/src/agent/index.ts b/agent-service/src/agent/index.ts
new file mode 100644
index 00000000000..d7b417b6ae7
--- /dev/null
+++ b/agent-service/src/agent/index.ts
@@ -0,0 +1,21 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export * from "./texera-agent";
+export * from "./prompts";
diff --git a/agent-service/src/agent/prompts.ts b/agent-service/src/agent/prompts.ts
new file mode 100644
index 00000000000..064eed2e3e5
--- /dev/null
+++ b/agent-service/src/agent/prompts.ts
@@ -0,0 +1,296 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { WorkflowSystemMetadata } from "./util/workflow-system-metadata";
+
+const PYTHON_UDF_OPERATOR_TYPES = ["PythonUDFV2"];
+const R_UDF_OPERATOR_TYPES = ["RUDF"];
+
+const PYTHON_UDF_INSTRUCTIONS = `## Python UDF Guide
+
+Python UDF operators run user-defined Python code. There are 2 APIs to process data:
+
+### Tuple API
+Takes one input tuple from a port at a time. Returns an iterator of optional TupleLike instances.
+Use cases: Functional operations applied to tuples one by one (map, reduce, filter).
+
+Template:
+\`\`\`python
+from pytexera import *
+
+class ProcessTupleOperator(UDFOperatorV2):
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ yield tuple_
+\`\`\`
+
+Example - Filter tuples by conditions:
+\`\`\`python
+from pytexera import *
+
+class ProcessTupleOperator(UDFOperatorV2):
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ q = tuple_["QUANTITY"]
+ oq = tuple_["ORDERED_QUANTITY"]
+ p = tuple_["UNIT_PRICE"]
+ if q is not None and oq is not None and p is not None:
+ if q <= oq and p >= 0:
+ yield tuple_
+\`\`\`
+
+### Table API
+Consumes a whole Table (pandas DataFrame) from a port. Returns an iterator of optional TableLike instances.
+Use cases: Blocking operations that consume the whole table.
+
+Template:
+\`\`\`python
+from pytexera import *
+
+class ProcessTableOperator(UDFTableOperator):
+ def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:
+ yield table
+\`\`\`
+
+Example - Filter DataFrame rows:
+\`\`\`python
+from pytexera import *
+import pandas as pd
+
+class ProcessTableOperator(UDFTableOperator):
+ def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:
+ df: pd.DataFrame = table
+ m1 = (df["KWMENG"].notna()) & (df["KBMENG"].notna()) & (df["KWMENG"] <= df["KBMENG"])
+ m2 = (df["NET_VALUE"].notna()) & (df["NET_VALUE"] >= 0)
+ yield df[m1 & m2]
+\`\`\`
+
+### Important Rules
+
+- DO NOT change the class name (ProcessTupleOperator or ProcessTableOperator).
+- Import packages explicitly (pandas, numpy, etc.).
+- Tuple is a Python dict. Access fields with tuple_["field"] ONLY (no .get/.set/.values).
+- Table is a pandas DataFrame.
+- Use yield to return results.
+- Handle None values carefully.
+- Do not cast types.
+- Keep each UDF focused on one task.
+- Only change the python code property, not other properties.
+- If adding extra columns, specify them in the Extra Output Columns property.
+- Prefer native operators over Python UDF when possible.`;
+
+const R_UDF_INSTRUCTIONS = `## R UDF Guide
+
+R UDF operators run user-defined R code. Two modes: Table API and Tuple API.
+
+### Table API
+Passes the entire input as an R data frame to your function and expects a data frame in return.
+
+Template:
+\`\`\`r
+function(table, port) {
+ return(table)
+}
+\`\`\`
+
+Example - Keep rows where quantities align and net value is valid:
+\`\`\`r
+function(table, port) {
+ valid_qty <- !is.na(table$KWMENG) & !is.na(table$KBMENG) & table$KWMENG <= table$KBMENG
+ valid_value <- !is.na(table$NET_VALUE) & table$NET_VALUE >= 0
+ valid_rows <- valid_qty & valid_value
+ return(table[valid_rows, , drop = FALSE])
+}
+\`\`\`
+
+### Tuple API
+Uses coro::generator to yield tuples (lists) one by one.
+
+Template:
+\`\`\`r
+library(coro)
+
+coro::generator(function(tuple, port) {
+ yield(tuple)
+})
+\`\`\`
+
+Example - Emit tuples that flag problematic status values:
+\`\`\`r
+library(coro)
+
+coro::generator(function(tuple, port) {
+ status <- tuple$STATUS
+ if (!is.null(status) && status == "ERROR") {
+ yield(tuple)
+ }
+})
+\`\`\`
+
+### Important Rules
+
+- Return a function(table, port) for Table API; use coro::generator(function(tuple, port) { ... }) for Tuple API.
+- Load libraries explicitly with library().
+- Handle NA with is.na() before comparisons.
+- Use yield() inside generators for each tuple to emit.
+- Keep output schema consistent with Retain input columns and Extra output columns settings.
+- Keep scripts focused on one task.
+- Only modify the script code field unless necessary.`;
+
+const SYSTEM_PROMPT_TEMPLATE = `You are a data science Copilot that helps users solve data-centric tasks by building dataflows.
+
+## What is Dataflow?
+
+Dataflow represents data analysis as a DAG (directed acyclic graph) where:
+- Each **operator** is a single step of data processing
+- Each **link** represents data dependency between operators
+- Each operator receives table(s) from input operator(s), processes them, and outputs a single table
+- The output table can be viewed via execution, or passed to downstream operators via links
+
+## Context Format
+
+Your conversation context is a single message with three top-level sections, in this order:
+
+- \`# Completed Tasks\` — previous tasks you've already finished (omitted if none)
+- \`# Ongoing Task\` — the current task, including turns you've taken so far
+- \`# Current Dataflow\` — the live DAG: every operator's current state
+
+**Overall layout:**
+
+\`\`\`
+# Completed Tasks
+
+## Task (completed)
+
+### User request
+
+
+
+### Turn 1
+Thought:
+- (succeeded)
+ - Summary:
+ - Output:
+
+## Task (completed)
+
+### User request
+
+
+
+### Turn 1
+...
+
+# Ongoing Task
+## Task (ongoing)
+
+### User request
+
+
+
+### Turn 1
+Thought: ...
+- (succeeded)
+ - Summary: ...
+ - Output: ...
+
+### Turn 2
+Thought: ...
+- (failed)
+ - Summary: ...
+ - Error:
+
+
+# Current Dataflow
+## Operators
+
+### Operator \`\` (, executed|failed|not-executed)
+Summary:
+Input Schema (port 0): [: , ...]
+Properties:
+ :
+Output Schema: [: , ...]
+Compilation Error:
+Result:
+
+
+### Operator \`\` ...
+...
+
+## Links
+- →
+\`\`\`
+
+## Key Principles
+
+- **Call tools only through the native protocol**: Invoke tools using the tool-call mechanism. Never emit \`\`, \`\`, \`\`, or any other tag-like structures in your response — those shapes appear in your input to describe past turns and existing state, never in your output.
+- **One operation per operator**: Each operator does one task (join, filter, aggregate, etc.). Use links to connect them.
+- **Build incrementally**: Link new operators to existing ones. Never recreate data already in the workflow.
+- **Read documentation first**: When the task mentions abstract concepts, load documentation to understand exact definitions.
+- **Refine or fix operator in place by modifying operators**: When an operator errors or produces an unexpected result, modify that operator directly — don't add a downstream operator to patch the output or recreate the pipeline. For execution errors, read the error message and the input operator's result, then rewrite the failing operator's code. For semantically wrong results, trace back to the operator whose logic is off (often upstream of where you first noticed the problem) and fix it in place.
+- **Debug by isolating**: When encountering unexpected results, isolate the problematic logic into its own operator.
+- **Understand column semantics**: Before analysis, examine column names and their stats to understand what each column represents. Columns may carry semantic meaning that affects how data should be filtered or interpreted — respect these signals and apply appropriate preprocessing before computing results.
+- **Normalize before grouping or joining**: String keys may contain naming variants such as special character delimiters, encoding differences, or duplicate entries across files. Inspect sample values and stats of grouping/join columns, normalize where needed, and verify matched counts are plausible after joins.
+- **Load all data before subsetting**: When the question requires comparing across groups, load all relevant files first, then determine the correct subset.
+- **Handle messy data files**: Load data files directly in a single operator. Real-world data files are often malformed — they may have wrong delimiters, missing or misplaced headers, metadata/comment rows, or multiple tables in one file. After loading, inspect the result. If column names look auto-generated (e.g., \`Unnamed: 0\`) or a data value appears as a header, adjust the loading parameters (e.g., \`header=\`, \`skiprows=\`, \`sep=\`) by modifying the data loading operator.
+- **Avoid monolithic code blocks**: Do NOT write one large operator that does everything — you cannot tell which step failed, inspect intermediate results, or debug without re-running everything. Instead, decompose into separate operators each doing ONE thing (e.g., filter → join → aggregate → filter → join → final filter). Each can be executed and verified independently.
+
+## Available Operators
+
+You have the following operators available:
+
+{{OPERATOR_SCHEMA}}
+`;
+
+function buildAllowedOperatorSchemas(
+ metadataStore: WorkflowSystemMetadata,
+ allowedOperatorTypes: string[] = []
+): string {
+ const schemas: string[] = [];
+
+ const operatorTypes =
+ allowedOperatorTypes.length > 0 ? allowedOperatorTypes : Object.keys(metadataStore.getAllOperatorTypes());
+
+ for (const operatorType of operatorTypes) {
+ const compactSchema = metadataStore.getCompactSchema(operatorType);
+ const description = metadataStore.getDescription(operatorType);
+
+ if (compactSchema) {
+ schemas.push(
+ `## ${operatorType}\n` +
+ (description ? `Description: ${description}\n` : "") +
+ `Schema:\n\`\`\`json\n${JSON.stringify(compactSchema, null, 2)}\n\`\`\``
+ );
+ }
+ }
+
+ return schemas.length > 0 ? schemas.join("\n\n") : "No operators available.";
+}
+
+export function buildSystemPrompt(metadataStore: WorkflowSystemMetadata, allowedOperatorTypes: string[] = []): string {
+ const operatorSchemas = buildAllowedOperatorSchemas(metadataStore, allowedOperatorTypes);
+ const allowsAll = allowedOperatorTypes.length === 0;
+ const pythonAllowed = allowsAll || allowedOperatorTypes.some(t => PYTHON_UDF_OPERATOR_TYPES.includes(t));
+ const rAllowed = allowsAll || allowedOperatorTypes.some(t => R_UDF_OPERATOR_TYPES.includes(t));
+
+ const extraSections: string[] = [];
+ if (pythonAllowed) extraSections.push(PYTHON_UDF_INSTRUCTIONS);
+ if (rAllowed) extraSections.push(R_UDF_INSTRUCTIONS);
+
+ const base = SYSTEM_PROMPT_TEMPLATE.replace("{{OPERATOR_SCHEMA}}", operatorSchemas);
+ return extraSections.length > 0 ? `${base}\n${extraSections.join("\n\n")}\n` : base;
+}
diff --git a/agent-service/src/agent/texera-agent.ts b/agent-service/src/agent/texera-agent.ts
new file mode 100644
index 00000000000..37eb12d8688
--- /dev/null
+++ b/agent-service/src/agent/texera-agent.ts
@@ -0,0 +1,840 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { generateText, type ModelMessage, type LanguageModel, stepCountIs } from "ai";
+import { Subscription } from "rxjs";
+import { debounceTime } from "rxjs/operators";
+import { WorkflowState } from "./workflow-state";
+import { WorkflowSystemMetadata } from "./util/workflow-system-metadata";
+import { WorkflowResultState } from "./workflow-result-state";
+import { formatOperatorResult } from "./tools/result-formatting";
+import type { AgentSettings, ReActStep, TokenUsage, UserInfo } from "../types/agent";
+import {
+ AgentState as AgentStateEnum,
+ DEFAULT_AGENT_SETTINGS,
+ OperatorResultSerializationMode,
+ INITIAL_STEP_ID,
+} from "../types/agent";
+import { buildSystemPrompt } from "./prompts";
+import {
+ createAddOperatorTool,
+ createModifyOperatorTool,
+ createDeleteOperatorTool,
+ TOOL_NAME_ADD_OPERATOR,
+ TOOL_NAME_MODIFY_OPERATOR,
+ TOOL_NAME_DELETE_OPERATOR,
+ type ToolContext,
+} from "./tools/workflow-crud-tools";
+import {
+ createExecuteOperatorTool,
+ executeOperatorAndFormat,
+ TOOL_NAME_EXECUTE_OPERATOR,
+ type ExecutionConfig,
+} from "./tools/workflow-execution-tools";
+import { assembleContext } from "./util/context-utils";
+import { compileWorkflowAsync, type WorkflowCompilationResponse } from "../api/compile-api";
+import { createLogger } from "../logger";
+import type { Logger } from "pino";
+
+const PERSIST_DEBOUNCE_MS = 500;
+
+export interface TexeraAgentConfig {
+ model: LanguageModel;
+ modelType: string;
+ agentId: string;
+ agentName?: string;
+ systemPrompt?: string;
+}
+
+export interface AgentMessageResult {
+ response: string;
+ messages: ModelMessage[];
+ usage: TokenUsage;
+ stopped: boolean;
+ error?: string;
+}
+
+type ReActStepCallback = (step: ReActStep) => void;
+
+/**
+ * A single Texera agent instance.
+ *
+ * Owns the conversation (ReAct step tree with HEAD/checkout semantics), the
+ * workflow being edited (`WorkflowState`), cached operator execution results
+ * (`WorkflowResultState`), and the tool surface exposed to the LLM. Each call
+ * to `sendMessage` drives one multi-step generation via the Vercel AI SDK,
+ * streaming step updates to subscribed websockets.
+ */
+export class TexeraAgent {
+ readonly agentId: string;
+ readonly agentName: string;
+ readonly modelType: string;
+ readonly createdAt: Date;
+
+ private state: AgentStateEnum = AgentStateEnum.AVAILABLE;
+ private workflowState: WorkflowState;
+ private metadataStore: WorkflowSystemMetadata;
+ private head: string = INITIAL_STEP_ID;
+ private stepsById: Map = new Map();
+ private stepCounter = 0;
+ private workflowResultState: WorkflowResultState;
+
+ private websockets: Set = new Set();
+
+ private model: LanguageModel;
+ private systemPrompt: string;
+ private settings: AgentSettings;
+
+ private reActStepsByMessageId: Map = new Map();
+
+ private currentMessageId: string | undefined = undefined;
+
+ private delegateConfig?: {
+ userToken: string;
+ userInfo?: UserInfo;
+ workflowId: number;
+ workflowName?: string;
+ computingUnitId?: number;
+ };
+
+ private stepCallback: ReActStepCallback | null = null;
+
+ private messageCounter = 0;
+
+ private tools: Record;
+
+ private abortController: AbortController | null = null;
+
+ private workflowChangeSubscription: Subscription | null = null;
+
+ private log: Logger;
+
+ constructor(config: TexeraAgentConfig) {
+ this.agentId = config.agentId;
+ this.agentName = config.agentName || `Agent-${config.agentId}`;
+ this.modelType = config.modelType;
+ this.createdAt = new Date();
+ this.model = config.model;
+ this.systemPrompt = config.systemPrompt || "";
+ this.log = createLogger("TexeraAgent", { agentId: this.agentId });
+
+ this.workflowState = new WorkflowState();
+ this.metadataStore = WorkflowSystemMetadata.getInstance();
+ this.workflowResultState = new WorkflowResultState(() => this.getAncestorPath());
+
+ const initialStep: ReActStep = {
+ id: INITIAL_STEP_ID,
+ messageId: "initial",
+ stepId: -1,
+ timestamp: Date.now(),
+ role: "user",
+ content: "",
+ isBegin: true,
+ isEnd: true,
+ parentId: undefined,
+ };
+ this.stepsById.set(INITIAL_STEP_ID, initialStep);
+
+ this.settings = {
+ ...DEFAULT_AGENT_SETTINGS,
+ systemPrompt: this.systemPrompt,
+ };
+
+ this.tools = this.createTools();
+ }
+
+ async initialize(): Promise {
+ try {
+ if (!this.metadataStore.isInitialized()) {
+ await this.metadataStore.initializeFromBackend();
+ }
+
+ this.rebuildSystemPrompt();
+
+ this.tools = this.createTools();
+ this.log.info({ operatorCount: this.metadataStore.getOperatorCount() }, "agent initialized");
+ } catch (error) {
+ this.log.error({ err: error }, "failed to initialize metadata");
+ }
+ }
+
+ private rebuildSystemPrompt(): void {
+ this.systemPrompt = buildSystemPrompt(this.metadataStore, this.settings.allowedOperatorTypes);
+ this.settings.systemPrompt = this.systemPrompt;
+ }
+
+ private buildExecutionConfig(): ExecutionConfig | undefined {
+ if (!this.delegateConfig) return undefined;
+ return {
+ userToken: this.delegateConfig.userToken,
+ workflowId: this.delegateConfig.workflowId,
+ computingUnitId: this.delegateConfig.computingUnitId,
+ maxOperatorResultCharLimit: this.settings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: this.settings.maxOperatorResultCellCharLimit,
+ executionTimeoutMs: this.settings.executionTimeoutMs,
+ };
+ }
+
+ private createTools(): Record {
+ const operatorSchemas = new Map();
+ for (const type of Object.keys(this.metadataStore.getAllOperatorTypes())) {
+ const jsonSchema = this.metadataStore.getSchema(type);
+ const additionalMetadata = this.metadataStore.getAdditionalMetadata(type);
+ if (jsonSchema) {
+ operatorSchemas.set(type, { jsonSchema, additionalMetadata });
+ }
+ }
+
+ const getExecutionConfig = this.delegateConfig ? () => this.buildExecutionConfig()! : undefined;
+
+ const context: ToolContext = {
+ metadataStore: this.metadataStore,
+ settings: {
+ maxOperatorResultCharLimit: this.settings.maxOperatorResultCharLimit,
+ toolTimeoutMs: this.settings.toolTimeoutMs,
+ executionTimeoutMs: this.settings.executionTimeoutMs,
+ },
+ };
+
+ const tools: Record = {
+ [TOOL_NAME_DELETE_OPERATOR]: createDeleteOperatorTool(this.workflowState, context),
+ [TOOL_NAME_ADD_OPERATOR]: createAddOperatorTool(this.workflowState, operatorSchemas, context),
+ [TOOL_NAME_MODIFY_OPERATOR]: createModifyOperatorTool(this.workflowState, context),
+ };
+
+ if (getExecutionConfig) {
+ tools[TOOL_NAME_EXECUTE_OPERATOR] = createExecuteOperatorTool(
+ this.workflowState,
+ getExecutionConfig,
+ (opId, operatorInfo) => {
+ this.workflowResultState.set(opId, this.head, operatorInfo);
+ }
+ );
+ }
+
+ return tools;
+ }
+
+ getState(): AgentStateEnum {
+ return this.state;
+ }
+
+ getWorkflowState(): WorkflowState {
+ return this.workflowState;
+ }
+
+ getMetadataStore(): WorkflowSystemMetadata {
+ return this.metadataStore;
+ }
+
+ getHead(): string {
+ return this.head;
+ }
+
+ getAncestorPath(stepId?: string): string[] {
+ const target = stepId ?? this.head;
+ const chain: string[] = [];
+ let current: string | undefined = target;
+ while (current) {
+ chain.unshift(current);
+ current = this.stepsById.get(current)?.parentId;
+ }
+ return chain;
+ }
+
+ getStepsById(): Map {
+ return this.stepsById;
+ }
+
+ getWorkflowResultState(): WorkflowResultState {
+ return this.workflowResultState;
+ }
+
+ getWebsockets(): Set {
+ return this.websockets;
+ }
+
+ addWebsocket(ws: any): void {
+ this.websockets.add(ws);
+ }
+
+ removeWebsocket(ws: any): void {
+ this.websockets.delete(ws);
+ }
+
+ getReActSteps(): ReActStep[] {
+ const all: ReActStep[] = [];
+ for (const steps of this.reActStepsByMessageId.values()) {
+ all.push(...steps);
+ }
+ return all;
+ }
+
+ getVisibleReActSteps(): ReActStep[] {
+ const path = this.getAncestorPath();
+ return path
+ .filter(id => id !== INITIAL_STEP_ID)
+ .map(id => this.stepsById.get(id)!)
+ .filter(Boolean);
+ }
+
+ getAllSteps(): ReActStep[] {
+ return Array.from(this.stepsById.values()).filter(s => s.id !== INITIAL_STEP_ID);
+ }
+
+ checkout(stepId: string): boolean {
+ const step = this.stepsById.get(stepId);
+ if (!step && stepId !== INITIAL_STEP_ID) return false;
+ this.head = stepId;
+ if (step?.afterWorkflowContent) {
+ this.workflowState.setWorkflowContent(step.afterWorkflowContent);
+ }
+ return true;
+ }
+
+ setStepCallback(callback: ReActStepCallback | null): void {
+ this.stepCallback = callback;
+ }
+
+ private generateStepId(): string {
+ return `step-${this.agentId}-${++this.stepCounter}-${Date.now()}`;
+ }
+
+ private addStep(step: ReActStep): void {
+ let steps = this.reActStepsByMessageId.get(step.messageId);
+ if (!steps) {
+ steps = [];
+ this.reActStepsByMessageId.set(step.messageId, steps);
+ }
+ steps.push(step);
+ this.stepsById.set(step.id, step);
+ if (this.stepCallback) {
+ this.stepCallback(step);
+ }
+ }
+
+ getSystemInfo(): {
+ systemPrompt: string;
+ tools: Array<{ name: string; description: string; inputSchema: any; enabled: boolean }>;
+ } {
+ const toolsInfo = Object.entries(this.tools).map(([name, toolDef]) => {
+ const description = toolDef.description || "";
+ const inputSchema = toolDef.parameters || {};
+ const enabled = !this.settings.disabledTools.has(name);
+
+ return {
+ name,
+ description,
+ inputSchema,
+ enabled,
+ };
+ });
+
+ return {
+ systemPrompt: this.systemPrompt,
+ tools: toolsInfo,
+ };
+ }
+
+ getSettings(): AgentSettings {
+ return { ...this.settings };
+ }
+
+ updateSettings(updates: {
+ maxOperatorResultCharLimit?: number;
+ maxOperatorResultCellCharLimit?: number;
+ operatorResultSerializationMode?: OperatorResultSerializationMode;
+ toolTimeoutMs?: number;
+ executionTimeoutMs?: number;
+ disabledTools?: Set;
+ maxSteps?: number;
+ allowedOperatorTypes?: string[];
+ }): void {
+ let promptNeedsRebuild = false;
+
+ if (updates.maxOperatorResultCharLimit !== undefined) {
+ this.settings.maxOperatorResultCharLimit = updates.maxOperatorResultCharLimit;
+ }
+ if (updates.maxOperatorResultCellCharLimit !== undefined) {
+ this.settings.maxOperatorResultCellCharLimit = updates.maxOperatorResultCellCharLimit;
+ }
+ if (updates.operatorResultSerializationMode !== undefined) {
+ this.settings.operatorResultSerializationMode = updates.operatorResultSerializationMode;
+ }
+ if (updates.toolTimeoutMs !== undefined) {
+ this.settings.toolTimeoutMs = updates.toolTimeoutMs;
+ }
+ if (updates.executionTimeoutMs !== undefined) {
+ this.settings.executionTimeoutMs = updates.executionTimeoutMs;
+ }
+ if (updates.disabledTools !== undefined) {
+ this.settings.disabledTools = updates.disabledTools;
+ }
+ if (updates.maxSteps !== undefined) {
+ this.settings.maxSteps = updates.maxSteps;
+ }
+ if (updates.allowedOperatorTypes !== undefined) {
+ this.settings.allowedOperatorTypes = updates.allowedOperatorTypes;
+ promptNeedsRebuild = true;
+ }
+
+ if (promptNeedsRebuild) {
+ this.rebuildSystemPrompt();
+ }
+
+ this.tools = this.createTools();
+ this.log.info(
+ {
+ maxOperatorResultCharLimit: this.settings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: this.settings.maxOperatorResultCellCharLimit,
+ },
+ "settings updated"
+ );
+ }
+
+ async refreshWorkflowFromBackend(): Promise {
+ // HEAD at a real step means the workflow is determined by that step's snapshot;
+ // only reload from backend when HEAD is the initial sentinel.
+ if (this.head !== INITIAL_STEP_ID) {
+ return;
+ }
+
+ if (!this.delegateConfig?.workflowId || !this.delegateConfig?.userToken) {
+ return;
+ }
+
+ try {
+ const { retrieveWorkflow } = await import("../api/workflow-api");
+ const workflow = await retrieveWorkflow(this.delegateConfig.userToken, this.delegateConfig.workflowId);
+ this.workflowState.setWorkflowContent(workflow.content);
+ this.log.debug({ workflowId: this.delegateConfig.workflowId }, "refreshed workflow from backend");
+ } catch (error) {
+ this.log.warn({ err: error }, "failed to refresh workflow from backend");
+ }
+ }
+
+ setDelegateConfig(config: {
+ userToken: string;
+ userInfo?: UserInfo;
+ workflowId: number;
+ workflowName?: string;
+ computingUnitId?: number;
+ }): void {
+ this.delegateConfig = config;
+
+ this.tools = this.createTools();
+
+ this.setupWorkflowChangeHandlers();
+ }
+
+ getDelegateConfig():
+ | { userToken: string; userInfo?: UserInfo; workflowId: number; workflowName?: string; computingUnitId?: number }
+ | undefined {
+ return this.delegateConfig;
+ }
+
+ private setupWorkflowChangeHandlers(): void {
+ if (this.workflowChangeSubscription) {
+ this.workflowChangeSubscription.unsubscribe();
+ }
+
+ const subscription = new Subscription();
+ const workflowChanged$ = this.workflowState.getWorkflowChangedStream();
+
+ if (this.delegateConfig?.workflowId && this.delegateConfig.userToken) {
+ const persistSubscription = workflowChanged$.pipe(debounceTime(PERSIST_DEBOUNCE_MS)).subscribe(async () => {
+ if (!this.delegateConfig?.workflowId || !this.delegateConfig.userToken) {
+ return;
+ }
+
+ try {
+ const { persistWorkflow } = await import("../api/workflow-api");
+ const workflowContent = this.workflowState.getWorkflowContent();
+ await persistWorkflow(
+ this.delegateConfig.userToken,
+ this.delegateConfig.workflowId,
+ this.delegateConfig.workflowName || "Agent Workflow",
+ workflowContent
+ );
+ this.log.debug({ workflowId: this.delegateConfig.workflowId }, "auto-persisted workflow");
+ } catch (error) {
+ this.log.error({ err: error }, "failed to auto-persist workflow");
+ }
+ });
+
+ subscription.add(persistSubscription);
+ }
+
+ this.workflowChangeSubscription = subscription;
+ this.workflowState.addSubscription(subscription);
+ }
+
+ async sendMessage(userMessage: string, messageSource?: "chat" | "feedback"): Promise {
+ const messageId = `msg-${this.agentId}-${++this.messageCounter}-${Date.now()}`;
+ let stepIndex = 0;
+
+ await this.refreshWorkflowFromBackend();
+
+ this.abortController = new AbortController();
+
+ this.state = AgentStateEnum.GENERATING;
+
+ this.currentMessageId = messageId;
+
+ try {
+ let beforeStepContent = this.workflowState.getWorkflowContent();
+
+ const estimatedInputTokens = Math.ceil(userMessage.length / 4);
+ const userStepId = this.generateStepId();
+ const userStep: ReActStep = {
+ id: userStepId,
+ parentId: this.head,
+ messageId,
+ stepId: 0,
+ timestamp: Date.now(),
+ role: "user",
+ content: userMessage,
+ isBegin: true,
+ isEnd: true,
+ messageSource,
+ beforeWorkflowContent: beforeStepContent,
+ afterWorkflowContent: beforeStepContent,
+ usage: {
+ inputTokens: estimatedInputTokens,
+ outputTokens: 0,
+ totalTokens: estimatedInputTokens,
+ },
+ };
+ this.addStep(userStep);
+ this.head = userStepId;
+
+ let isFirstStep = true;
+ let lastPreparedMessages: ModelMessage[] | undefined;
+
+ // Pass only the current user turn; prepareStep rebuilds full context each step
+ // (historical interactions + DAG + this message).
+ const currentUserMessage: ModelMessage[] = [{ role: "user", content: userMessage }];
+ const result = await generateText({
+ model: this.model,
+ system: this.systemPrompt,
+ messages: currentUserMessage,
+ tools: this.tools,
+ temperature: 0.2,
+ stopWhen: stepCountIs(this.settings.maxSteps),
+ prepareStep: async ({ stepNumber, messages: currentMessages }) => {
+ let compilationResult: WorkflowCompilationResponse | null = null;
+ if (this.workflowState.getAllOperators().length > 0) {
+ try {
+ const logicalPlan = this.workflowState.toLogicalPlan();
+ compilationResult = await compileWorkflowAsync(logicalPlan);
+ } catch (e: any) {
+ this.log.warn({ err: e?.message || e }, "compilation failed; proceeding without schemas");
+ }
+ }
+
+ const visibleSteps = this.getVisibleReActSteps();
+ const processed = assembleContext(
+ visibleSteps,
+ this.workflowState,
+ this.getFormattedResultsForDAG(),
+ false,
+ compilationResult
+ );
+ lastPreparedMessages = processed;
+ return { messages: processed };
+ },
+ abortSignal: this.abortController?.signal,
+ // reasoning_effort is configured per-model in litellm-config.yaml via extra_body
+ // to bypass LiteLLM's param validation — do not pass it here.
+ providerOptions: {
+ openai: { parallelToolCalls: false },
+ anthropic: { disableParallelToolUse: true },
+ mistral: { parallelToolCalls: false },
+ },
+ onStepFinish: async ({ text, toolCalls, toolResults, usage }) => {
+ stepIndex++;
+
+ const formattedToolCalls = toolCalls?.map(tc => ({
+ toolName: tc.toolName,
+ toolCallId: tc.toolCallId,
+ input: tc.input,
+ }));
+
+ const formattedToolResults = toolResults?.map(tr => ({
+ toolCallId: tr.toolCallId,
+ output: tr.output,
+ isError: !!(tr.output as any)?.error,
+ }));
+
+ const afterStepContent = this.workflowState.getWorkflowContent();
+
+ const agentStepId = this.generateStepId();
+ const agentStep: ReActStep = {
+ id: agentStepId,
+ parentId: this.head,
+ messageId,
+ stepId: stepIndex,
+ timestamp: Date.now(),
+ role: "agent",
+ content: text || "",
+ isBegin: isFirstStep,
+ isEnd: false,
+ toolCalls: formattedToolCalls,
+ toolResults: formattedToolResults,
+ usage: usage
+ ? {
+ inputTokens: usage.inputTokens,
+ outputTokens: usage.outputTokens,
+ totalTokens: usage.totalTokens,
+ }
+ : undefined,
+ inputMessages: lastPreparedMessages,
+ beforeWorkflowContent: beforeStepContent,
+ afterWorkflowContent: afterStepContent,
+ };
+ lastPreparedMessages = undefined;
+ this.addStep(agentStep);
+ this.head = agentStepId;
+
+ const execConfig = this.buildExecutionConfig();
+ if (execConfig && toolCalls && toolResults) {
+ const EXECUTE_AFTER_TOOLS = new Set([TOOL_NAME_ADD_OPERATOR, TOOL_NAME_MODIFY_OPERATOR]);
+
+ for (let i = 0; i < toolCalls.length; i++) {
+ const tc = toolCalls[i];
+ const tr = toolResults[i];
+ if (!EXECUTE_AFTER_TOOLS.has(tc.toolName)) continue;
+
+ const resultText = typeof tr?.output === "string" ? tr.output : String(tr?.output ?? "");
+ if (resultText.startsWith("[ERROR]")) continue;
+
+ const operatorId = (tc.input as any)?.operatorId;
+ if (!operatorId) continue;
+
+ try {
+ await executeOperatorAndFormat(this.workflowState, execConfig, operatorId, {
+ abortSignal: this.abortController?.signal,
+ onResult: (opId, operatorInfo) => {
+ this.workflowResultState.set(opId, this.head, operatorInfo);
+ },
+ });
+ } catch (e: any) {
+ this.log.warn({ operatorId, err: e?.message || e }, "post-step execution failed");
+ }
+ }
+ }
+
+ beforeStepContent = afterStepContent;
+ isFirstStep = false;
+ },
+ });
+
+ const msgSteps = this.reActStepsByMessageId.get(messageId);
+ if (msgSteps && msgSteps.length > 0) {
+ const lastStep = msgSteps[msgSteps.length - 1];
+ if (lastStep.role === "agent") {
+ lastStep.isEnd = true;
+ }
+ }
+
+ const finalUsage = (result as any).totalUsage || result.usage;
+ const usage: TokenUsage = {
+ inputTokens: finalUsage?.inputTokens ?? finalUsage?.promptTokens ?? 0,
+ outputTokens: finalUsage?.outputTokens ?? finalUsage?.completionTokens ?? 0,
+ totalTokens: finalUsage?.totalTokens ?? 0,
+ };
+
+ return {
+ response: result.text,
+ messages: result.response.messages,
+ usage,
+ stopped: false,
+ };
+ } catch (error: any) {
+ const isAborted = error.name === "AbortError" || this.abortController?.signal.aborted;
+
+ if (isAborted) {
+ stepIndex++;
+ const stoppedStepId = this.generateStepId();
+ const stoppedStep: ReActStep = {
+ id: stoppedStepId,
+ parentId: this.head,
+ messageId,
+ stepId: stepIndex,
+ timestamp: Date.now(),
+ role: "agent",
+ content: "Generation stopped by user.",
+ isBegin: false,
+ isEnd: true,
+ };
+ this.addStep(stoppedStep);
+ this.head = stoppedStepId;
+
+ return {
+ response: "",
+ messages: [],
+ usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
+ stopped: true,
+ };
+ }
+
+ stepIndex++;
+ const errorStepId = this.generateStepId();
+ const errorStep: ReActStep = {
+ id: errorStepId,
+ parentId: this.head,
+ messageId,
+ stepId: stepIndex,
+ timestamp: Date.now(),
+ role: "agent",
+ content: `Error: ${error.message || String(error)}`,
+ isBegin: false,
+ isEnd: true,
+ };
+ this.addStep(errorStep);
+ this.head = errorStepId;
+
+ return {
+ response: "",
+ messages: [],
+ usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
+ stopped: false,
+ error: error.message || String(error),
+ };
+ } finally {
+ this.abortController = null;
+ this.currentMessageId = undefined;
+ this.state = AgentStateEnum.AVAILABLE;
+ }
+ }
+
+ private getFormattedResultsForDAG(): Map {
+ const result = new Map();
+ const visible = this.workflowResultState.getAllVisible();
+ for (const [operatorId, entry] of visible) {
+ result.set(operatorId, formatOperatorResult(operatorId, entry.operatorInfo, this.workflowState));
+ }
+ return result;
+ }
+
+ stop(): void {
+ this.state = AgentStateEnum.STOPPING;
+ if (this.abortController) {
+ this.abortController.abort();
+ }
+ }
+
+ clearHistory(): void {
+ this.reActStepsByMessageId.clear();
+ this.stepsById.clear();
+ this.currentMessageId = undefined;
+ this.head = INITIAL_STEP_ID;
+ const initialStep: ReActStep = {
+ id: INITIAL_STEP_ID,
+ messageId: "initial",
+ stepId: -1,
+ timestamp: Date.now(),
+ role: "user",
+ content: "",
+ isBegin: true,
+ isEnd: true,
+ };
+ this.stepsById.set(INITIAL_STEP_ID, initialStep);
+ }
+
+ private getOperatorIdsFromStep(step: ReActStep): { added: string[]; modified: string[] } {
+ const added: string[] = [];
+ const modified: string[] = [];
+
+ if (!step.toolResults) {
+ return { added, modified };
+ }
+
+ for (const result of step.toolResults) {
+ if (result.isError || !result.output) continue;
+
+ const toolCall = step.toolCalls?.find(tc => tc.toolCallId === result.toolCallId);
+ const toolName = toolCall?.toolName || "";
+
+ const outputStr = typeof result.output === "string" ? result.output : JSON.stringify(result.output);
+
+ const addedMatch = outputStr.match(/Added operator ([a-zA-Z0-9_-]+)/);
+ if (addedMatch && (toolName === "addOperator" || toolName.toLowerCase().includes("add"))) {
+ added.push(addedMatch[1]);
+ continue;
+ }
+
+ const modifiedMatch = outputStr.match(/Operator ([a-zA-Z0-9_-]+) modified/);
+ if (modifiedMatch && (toolName === "modifyOperator" || toolName.toLowerCase().includes("modify"))) {
+ modified.push(modifiedMatch[1]);
+ continue;
+ }
+
+ try {
+ const output = JSON.parse(outputStr);
+ if (output.operatorId) {
+ if (toolName === "addOperator" || toolName === "addCodeOperator") {
+ added.push(output.operatorId);
+ } else if (toolName === "modifyOperator" || toolName === "modifyCodeOperator") {
+ modified.push(output.operatorId);
+ }
+ }
+ } catch {}
+ }
+
+ return { added, modified };
+ }
+
+ public getReActStepsByOperatorIds(operatorIds: string[]): ReActStep[] {
+ const allSteps = this.getReActSteps();
+ if (!operatorIds || operatorIds.length === 0) {
+ return allSteps;
+ }
+
+ const operatorIdSet = new Set(operatorIds);
+ const relevantSteps: ReActStep[] = [];
+
+ for (const step of allSteps) {
+ const { added, modified } = this.getOperatorIdsFromStep(step);
+
+ const affectsOperator = [...added, ...modified].some(id => operatorIdSet.has(id));
+
+ if (affectsOperator) {
+ relevantSteps.push(step);
+ }
+ }
+
+ return relevantSteps;
+ }
+
+ destroy(): void {
+ if (this.workflowChangeSubscription) {
+ this.workflowChangeSubscription.unsubscribe();
+ this.workflowChangeSubscription = null;
+ }
+
+ this.workflowState.destroy();
+
+ this.websockets.clear();
+
+ this.reActStepsByMessageId.clear();
+ this.stepsById.clear();
+ this.currentMessageId = undefined;
+ }
+}
diff --git a/agent-service/src/agent/tools/index.ts b/agent-service/src/agent/tools/index.ts
new file mode 100644
index 00000000000..7e2d9570703
--- /dev/null
+++ b/agent-service/src/agent/tools/index.ts
@@ -0,0 +1,22 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export * from "./tools-utility";
+export * from "./workflow-crud-tools";
+export * from "./workflow-execution-tools";
diff --git a/agent-service/src/agent/tools/result-formatting.ts b/agent-service/src/agent/tools/result-formatting.ts
new file mode 100644
index 00000000000..9a11ba50853
--- /dev/null
+++ b/agent-service/src/agent/tools/result-formatting.ts
@@ -0,0 +1,138 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import type { OperatorInfo } from "../../types/execution";
+import type { WorkflowState } from "../workflow-state";
+import { formatExecuteOperatorResult } from "./tools-utility";
+
+export function formatOperatorResult(operatorId: string, opInfo: OperatorInfo, workflowState: WorkflowState): string {
+ if (opInfo.error) {
+ return `[ERROR] ${opInfo.error}`;
+ }
+
+ if (!opInfo.result || !Array.isArray(opInfo.result)) {
+ return "(no result data)";
+ }
+
+ const jsonArray = opInfo.result as Record[];
+ const headers =
+ jsonArray.length > 0
+ ? Object.keys(jsonArray[0]).filter(k => k !== "__row_index__" && k !== "__is_visualization__")
+ : [];
+ const columns = headers.length;
+
+ const isViz = jsonArray.length > 0 && jsonArray[0]["__is_visualization__"] === true;
+ const serializableArray = isViz
+ ? jsonArray.map(row => {
+ const cleaned: Record = {};
+ for (const key of Object.keys(row)) {
+ if (key === "__is_visualization__") continue;
+ if (key === "html-content" || key === "json-content") {
+ cleaned[key] = "";
+ } else {
+ cleaned[key] = row[key];
+ }
+ }
+ return cleaned;
+ })
+ : jsonArray;
+
+ const dataString = jsonToTableFormat(serializableArray);
+
+ const metadataLines = [
+ formatInputOutputMetadata(workflowState, operatorId, opInfo, columns),
+ ...(opInfo.warnings ?? []),
+ ].filter(Boolean);
+
+ const briefSummary = formatExecuteOperatorResult(operatorId);
+ return [briefSummary, ...metadataLines, dataString].filter(Boolean).join("\n");
+}
+
+function formatInputOutputMetadata(
+ workflowState: WorkflowState,
+ operatorId: string,
+ opInfo: OperatorInfo,
+ outputColumns: number
+): string {
+ const outputRows = opInfo.totalRowCount ?? opInfo.outputTuples;
+ const outputLine = `Output table shape: (${outputRows}, ${outputColumns})`;
+
+ const inputShapes = opInfo.inputPortShapes;
+ if (!inputShapes || inputShapes.length === 0) {
+ return outputLine;
+ }
+
+ const inputLinks = workflowState.getAllLinks().filter(l => l.target.operatorID === operatorId);
+ const portIndexToUpstream = new Map();
+ const op = workflowState.getOperator(operatorId);
+ for (const link of inputLinks) {
+ const portIdx = op?.inputPorts.findIndex(p => p.portID === link.target.portID) ?? -1;
+ if (portIdx >= 0) {
+ portIndexToUpstream.set(portIdx, link.source.operatorID);
+ }
+ }
+
+ const inputPart = inputShapes
+ .sort((a, b) => a.portIndex - b.portIndex)
+ .map(p => {
+ const name = portIndexToUpstream.get(p.portIndex) ?? `input${p.portIndex}`;
+ return `${name}(${p.rows}, ${p.columns})`;
+ })
+ .join(", ");
+
+ return `Input operator(table shape): ${inputPart}\n${outputLine}`;
+}
+
+function jsonToTableFormat(jsonResult: Record[]): string {
+ if (!jsonResult || jsonResult.length === 0) return "";
+
+ const hasRowIndex = "__row_index__" in jsonResult[0];
+ const headers = Object.keys(jsonResult[0]).filter(h => h !== "__row_index__");
+ if (headers.length === 0) return "";
+
+ const headerLine = "\t" + headers.join("\t");
+ const formattedRows: string[] = [];
+ let prevIndex = -1;
+
+ for (let i = 0; i < jsonResult.length; i++) {
+ const row = jsonResult[i];
+ const rowIndex = hasRowIndex ? (row["__row_index__"] as number) : i;
+
+ if (prevIndex >= 0 && rowIndex > prevIndex + 1) {
+ const dots = headers.map(() => "...").join("\t");
+ formattedRows.push(`...\t${dots}`);
+ }
+ prevIndex = rowIndex;
+
+ const cells = headers.map(h => {
+ const val = row[h];
+ if (val === null) return "NaN";
+ if (val === undefined) return "";
+ if (typeof val === "number" || typeof val === "boolean") return String(val);
+ if (typeof val === "string") {
+ if (val === "NULL") return "NaN";
+ return val.replace(/\t/g, "\\t").replace(/\n/g, "\\n");
+ }
+ return JSON.stringify(val);
+ });
+ formattedRows.push(`${rowIndex}\t${cells.join("\t")}`);
+ }
+
+ return [headerLine, ...formattedRows].join("\n");
+}
diff --git a/agent-service/src/agent/tools/tools-utility.ts b/agent-service/src/agent/tools/tools-utility.ts
new file mode 100644
index 00000000000..3bc1a1d4312
--- /dev/null
+++ b/agent-service/src/agent/tools/tools-utility.ts
@@ -0,0 +1,70 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export function createToolResult(message: string): string {
+ return message;
+}
+
+export function createErrorResult(error: string): string {
+ return `[ERROR] ${error}`;
+}
+
+function formatLinkDescription(sourceOperatorId: string, targetOperatorId: string): string {
+ return `${sourceOperatorId} --> ${targetOperatorId}`;
+}
+
+export function formatAddOperatorResult(
+ operatorId: string,
+ numInputPorts: number,
+ numOutputPorts: number,
+ createdLinks?: { source: string; target: string }[],
+ deletedLinks?: { source: string; target: string }[]
+): string {
+ let summary = `Added operator ${operatorId}, input ports: ${numInputPorts}, output ports: ${numOutputPorts}`;
+ if (deletedLinks && deletedLinks.length > 0) {
+ summary += `, deleted links: [${deletedLinks.map(l => formatLinkDescription(l.source, l.target)).join(", ")}]`;
+ }
+ if (createdLinks && createdLinks.length > 0) {
+ summary += `, created links: [${createdLinks.map(l => formatLinkDescription(l.source, l.target)).join(", ")}]`;
+ }
+ return summary;
+}
+
+export function formatModifyOperatorResult(
+ operatorId: string,
+ createdLinks?: { source: string; target: string }[],
+ deletedLinks?: { source: string; target: string }[]
+): string {
+ let summary = `Operator ${operatorId} modified`;
+ if (deletedLinks && deletedLinks.length > 0) {
+ summary += `, deleted links: [${deletedLinks.map(l => formatLinkDescription(l.source, l.target)).join(", ")}]`;
+ }
+ if (createdLinks && createdLinks.length > 0) {
+ summary += `, created links: [${createdLinks.map(l => formatLinkDescription(l.source, l.target)).join(", ")}]`;
+ }
+ return summary;
+}
+
+export function formatExecuteOperatorResult(operatorId: string): string {
+ return `Executed operator ${operatorId}`;
+}
+
+export function formatOperatorError(operatorId: string, error: string): string {
+ return `Error on operator ${operatorId}: ${error}`;
+}
diff --git a/agent-service/src/agent/tools/workflow-crud-tools.ts b/agent-service/src/agent/tools/workflow-crud-tools.ts
new file mode 100644
index 00000000000..54a0b29db99
--- /dev/null
+++ b/agent-service/src/agent/tools/workflow-crud-tools.ts
@@ -0,0 +1,346 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { z } from "zod";
+import { tool } from "ai";
+import { WorkflowState } from "../workflow-state";
+import { autoLayoutWorkflow } from "../util/auto-layout";
+import { WorkflowUtilService } from "../util/workflow-utils";
+import type { OperatorLink } from "../../types/workflow";
+import {
+ createToolResult,
+ createErrorResult,
+ formatAddOperatorResult,
+ formatModifyOperatorResult,
+ formatOperatorError,
+} from "./tools-utility";
+import {
+ type WorkflowSystemMetadata,
+ formatValidationErrors,
+ formatCompactSchemaForError,
+} from "../util/workflow-system-metadata";
+
+export interface ToolContext {
+ metadataStore?: WorkflowSystemMetadata;
+ settings?: {
+ maxOperatorResultCharLimit?: number;
+ toolTimeoutMs?: number;
+ executionTimeoutMs?: number;
+ };
+}
+
+export const TOOL_NAME_ADD_OPERATOR = "addOperator";
+export const TOOL_NAME_MODIFY_OPERATOR = "modifyOperator";
+export const TOOL_NAME_DELETE_OPERATOR = "deleteOperator";
+
+function formatInputArgs(args: Record): string {
+ const compact: Record = {};
+ for (const [key, value] of Object.entries(args)) {
+ if (value !== undefined) compact[key] = value;
+ }
+ return `Input: ${JSON.stringify(compact)}`;
+}
+
+export function createAddOperatorTool(
+ workflowState: WorkflowState,
+ operatorSchemas: Map,
+ context?: ToolContext
+) {
+ const workflowUtil = context?.metadataStore ? new WorkflowUtilService(context.metadataStore, workflowState) : null;
+
+ return tool({
+ description: `Add a new operator to the workflow. Use getOperatorSchema first to understand required properties.
+
+Examples:
+1. Add a source operator (no inputs):
+ { "operatorId": "op1", "operatorType": "TableFileScan", "properties": { "fileName": "data.csv" }, "summary": "Load CSV data" }
+
+2. Add an operator with input connections:
+ { "operatorId": "op2", "operatorType": "TableFilter", "properties": { "predicates": [...] }, "inputOperatorIds": { "0": ["op1"] }, "summary": "Filter rows by condition" }`,
+ inputSchema: z.object({
+ operatorId: z
+ .string()
+ .describe(
+ "Name of Operator. Use the format 'op' followed by an incrementing number starting from 1 (e.g., op1, op2, op3)."
+ ),
+ operatorType: z.string().describe("The operator type (e.g., 'DataProcessing', 'Aggregate')"),
+ properties: z.record(z.any()).describe("Properties to set on the operator"),
+ inputOperatorIds: z
+ .record(z.array(z.string()))
+ .optional()
+ .describe(
+ "Mapping from input port index to an ordered list of source operator IDs that connect to that port. " +
+ 'E.g. {"0": ["opA", "opB"], "1": ["opC"]} connects opA and opB to input port 0, opC to input port 1. ' +
+ "Source operators that load files (e.g. CSVFileScan) should NOT have any input operators."
+ ),
+ summary: z.string().describe("Very brief summary of operator behavior. Within 5 words"),
+ }),
+ execute: async (args: {
+ operatorId: string;
+ operatorType: string;
+ properties?: Record;
+ inputOperatorIds?: Record;
+ summary: string;
+ }) => {
+ try {
+ const inputInfo = formatInputArgs(args);
+
+ const schemaEntry = operatorSchemas.get(args.operatorType);
+ if (!schemaEntry) {
+ return createErrorResult(
+ `Unknown operator type: "${args.operatorType}". Available types: ${[...operatorSchemas.keys()].join(", ")}. ${inputInfo}`
+ );
+ }
+
+ if (context?.metadataStore && args.properties) {
+ const validation = context.metadataStore.validateOperatorProperties(args.operatorType, args.properties);
+ if (!validation.isValid) {
+ const compactSchema = context.metadataStore.getCompactSchema(args.operatorType);
+ const schemaStr = compactSchema ? ` Expected: ${formatCompactSchemaForError(compactSchema)}.` : "";
+ return createErrorResult(
+ `Invalid properties for "${args.operatorType}": ${formatValidationErrors(validation)}.${schemaStr} ${inputInfo}`
+ );
+ }
+ }
+
+ if (!workflowUtil) {
+ return createErrorResult(`Metadata store not available for operator creation. ${inputInfo}`);
+ }
+
+ if (!/^op\d+$/.test(args.operatorId)) {
+ return createErrorResult(
+ `Invalid operatorId: "${args.operatorId}". Must follow the format "op" followed by a number (e.g., op1, op2, op3). ${inputInfo}`
+ );
+ }
+
+ const existing = workflowState.getOperator(args.operatorId);
+ if (existing) {
+ return createErrorResult(
+ `Operator with ID "${args.operatorId}" already exists. Use modifyOperator to update it, or choose a different ID. ${inputInfo}`
+ );
+ }
+
+ let operator = workflowUtil.getNewOperatorPredicate(args.operatorType, args.summary);
+ operator = {
+ ...operator,
+ operatorID: args.operatorId,
+ operatorProperties: { ...operator.operatorProperties, ...args.properties },
+ };
+
+ workflowState.addOperator(operator);
+
+ const createdLinkPairs: { source: string; target: string }[] = [];
+ if (args.inputOperatorIds) {
+ const addedOperator = workflowState.getOperator(operator.operatorID)!;
+ for (const [portIndexStr, sourceOpIds] of Object.entries(args.inputOperatorIds)) {
+ const targetPortIdx = parseInt(portIndexStr, 10);
+ if (isNaN(targetPortIdx) || targetPortIdx < 0) {
+ return createErrorResult(
+ `Invalid input port index: "${portIndexStr}". Must be a non-negative integer. ${inputInfo}`
+ );
+ }
+ if (targetPortIdx >= addedOperator.inputPorts.length) {
+ return createErrorResult(
+ `Input port index ${targetPortIdx} out of range. Operator "${args.operatorId}" has ${addedOperator.inputPorts.length} input port(s). ${inputInfo}`
+ );
+ }
+ const targetPortId = addedOperator.inputPorts[targetPortIdx].portID;
+
+ for (const sourceOpId of sourceOpIds) {
+ const sourceOp = workflowState.getOperator(sourceOpId);
+ if (!sourceOp) {
+ return createErrorResult(
+ `Source operator "${sourceOpId}" not found. Make sure it exists before referencing it in inputOperatorIds. ${inputInfo}`
+ );
+ }
+ const sourcePortId = sourceOp.outputPorts.length > 0 ? sourceOp.outputPorts[0].portID : "output-0";
+
+ const linkId = workflowState.generateLinkId();
+ const link: OperatorLink = {
+ linkID: linkId,
+ source: { operatorID: sourceOpId, portID: sourcePortId },
+ target: { operatorID: args.operatorId, portID: targetPortId },
+ };
+ workflowState.addLink(link);
+ createdLinkPairs.push({ source: sourceOpId, target: args.operatorId });
+ }
+ }
+ }
+
+ autoLayoutWorkflow(workflowState);
+
+ const finalOperator = workflowState.getOperator(operator.operatorID) || operator;
+ const numInputPorts = finalOperator.inputPorts.length;
+ const numOutputPorts = finalOperator.outputPorts.length;
+
+ let resultMsg = formatAddOperatorResult(
+ operator.operatorID,
+ numInputPorts,
+ numOutputPorts,
+ createdLinkPairs.length > 0 ? createdLinkPairs : undefined
+ );
+
+ return createToolResult(resultMsg);
+ } catch (error: any) {
+ return createErrorResult(error.message || String(error));
+ }
+ },
+ });
+}
+
+export function createModifyOperatorTool(workflowState: WorkflowState, context?: ToolContext) {
+ return tool({
+ description: `Modify an existing operator's properties, input links, or both.
+
+Examples:
+1. Modify properties only:
+ { "operatorId": "agg", "properties": { "groupByKeys": ["city"] }, "summary": "Group by city" }
+
+2. Modify input links only (replaces all existing incoming links):
+ { "operatorId": "join_op", "inputOperatorIds": { "0": ["users"], "1": ["orders"] }, "summary": "Re-link join inputs" }
+
+3. Modify both properties and links:
+ { "operatorId": "filter", "properties": { "predicates": [...] }, "inputOperatorIds": { "0": ["cleaned"] }, "summary": "Update filter and re-link" }`,
+ inputSchema: z.object({
+ operatorId: z.string().describe("ID of the operator to modify"),
+ properties: z.record(z.any()).optional().describe("Properties to update (merged with existing)"),
+ inputOperatorIds: z
+ .record(z.array(z.string()))
+ .optional()
+ .describe(
+ "Mapping from input port index to an ordered list of source operator IDs. " +
+ "If provided, all existing incoming links are deleted and replaced with these. " +
+ 'E.g. {"0": ["opA", "opB"], "1": ["opC"]} connects opA and opB to input port 0, opC to input port 1.'
+ ),
+ summary: z.string().describe("Very brief summary of operator behavior after your modification. Within 5 words"),
+ }),
+ execute: async (args: {
+ operatorId: string;
+ properties?: Record;
+ inputOperatorIds?: Record;
+ summary?: string;
+ }) => {
+ try {
+ const inputInfo = formatInputArgs(args);
+
+ const operator = workflowState.getOperator(args.operatorId);
+ if (!operator) return createErrorResult(`Operator ${args.operatorId} not found. ${inputInfo}`);
+
+ if (args.properties && context?.metadataStore) {
+ const mergedProperties = { ...operator.operatorProperties, ...args.properties };
+ const validation = context.metadataStore.validateOperatorProperties(operator.operatorType, mergedProperties);
+ if (!validation.isValid) {
+ const compactSchema = context.metadataStore.getCompactSchema(operator.operatorType);
+ const schemaStr = compactSchema ? ` Expected: ${formatCompactSchemaForError(compactSchema)}.` : "";
+ return createErrorResult(
+ `Invalid properties for "${operator.operatorType}": ${formatValidationErrors(validation)}.${schemaStr} ${inputInfo}`
+ );
+ }
+ }
+
+ const createdLinkPairs: { source: string; target: string }[] = [];
+ const deletedLinkPairs: { source: string; target: string }[] = [];
+
+ if (args.properties) {
+ workflowState.updateOperatorProperties(args.operatorId, args.properties);
+ }
+
+ if (args.summary) {
+ workflowState.updateOperatorDisplayName(args.operatorId, args.summary);
+ }
+
+ if (args.inputOperatorIds) {
+ const currentLinks = workflowState
+ .getLinksConnectedToOperator(args.operatorId)
+ .filter(link => link.target.operatorID === args.operatorId);
+ for (const link of currentLinks) {
+ deletedLinkPairs.push({ source: link.source.operatorID, target: link.target.operatorID });
+ workflowState.deleteLink(link.linkID);
+ }
+
+ for (const [portIndexStr, sourceOpIds] of Object.entries(args.inputOperatorIds)) {
+ const targetPortIdx = parseInt(portIndexStr, 10);
+ if (isNaN(targetPortIdx) || targetPortIdx < 0) {
+ return createErrorResult(
+ `Invalid input port index: "${portIndexStr}". Must be a non-negative integer. ${inputInfo}`
+ );
+ }
+ if (targetPortIdx >= operator.inputPorts.length) {
+ return createErrorResult(
+ `Input port index ${targetPortIdx} out of range. Operator "${args.operatorId}" has ${operator.inputPorts.length} input port(s). ${inputInfo}`
+ );
+ }
+ const targetPortId = operator.inputPorts[targetPortIdx].portID;
+
+ for (const sourceOpId of sourceOpIds) {
+ const sourceOp = workflowState.getOperator(sourceOpId);
+ if (!sourceOp) {
+ return createErrorResult(
+ `Source operator "${sourceOpId}" not found. Make sure it exists before referencing it in inputOperatorIds. ${inputInfo}`
+ );
+ }
+ const sourcePortId = sourceOp.outputPorts.length > 0 ? sourceOp.outputPorts[0].portID : "output-0";
+
+ const linkId = workflowState.generateLinkId();
+ const link: OperatorLink = {
+ linkID: linkId,
+ source: { operatorID: sourceOpId, portID: sourcePortId },
+ target: { operatorID: args.operatorId, portID: targetPortId },
+ };
+ workflowState.addLink(link);
+ createdLinkPairs.push({ source: sourceOpId, target: args.operatorId });
+ }
+ }
+
+ autoLayoutWorkflow(workflowState);
+ }
+
+ let resultMsg = formatModifyOperatorResult(
+ args.operatorId,
+ createdLinkPairs.length > 0 ? createdLinkPairs : undefined,
+ deletedLinkPairs.length > 0 ? deletedLinkPairs : undefined
+ );
+
+ return createToolResult(resultMsg);
+ } catch (error: any) {
+ return createErrorResult(formatOperatorError(args.operatorId, error.message || String(error)));
+ }
+ },
+ });
+}
+
+export function createDeleteOperatorTool(workflowState: WorkflowState, _context?: ToolContext) {
+ return tool({
+ description: "Delete an operator from the workflow. This also deletes all connected links.",
+ inputSchema: z.object({
+ operatorId: z.string().describe("ID of the operator to delete"),
+ }),
+ execute: async (args: { operatorId: string }) => {
+ try {
+ const deleted = workflowState.deleteOperator(args.operatorId);
+ if (!deleted) {
+ return createErrorResult(`Operator ${args.operatorId} not found`);
+ }
+ return createToolResult(`Deleted operator: ${args.operatorId}`);
+ } catch (error: any) {
+ return createErrorResult(error.message || String(error));
+ }
+ },
+ });
+}
diff --git a/agent-service/src/agent/tools/workflow-execution-tools.ts b/agent-service/src/agent/tools/workflow-execution-tools.ts
new file mode 100644
index 00000000000..15fa81ff977
--- /dev/null
+++ b/agent-service/src/agent/tools/workflow-execution-tools.ts
@@ -0,0 +1,604 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { z } from "zod";
+import { tool } from "ai";
+import { createErrorResult, formatExecuteOperatorResult } from "./tools-utility";
+import type { WorkflowState } from "../workflow-state";
+import { getBackendConfig } from "../../api/backend-api";
+import { env } from "../../config/env";
+import type { LogicalPlan, LogicalLink } from "../../api/execution-api";
+import type { OperatorInfo, SyncExecutionResult } from "../../types/execution";
+import { WorkflowSystemMetadata } from "../util/workflow-system-metadata";
+import { DEFAULT_AGENT_SETTINGS } from "../../types/agent";
+import { createLogger } from "../../logger";
+
+const log = createLogger("ExecutionTools");
+
+export const TOOL_NAME_EXECUTE_OPERATOR = "executeOperator";
+
+export interface ExecutionConfig {
+ userToken: string;
+ workflowId: number;
+ computingUnitId?: number;
+ maxOperatorResultCharLimit?: number;
+ maxOperatorResultCellCharLimit?: number;
+ executionTimeoutMs?: number;
+}
+
+/**
+ * FIFO async lock used to serialize workflow executions per workflow id.
+ *
+ * `acquire()` resolves with a release function once prior holders have
+ * released. Callers must invoke the release in a `finally` to avoid
+ * deadlocking subsequent waiters.
+ */
+class AsyncMutex {
+ private queue: Promise = Promise.resolve();
+
+ async acquire(): Promise<() => void> {
+ let release: () => void;
+ const currentQueue = this.queue;
+
+ this.queue = new Promise(resolve => {
+ release = resolve;
+ });
+
+ await currentQueue;
+
+ return release!;
+ }
+}
+
+const workflowMutexes = new Map();
+
+function getWorkflowMutex(workflowId: number): AsyncMutex {
+ let mutex = workflowMutexes.get(workflowId);
+ if (!mutex) {
+ mutex = new AsyncMutex();
+ workflowMutexes.set(workflowId, mutex);
+ }
+ return mutex;
+}
+
+interface WorkflowValidationResult {
+ isValid: boolean;
+ errors: Record>;
+}
+
+interface OperatorValidation {
+ isValid: boolean;
+ messages: Record;
+}
+
+function validateOperatorSchema(operatorType: string, operatorProperties: Record): OperatorValidation {
+ const metadataStore = WorkflowSystemMetadata.getInstance();
+ const validation = metadataStore.validateOperatorProperties(operatorType, operatorProperties);
+ return validation.isValid ? { isValid: true, messages: {} } : { isValid: false, messages: validation.messages };
+}
+
+function validateOperatorConnection(operatorId: string, workflowState: WorkflowState): OperatorValidation {
+ const operator = workflowState.getOperator(operatorId);
+ if (!operator) {
+ return { isValid: false, messages: { error: `Operator ${operatorId} not found` } };
+ }
+
+ const numInputLinksByPort = new Map();
+ const allLinks = workflowState.getAllLinks();
+
+ for (const link of allLinks) {
+ if (link.target.operatorID === operatorId) {
+ const portID = link.target.portID;
+ numInputLinksByPort.set(portID, (numInputLinksByPort.get(portID) ?? 0) + 1);
+ }
+ }
+
+ let satisfyInput = true;
+ let violationMessage = "";
+
+ for (const port of operator.inputPorts) {
+ const portNumInputs = numInputLinksByPort.get(port.portID) ?? 0;
+
+ if (port.disallowMultiInputs) {
+ if (portNumInputs !== 1) {
+ satisfyInput = false;
+ violationMessage += `${port.displayName ?? port.portID} requires 1 input, has ${portNumInputs}. `;
+ }
+ } else {
+ if (portNumInputs < 1) {
+ satisfyInput = false;
+ violationMessage += `${port.displayName ?? port.portID} requires at least 1 input, has ${portNumInputs}. `;
+ }
+ }
+ }
+
+ return satisfyInput
+ ? { isValid: true, messages: {} }
+ : { isValid: false, messages: { inputs: violationMessage.trim() } };
+}
+
+function combineValidations(...validations: OperatorValidation[]): OperatorValidation {
+ let isValid = true;
+ let messages: Record = {};
+
+ for (const validation of validations) {
+ if (!validation.isValid) {
+ isValid = false;
+ messages = { ...messages, ...validation.messages };
+ }
+ }
+
+ return { isValid, messages };
+}
+
+function validateWorkflow(workflowState: WorkflowState): WorkflowValidationResult {
+ const errors: Record> = {};
+
+ for (const operator of workflowState.getAllEnabledOperators()) {
+ const schemaValidation = validateOperatorSchema(operator.operatorType, operator.operatorProperties);
+ const connectionValidation = validateOperatorConnection(operator.operatorID, workflowState);
+ const combined = combineValidations(schemaValidation, connectionValidation);
+
+ if (!combined.isValid) {
+ errors[operator.operatorID] = combined.messages;
+ }
+ }
+
+ return {
+ isValid: Object.keys(errors).length === 0,
+ errors,
+ };
+}
+
+function formatWorkflowValidationErrors(validationResult: WorkflowValidationResult): string {
+ if (validationResult.isValid) return "";
+
+ const lines: string[] = ["Workflow validation failed:"];
+ for (const [operatorId, fieldErrors] of Object.entries(validationResult.errors)) {
+ lines.push(` Operator ${operatorId}:`);
+ for (const [field, message] of Object.entries(fieldErrors)) {
+ lines.push(` - ${field}: ${message}`);
+ }
+ }
+ return lines.join("\n");
+}
+
+function buildLogicalPlan(workflowState: WorkflowState, opsToViewResult?: string[]): LogicalPlan {
+ const useSubDAG = opsToViewResult && opsToViewResult.length === 1;
+ const targetOperatorId = useSubDAG ? opsToViewResult[0] : undefined;
+
+ let operatorsList: { operatorID: string; operatorType: string; [key: string]: any }[];
+ let linksList: LogicalLink[];
+
+ const getInputPortOrdinal = (operatorID: string, inputPortID: string): number => {
+ const op = workflowState.getOperator(operatorID);
+ if (!op) return 0;
+ const idx = op.inputPorts.findIndex(port => port.portID === inputPortID);
+ return idx >= 0 ? idx : 0;
+ };
+
+ const getOutputPortOrdinal = (operatorID: string, outputPortID: string): number => {
+ const op = workflowState.getOperator(operatorID);
+ if (!op) return 0;
+ const idx = op.outputPorts.findIndex(port => port.portID === outputPortID);
+ return idx >= 0 ? idx : 0;
+ };
+
+ if (targetOperatorId) {
+ const subDAG = workflowState.getSubDAG(targetOperatorId);
+
+ operatorsList = subDAG.operators.map(op => ({
+ ...op.operatorProperties,
+ operatorID: op.operatorID,
+ operatorType: op.operatorType,
+ inputPorts: op.inputPorts,
+ outputPorts: op.outputPorts,
+ }));
+
+ linksList = subDAG.links.map(link => ({
+ fromOpId: link.source.operatorID,
+ fromPortId: { id: getOutputPortOrdinal(link.source.operatorID, link.source.portID), internal: false },
+ toOpId: link.target.operatorID,
+ toPortId: { id: getInputPortOrdinal(link.target.operatorID, link.target.portID), internal: false },
+ }));
+ } else {
+ operatorsList = workflowState.getAllEnabledOperators().map(op => ({
+ ...op.operatorProperties,
+ operatorID: op.operatorID,
+ operatorType: op.operatorType,
+ inputPorts: op.inputPorts,
+ outputPorts: op.outputPorts,
+ }));
+
+ linksList = workflowState.getAllLinks().map(link => ({
+ fromOpId: link.source.operatorID,
+ fromPortId: { id: getOutputPortOrdinal(link.source.operatorID, link.source.portID), internal: false },
+ toOpId: link.target.operatorID,
+ toPortId: { id: getInputPortOrdinal(link.target.operatorID, link.target.portID), internal: false },
+ }));
+ }
+
+ let allOpsToView: string[];
+ if (opsToViewResult && opsToViewResult.length > 0) {
+ const operatorIds = new Set(operatorsList.map(op => op.operatorID));
+ allOpsToView = opsToViewResult.filter(id => operatorIds.has(id));
+ } else {
+ allOpsToView = operatorsList
+ .filter(op => !linksList.some(link => link.fromOpId === op.operatorID))
+ .map(op => op.operatorID);
+ }
+
+ return {
+ operators: operatorsList,
+ links: linksList,
+ opsToViewResult: allOpsToView,
+ };
+}
+
+async function executeWorkflowHttp(
+ config: ExecutionConfig,
+ logicalPlan: LogicalPlan,
+ options: { abortSignal?: AbortSignal } = {}
+): Promise {
+ const backendConfig = getBackendConfig();
+
+ const workflowId = config.workflowId;
+ const computingUnitId = config.computingUnitId ?? 0;
+
+ // In k8s each computing unit is a separate pod, so the endpoint varies per cuid.
+ const executionEndpoint = env.EXECUTION_ENDPOINT_TEMPLATE
+ ? env.EXECUTION_ENDPOINT_TEMPLATE.replace("{cuid}", String(computingUnitId))
+ : backendConfig.executionEndpoint;
+
+ const url = `${executionEndpoint}/api/execution/${workflowId}/${computingUnitId}/run`;
+
+ const timeoutSeconds = config.executionTimeoutMs
+ ? Math.ceil(config.executionTimeoutMs / 1000)
+ : Math.ceil(DEFAULT_AGENT_SETTINGS.executionTimeoutMs / 1000);
+
+ const request = {
+ executionName: "agent-execution",
+ logicalPlan: {
+ operators: logicalPlan.operators,
+ links: logicalPlan.links,
+ opsToViewResult: logicalPlan.opsToViewResult || [],
+ opsToReuseResult: [],
+ },
+ targetOperatorIds: logicalPlan.opsToViewResult || [],
+ timeoutSeconds,
+ maxOperatorResultCharLimit: config.maxOperatorResultCharLimit ?? DEFAULT_AGENT_SETTINGS.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit:
+ config.maxOperatorResultCellCharLimit ?? DEFAULT_AGENT_SETTINGS.maxOperatorResultCellCharLimit,
+ };
+
+ log.debug(
+ {
+ url,
+ maxOperatorResultCharLimit: request.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: request.maxOperatorResultCellCharLimit,
+ },
+ "executing workflow"
+ );
+
+ try {
+ const response = await fetch(url, {
+ method: "POST",
+ headers: {
+ "Content-Type": "application/json",
+ Authorization: `Bearer ${config.userToken}`,
+ },
+ body: JSON.stringify(request),
+ signal: options.abortSignal,
+ });
+
+ if (!response.ok) {
+ const errorText = await response.text();
+ throw new Error(`Execution request failed: ${response.status} ${response.statusText} - ${errorText}`);
+ }
+
+ return (await response.json()) as SyncExecutionResult;
+ } catch (error) {
+ if (error instanceof Error && error.name === "AbortError") {
+ throw error;
+ }
+ log.error({ err: error }, "execution failed");
+ return {
+ success: false,
+ state: "Error",
+ operators: {},
+ errors: [error instanceof Error ? error.message : "Unknown error"],
+ };
+ }
+}
+
+function formatInputOutput(
+ workflowState: WorkflowState,
+ operatorId: string,
+ opInfo: OperatorInfo,
+ outputColumns: number
+): string {
+ const outputRows = opInfo.totalRowCount ?? opInfo.outputTuples;
+ const outputLine = `Output table shape: (${outputRows}, ${outputColumns})`;
+
+ const inputShapes = opInfo.inputPortShapes;
+ if (!inputShapes || inputShapes.length === 0) {
+ return outputLine;
+ }
+
+ const inputLinks = workflowState.getAllLinks().filter(l => l.target.operatorID === operatorId);
+ const portIndexToUpstream = new Map();
+ const op = workflowState.getOperator(operatorId);
+ for (const link of inputLinks) {
+ const portIdx = op?.inputPorts.findIndex(p => p.portID === link.target.portID) ?? -1;
+ if (portIdx >= 0) {
+ portIndexToUpstream.set(portIdx, link.source.operatorID);
+ }
+ }
+
+ const inputPart = inputShapes
+ .sort((a, b) => a.portIndex - b.portIndex)
+ .map(p => {
+ const name = portIndexToUpstream.get(p.portIndex) ?? `input${p.portIndex}`;
+ return `${name}(${p.rows}, ${p.columns})`;
+ })
+ .join(", ");
+
+ return `Input operator(table shape): ${inputPart}\n${outputLine}`;
+}
+
+function formatExecutionError(
+ compilationErrors?: Record,
+ operatorErrors?: Array<{ operatorId: string; error: string }>,
+ generalErrors?: string[]
+): string {
+ const lines: string[] = ["Execution failed due to the following error:"];
+
+ if (compilationErrors && Object.keys(compilationErrors).length > 0) {
+ lines.push("Compilation error:");
+ for (const [key, value] of Object.entries(compilationErrors)) {
+ lines.push(` ${key}: ${value}`);
+ }
+ }
+
+ if (operatorErrors && operatorErrors.length > 0) {
+ lines.push("Execution error:");
+ for (const { operatorId, error } of operatorErrors) {
+ lines.push(` ${operatorId}: ${error}`);
+ }
+ }
+
+ if (generalErrors && generalErrors.length > 0) {
+ lines.push("Error:");
+ for (const error of generalErrors) {
+ lines.push(` ${error}`);
+ }
+ }
+
+ return lines.join("\n");
+}
+
+function jsonToTableFormat(jsonResult: Record[]): string {
+ if (!jsonResult || jsonResult.length === 0) return "";
+
+ const hasRowIndex = jsonResult.length > 0 && "__row_index__" in jsonResult[0];
+ const headers = Object.keys(jsonResult[0]).filter(h => h !== "__row_index__");
+ // Leading tab aligns headers with the index column (pandas __repr__ style).
+ const headerLine = "\t" + headers.join("\t");
+
+ const formattedRows: string[] = [];
+ let prevIndex = -1;
+
+ for (let i = 0; i < jsonResult.length; i++) {
+ const row = jsonResult[i];
+ const rowIndex = hasRowIndex ? (row["__row_index__"] as number) : i;
+
+ if (prevIndex >= 0 && rowIndex > prevIndex + 1) {
+ const dots = headers.map(() => "...").join("\t");
+ formattedRows.push(`...\t${dots}`);
+ }
+ prevIndex = rowIndex;
+
+ const cells = headers.map(h => {
+ const val = row[h];
+ if (val === null) return "NaN";
+ if (val === undefined) return "";
+ if (typeof val === "number" || typeof val === "boolean") return String(val);
+ if (typeof val === "string") {
+ if (val === "NULL") return "NaN";
+ return val.replace(/\t/g, "\\t").replace(/\n/g, "\\n");
+ }
+ return JSON.stringify(val);
+ });
+ formattedRows.push(`${rowIndex}\t${cells.join("\t")}`);
+ }
+
+ return [headerLine, ...formattedRows].join("\n");
+}
+
+export async function executeOperatorAndFormat(
+ workflowState: WorkflowState,
+ config: ExecutionConfig,
+ operatorId: string,
+ options: {
+ abortSignal?: AbortSignal;
+ onResult?: (operatorId: string, operatorInfo: OperatorInfo) => void;
+ onResultLegacy?: (operatorId: string, backendStats?: Record) => void;
+ } = {}
+): Promise {
+ // Serialize executions per workflow to avoid ConcurrentModificationException on the backend.
+ const release = await getWorkflowMutex(config.workflowId).acquire();
+
+ try {
+ const logicalPlan = buildLogicalPlan(workflowState, [operatorId]);
+
+ if (logicalPlan.operators.length === 0) {
+ return createErrorResult("Cannot execute: workflow has no operators.");
+ }
+
+ // Only block on the target operator's validation errors; upstream issues will
+ // surface as runtime errors that correctly identify the failing operator.
+ const validationResult = validateWorkflow(workflowState);
+ if (!validationResult.isValid) {
+ const targetErrors = validationResult.errors[operatorId];
+ if (targetErrors) {
+ const lines = [`Operator ${operatorId}:`];
+ for (const [field, message] of Object.entries(targetErrors)) {
+ lines.push(` - ${field}: ${message}`);
+ }
+ return createErrorResult(lines.join("\n"));
+ }
+ }
+
+ const result: SyncExecutionResult = await executeWorkflowHttp(config, logicalPlan, {
+ abortSignal: options.abortSignal,
+ });
+
+ if (!result.success) {
+ const compilationErrors =
+ result.state === "CompilationFailed" || result.state === "ValidationFailed"
+ ? result.compilationErrors
+ : undefined;
+
+ const operatorErrors =
+ result.state === "Failed"
+ ? Object.entries(result.operators)
+ .filter(([_, op]) => op.error)
+ .map(([opId, op]) => ({ operatorId: opId, error: op.error! }))
+ : undefined;
+
+ const generalErrors = result.state === "Killed" ? ["Workflow execution was killed (timeout)."] : result.errors;
+
+ const errorText = formatExecutionError(compilationErrors, operatorErrors, generalErrors);
+
+ if (options.onResult) {
+ const errorInfo: OperatorInfo = {
+ state: result.state,
+ inputTuples: 0,
+ outputTuples: 0,
+ resultMode: "table",
+ error: errorText,
+ };
+ options.onResult(operatorId, errorInfo);
+ }
+
+ return createErrorResult(errorText);
+ }
+
+ const opInfo = result.operators[operatorId];
+ if (!opInfo) {
+ return createErrorResult(
+ formatExecutionError(undefined, undefined, [`No result found for operator: ${operatorId}`])
+ );
+ }
+
+ if (opInfo.error) {
+ if (options.onResult) {
+ options.onResult(operatorId, opInfo);
+ }
+ return createErrorResult(formatExecutionError(undefined, [{ operatorId, error: opInfo.error }]));
+ }
+
+ if (!opInfo.result || !Array.isArray(opInfo.result)) {
+ return "(no result data)";
+ }
+
+ const jsonArray = opInfo.result as Record[];
+ const headers = jsonArray.length > 0 ? Object.keys(jsonArray[0]).filter(k => k !== "__row_index__") : [];
+ const columns = headers.length;
+
+ // Notify for every operator in the execution so upstream stats are also stored.
+ if (options.onResult) {
+ for (const [opId, info] of Object.entries(result.operators)) {
+ if (info && !info.error) {
+ options.onResult(opId, info);
+ }
+ }
+ }
+
+ let dataString = jsonToTableFormat(jsonArray);
+
+ // Safety-net: TSV serialization may add padding beyond backend's raw-record budget.
+ const charLimit = config.maxOperatorResultCharLimit ?? DEFAULT_AGENT_SETTINGS.maxOperatorResultCharLimit;
+
+ if (dataString.length > charLimit) {
+ const allLines = dataString.split("\n");
+ const headerLine = allLines[0];
+ const dataRows = allLines.slice(1);
+
+ const reservedSize = headerLine.length + 1;
+
+ const halfLimit = Math.floor((charLimit - reservedSize) / 2);
+
+ let frontSize = 0;
+ const frontRows: string[] = [];
+ for (const row of dataRows) {
+ const rowLen = row.length + 1;
+ if (frontSize + rowLen > halfLimit && frontRows.length > 0) break;
+ frontRows.push(row);
+ frontSize += rowLen;
+ }
+
+ let backSize = 0;
+ const backRows: string[] = [];
+ for (let i = dataRows.length - 1; i >= frontRows.length; i--) {
+ const rowLen = dataRows[i].length + 1;
+ if (backSize + rowLen > halfLimit && backRows.length > 0) break;
+ backRows.unshift(dataRows[i]);
+ backSize += rowLen;
+ }
+
+ const keptRows = [...frontRows, ...backRows];
+ dataString = [headerLine, ...keptRows].join("\n");
+ }
+
+ const shapeLine = formatInputOutput(workflowState, operatorId, opInfo, columns);
+
+ const warningLines = opInfo.warnings?.map(w => w) ?? [];
+
+ const metadataLines = [shapeLine, ...warningLines].filter(Boolean);
+
+ const briefSummary = formatExecuteOperatorResult(operatorId);
+ return [briefSummary, ...metadataLines, dataString].filter(Boolean).join("\n");
+ } catch (error: any) {
+ if (error.name === "AbortError") {
+ throw error;
+ }
+ return createErrorResult(`Execution failed: ${error.message || String(error)}`);
+ } finally {
+ release();
+ }
+}
+
+export function createExecuteOperatorTool(
+ workflowState: WorkflowState,
+ getConfig: () => ExecutionConfig,
+ onResult?: (operatorId: string, operatorInfo: OperatorInfo) => void
+) {
+ return tool({
+ description:
+ "Execute the workflow and get the specified operator's result. The execution result(if succeeded) includes the shape of the input tables(if any) and output table, and the records in the output table",
+ inputSchema: z.object({
+ operatorId: z.string().describe("The operator ID to view result for."),
+ }),
+ execute: async (args: { operatorId: string }, options: { abortSignal?: AbortSignal }) => {
+ const config = getConfig();
+ return await executeOperatorAndFormat(workflowState, config, args.operatorId, { ...options, onResult });
+ },
+ });
+}
diff --git a/agent-service/src/agent/util/auto-layout.ts b/agent-service/src/agent/util/auto-layout.ts
new file mode 100644
index 00000000000..b17fed4e084
--- /dev/null
+++ b/agent-service/src/agent/util/auto-layout.ts
@@ -0,0 +1,70 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import dagre from "dagre";
+import type { WorkflowState } from "../workflow-state";
+
+// Values mirror frontend joint-graph-wrapper.ts so agent-generated and
+// user-generated layouts visually match.
+const LAYOUT_CONFIG: dagre.GraphLabel = {
+ nodesep: 100,
+ edgesep: 150,
+ ranksep: 100,
+ ranker: "tight-tree",
+ rankdir: "LR",
+};
+
+const NODE_WIDTH = 200;
+const NODE_HEIGHT = 80;
+
+export function autoLayoutWorkflow(workflowState: WorkflowState): void {
+ const operators = workflowState.getAllOperators();
+ const links = workflowState.getAllLinks();
+
+ if (operators.length === 0) {
+ return;
+ }
+
+ const graph = new dagre.graphlib.Graph();
+ graph.setGraph(LAYOUT_CONFIG);
+ graph.setDefaultEdgeLabel(() => ({}));
+
+ for (const operator of operators) {
+ graph.setNode(operator.operatorID, {
+ width: NODE_WIDTH,
+ height: NODE_HEIGHT,
+ });
+ }
+
+ for (const link of links) {
+ graph.setEdge(link.source.operatorID, link.target.operatorID);
+ }
+
+ dagre.layout(graph);
+
+ for (const operator of operators) {
+ const node = graph.node(operator.operatorID);
+ if (node) {
+ workflowState.updateOperatorPosition(operator.operatorID, {
+ x: node.x,
+ y: node.y,
+ });
+ }
+ }
+}
diff --git a/agent-service/src/agent/util/context-utils.ts b/agent-service/src/agent/util/context-utils.ts
new file mode 100644
index 00000000000..195692cbf50
--- /dev/null
+++ b/agent-service/src/agent/util/context-utils.ts
@@ -0,0 +1,288 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+// Output uses plain markdown rather than XML-like tags to reduce format
+// mimicry, where the model echoes the context shape into its output instead
+// of calling tools via the native protocol.
+
+import type { ModelMessage } from "ai";
+import type { WorkflowState } from "../workflow-state";
+import type { OperatorPredicate, OperatorPortSchemaMap, PortSchema } from "../../types/workflow";
+import type { ReActStep } from "../../types/agent";
+import type { WorkflowCompilationResponse, WorkflowFatalError } from "../../api/compile-api";
+import { extractOperatorInputPortSchemaMap } from "./workflow-utils";
+import { createLogger } from "../../logger";
+
+const log = createLogger("ContextAssembler");
+
+export function assembleContext(
+ visibleSteps: ReActStep[],
+ workflowState: WorkflowState,
+ operatorExecutionResults: Map,
+ useRedact: boolean = false,
+ compilationResult?: WorkflowCompilationResponse | null
+): ModelMessage[] {
+ const messageIds: string[] = [];
+ const stepsByMessage = new Map();
+ for (const step of visibleSteps) {
+ let group = stepsByMessage.get(step.messageId);
+ if (!group) {
+ group = [];
+ stepsByMessage.set(step.messageId, group);
+ messageIds.push(step.messageId);
+ }
+ group.push(step);
+ }
+
+ const sections: string[] = [];
+ let completedCount = 0;
+ let hasOngoing = false;
+
+ for (const msgId of messageIds) {
+ const steps = stepsByMessage.get(msgId)!;
+ // A task is completed only when an *agent* step has isEnd=true; user steps
+ // always have isEnd=true because they are single-step messages.
+ const isCompleted = steps.some(s => s.role === "agent" && s.isEnd);
+
+ if (isCompleted) {
+ if (completedCount === 0) {
+ sections.push("# Completed Tasks");
+ }
+ sections.push("");
+ sections.push(serializeTask(steps, "completed"));
+ completedCount++;
+ } else {
+ hasOngoing = true;
+ sections.push("");
+ sections.push("# Ongoing Task");
+ sections.push(serializeTask(steps, "ongoing"));
+ sections.push("");
+ sections.push(
+ "Above is user's request and the steps you already took. You as an assistant please keep working on solving user's request based on the progress of current workflow."
+ );
+ }
+ }
+
+ const dagSection = serializeDag(workflowState, operatorExecutionResults, useRedact, compilationResult);
+ if (dagSection) {
+ sections.push("");
+ sections.push("# Current Dataflow");
+ sections.push(dagSection);
+ }
+
+ const content = sections.join("\n");
+
+ log.debug(
+ {
+ completed: completedCount,
+ ongoing: hasOngoing ? 1 : 0,
+ operatorResults: operatorExecutionResults.size,
+ useRedact,
+ },
+ "built context"
+ );
+
+ return [{ role: "user", content }];
+}
+
+function serializeTask(steps: ReActStep[], status: "completed" | "ongoing"): string {
+ const lines: string[] = [];
+ lines.push(`## Task (${status})`);
+ lines.push("");
+
+ const userStep = steps.find(s => s.role === "user");
+ const assistantSteps = steps.filter(s => s.role === "agent");
+
+ if (userStep) {
+ lines.push("### User request");
+ lines.push("");
+ lines.push(userStep.content);
+ lines.push("");
+ }
+
+ for (const step of assistantSteps) {
+ lines.push(`### Turn ${step.stepId}`);
+ if (step.content) {
+ lines.push(`Thought: ${step.content}`);
+ }
+ if (step.toolCalls && step.toolCalls.length > 0) {
+ for (let i = 0; i < step.toolCalls.length; i++) {
+ const tc = step.toolCalls[i];
+ const tr = step.toolResults?.[i];
+ const statusAttr = tr?.isError ? "failed" : "succeeded";
+ lines.push(`- ${tc.toolName} (${statusAttr})`);
+ }
+ }
+ lines.push("");
+ }
+
+ return lines.join("\n").trimEnd();
+}
+
+function serializeDag(
+ workflowState: WorkflowState,
+ operatorExecutionResults: Map,
+ useRedact: boolean,
+ compilationResult?: WorkflowCompilationResponse | null
+): string | null {
+ const allOperators = workflowState.getAllOperators();
+ if (allOperators.length === 0) return null;
+
+ const lines: string[] = [];
+
+ const allLinks = workflowState.getAllLinks();
+ const opIds = new Set(allOperators.map(op => op.operatorID));
+ const inDegree = new Map();
+ const children = new Map();
+ for (const id of opIds) {
+ inDegree.set(id, 0);
+ children.set(id, []);
+ }
+ for (const link of allLinks) {
+ children.get(link.source.operatorID)?.push(link.target.operatorID);
+ inDegree.set(link.target.operatorID, (inDegree.get(link.target.operatorID) ?? 0) + 1);
+ }
+ const queue: string[] = [...opIds].filter(id => (inDegree.get(id) ?? 0) === 0);
+ const topoOrder = new Map();
+ let rank = 0;
+ while (queue.length > 0) {
+ const node = queue.shift()!;
+ topoOrder.set(node, rank++);
+ for (const child of children.get(node) ?? []) {
+ const newDeg = (inDegree.get(child) ?? 1) - 1;
+ inDegree.set(child, newDeg);
+ if (newDeg === 0) queue.push(child);
+ }
+ }
+
+ const sortedOps = [...allOperators].sort(
+ (a, b) => (topoOrder.get(a.operatorID) ?? 0) - (topoOrder.get(b.operatorID) ?? 0)
+ );
+
+ const outputSchemas = compilationResult?.operatorOutputSchemas ?? {};
+ const compilationErrors = compilationResult?.operatorErrors ?? {};
+
+ lines.push("## Operators");
+ lines.push("");
+
+ for (const op of sortedOps) {
+ const inputSchemaMap = extractOperatorInputPortSchemaMap(op.operatorID, op, outputSchemas, allLinks);
+ const outputSchemaMap = outputSchemas[op.operatorID];
+ const compilationError = compilationErrors[op.operatorID];
+ lines.push(
+ serializeOperator(
+ op,
+ operatorExecutionResults.get(op.operatorID),
+ useRedact,
+ inputSchemaMap,
+ outputSchemaMap,
+ compilationError
+ )
+ );
+ lines.push("");
+ }
+
+ if (allLinks.length > 0) {
+ const sortedLinks = [...allLinks].sort((a, b) => {
+ const srcA = topoOrder.get(a.source.operatorID) ?? 0;
+ const srcB = topoOrder.get(b.source.operatorID) ?? 0;
+ if (srcA !== srcB) return srcA - srcB;
+ return (topoOrder.get(a.target.operatorID) ?? 0) - (topoOrder.get(b.target.operatorID) ?? 0);
+ });
+
+ lines.push("## Links");
+ for (const link of sortedLinks) {
+ lines.push(`- ${link.source.operatorID} → ${link.target.operatorID}`);
+ }
+ }
+
+ return lines.join("\n").trimEnd();
+}
+
+function serializeOperator(
+ op: OperatorPredicate,
+ execResult: string | undefined,
+ useRedact: boolean,
+ inputSchemaMap?: OperatorPortSchemaMap,
+ outputSchemaMap?: OperatorPortSchemaMap,
+ compilationError?: WorkflowFatalError
+): string {
+ const hasError = execResult !== undefined && execResult.includes("[ERROR]");
+ const status = execResult ? (hasError ? "failed" : "executed") : "not-executed";
+
+ const summary = op.customDisplayName || op.operatorID;
+ const showProperties = !useRedact || hasError;
+
+ const lines: string[] = [];
+ lines.push(`### Operator \`${op.operatorID}\` (${op.operatorType}, ${status})`);
+ lines.push(`Summary: ${summary}`);
+
+ if (inputSchemaMap) {
+ for (const [portId, schema] of Object.entries(inputSchemaMap)) {
+ if (schema) {
+ lines.push(`Input Schema (port ${parsePortIndex(portId)}): ${formatSchema(schema)}`);
+ }
+ }
+ }
+
+ if (showProperties) {
+ const props = op.operatorProperties;
+ if (props && Object.keys(props).length > 0) {
+ lines.push("Properties:");
+ for (const [key, value] of Object.entries(props)) {
+ if (value !== undefined && value !== null && value !== "") {
+ const valueStr = typeof value === "string" ? value : JSON.stringify(value);
+ lines.push(` ${key}: ${valueStr}`);
+ }
+ }
+ }
+ }
+
+ if (outputSchemaMap) {
+ const firstSchema = Object.values(outputSchemaMap).find(s => s !== undefined);
+ if (firstSchema) {
+ lines.push(`Output Schema: ${formatSchema(firstSchema)}`);
+ }
+ }
+
+ if (compilationError) {
+ lines.push(`Compilation Error: ${compilationError.message}`);
+ }
+
+ if (execResult) {
+ lines.push("Result:");
+ const indented = execResult
+ .split("\n")
+ .map(l => " " + l)
+ .join("\n");
+ lines.push(indented);
+ }
+
+ return lines.join("\n");
+}
+
+function formatSchema(schema: PortSchema): string {
+ const attrs = schema.map(a => `${a.attributeName}: ${a.attributeType}`);
+ return `[${attrs.join(", ")}]`;
+}
+
+function parsePortIndex(portId: string): string {
+ const idx = portId.indexOf("_");
+ return idx >= 0 ? portId.substring(0, idx) : portId;
+}
diff --git a/agent-service/src/agent/util/workflow-system-metadata.ts b/agent-service/src/agent/util/workflow-system-metadata.ts
new file mode 100644
index 00000000000..9269a0cff7c
--- /dev/null
+++ b/agent-service/src/agent/util/workflow-system-metadata.ts
@@ -0,0 +1,268 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import Ajv from "ajv";
+import { fetchOperatorMetadata, type OperatorSchema, type OperatorMetadata } from "../../api/backend-api";
+import type { ValidationError, Validation } from "../../types/workflow";
+import { createLogger } from "../../logger";
+
+const log = createLogger("WorkflowSystemMetadata");
+
+export type { ValidationError, Validation } from "../../types/workflow";
+
+interface OperatorSchemaInfo {
+ properties: any;
+ required: any;
+ definitions: any;
+}
+
+interface CompactOperatorSchema {
+ properties: Record;
+ required: string[];
+}
+
+const FILTERED_PROPERTY_KEYS = ["dummyPropertyList"];
+
+const FILTERED_DEFINITION_KEYS = [
+ "DummyProperties",
+ "PortDescription",
+ "HashPartition",
+ "RangePartition",
+ "SinglePartition",
+ "BroadcastPartition",
+ "UnknownPartition",
+];
+
+const COMPACT_SCHEMA_EXCLUDED_KEYS = ["propertyOrder", "autofill", "autofillAttributeOnPort", "attributeTypeRules"];
+
+function filterObjectKeys(obj: any, keysToExclude: string[]): any {
+ if (!obj || typeof obj !== "object") {
+ return obj;
+ }
+ const filtered: any = {};
+ for (const key of Object.keys(obj)) {
+ if (!keysToExclude.includes(key)) {
+ filtered[key] = obj[key];
+ }
+ }
+ return filtered;
+}
+
+function inlineRefs(schema: any, definitions: Record): any {
+ if (!schema || typeof schema !== "object") {
+ return schema;
+ }
+
+ if (schema.$ref && typeof schema.$ref === "string") {
+ const refPath = schema.$ref.replace("#/definitions/", "");
+ const refDef = definitions[refPath];
+ if (refDef) {
+ return inlineRefs(refDef, definitions);
+ }
+ return schema;
+ }
+
+ if (Array.isArray(schema)) {
+ return schema.map(item => inlineRefs(item, definitions));
+ }
+
+ const result: any = {};
+ for (const [key, value] of Object.entries(schema)) {
+ if (COMPACT_SCHEMA_EXCLUDED_KEYS.includes(key)) {
+ continue;
+ }
+ if (typeof value === "object" && value !== null) {
+ result[key] = inlineRefs(value, definitions);
+ } else {
+ result[key] = value;
+ }
+ }
+ return result;
+}
+
+function getCompactSchema(jsonSchema: any): CompactOperatorSchema | null {
+ try {
+ const properties = filterObjectKeys(jsonSchema.properties, FILTERED_PROPERTY_KEYS);
+ const definitions = filterObjectKeys(jsonSchema.definitions, FILTERED_DEFINITION_KEYS) || {};
+
+ const compactProperties: Record = {};
+ for (const [propName, propSchema] of Object.entries(properties || {})) {
+ compactProperties[propName] = inlineRefs(propSchema, definitions);
+ }
+
+ return {
+ properties: compactProperties,
+ required: jsonSchema.required || [],
+ };
+ } catch {
+ return null;
+ }
+}
+
+// Matches the frontend ValidationWorkflowService Ajv configuration.
+const ajv = new Ajv({ allErrors: true, strict: false });
+
+/**
+ * Process-wide singleton cache of operator metadata fetched from the backend.
+ *
+ * Holds each operator type's JSON schema, description, and additional
+ * metadata, plus a compact schema variant used in system prompts and error
+ * messages. Exposes Ajv-backed property validation that matches the
+ * frontend's `ValidationWorkflowService` configuration.
+ */
+export class WorkflowSystemMetadata {
+ private static instance: WorkflowSystemMetadata | null = null;
+
+ static getInstance(): WorkflowSystemMetadata {
+ if (!WorkflowSystemMetadata.instance) {
+ WorkflowSystemMetadata.instance = new WorkflowSystemMetadata();
+ }
+ return WorkflowSystemMetadata.instance;
+ }
+
+ static async initializeGlobal(): Promise {
+ const instance = WorkflowSystemMetadata.getInstance();
+ if (!instance.isInitialized()) {
+ await instance.initializeFromBackend();
+ }
+ return instance;
+ }
+
+ private schemas: Map = new Map();
+ private descriptions: Map = new Map();
+ private additionalMetadata: Map = new Map();
+ private initialized = false;
+
+ async initializeFromBackend(): Promise {
+ try {
+ const metadata = await fetchOperatorMetadata();
+ this.loadFromMetadata(metadata);
+ this.initialized = true;
+ log.info({ operatorCount: this.schemas.size }, "loaded operators from backend");
+ } catch (error) {
+ log.warn({ err: error }, "failed to fetch from backend");
+ throw error;
+ }
+ }
+
+ loadFromMetadata(metadata: OperatorMetadata): void {
+ for (const op of metadata.operators) {
+ this.schemas.set(op.operatorType, op.jsonSchema);
+ this.descriptions.set(
+ op.operatorType,
+ op.additionalMetadata.operatorDescription || op.additionalMetadata.userFriendlyName
+ );
+ this.additionalMetadata.set(op.operatorType, op.additionalMetadata);
+ }
+ }
+
+ isInitialized(): boolean {
+ return this.initialized;
+ }
+
+ getSchema(operatorType: string): any | undefined {
+ return this.schemas.get(operatorType);
+ }
+
+ getDescription(operatorType: string): string {
+ return this.descriptions.get(operatorType) || "";
+ }
+
+ getAdditionalMetadata(operatorType: string): any | undefined {
+ return this.additionalMetadata.get(operatorType);
+ }
+
+ getAllOperatorTypes(): Record {
+ const result: Record = {};
+ for (const [type, desc] of this.descriptions) {
+ result[type] = desc;
+ }
+ return result;
+ }
+
+ getCompactSchema(operatorType: string): CompactOperatorSchema | null {
+ const schema = this.schemas.get(operatorType);
+ if (!schema) return null;
+ return getCompactSchema(schema);
+ }
+
+ getAllSchemasAsJson(): string {
+ const result: Record = {};
+ for (const [type, schema] of this.schemas) {
+ result[type] = {
+ properties: filterObjectKeys(schema.properties, FILTERED_PROPERTY_KEYS),
+ required: schema.required,
+ definitions: filterObjectKeys(schema.definitions, FILTERED_DEFINITION_KEYS),
+ };
+ }
+ return JSON.stringify(result, null, 2);
+ }
+
+ getOperatorCount(): number {
+ return this.schemas.size;
+ }
+
+ operatorTypeExists(operatorType: string): boolean {
+ return this.schemas.has(operatorType);
+ }
+
+ validateOperatorProperties(operatorType: string, properties: Record): Validation {
+ const schema = this.schemas.get(operatorType);
+ if (!schema) {
+ return { isValid: false, messages: { error: `Unknown operator type: ${operatorType}` } };
+ }
+
+ try {
+ const isValid = ajv.validate(schema, properties);
+
+ if (isValid) {
+ return { isValid: true };
+ }
+
+ const messages: Record = {};
+ if (ajv.errors) {
+ for (const error of ajv.errors) {
+ const key = error.instancePath
+ ? error.instancePath.replace(/^\//, "").replace(/\//g, ".")
+ : (error.params as any)?.missingProperty || error.keyword;
+ messages[key] = error.message || "Validation failed";
+ }
+ }
+ return { isValid: false, messages };
+ } catch (e) {
+ return { isValid: false, messages: { error: `Validation error: ${e}` } };
+ }
+ }
+}
+
+export function formatValidationErrors(validation: Validation): string {
+ if (validation.isValid) return "";
+ const errorMessages = Object.entries(validation.messages).map(([key, msg]) => `${key}: ${msg}`);
+ return errorMessages.join("; ");
+}
+
+export function formatCompactSchemaForError(compactSchema: CompactOperatorSchema): string {
+ const requiredProps: Record = {};
+ for (const key of compactSchema.required) {
+ if (compactSchema.properties[key]) {
+ requiredProps[key] = compactSchema.properties[key];
+ }
+ }
+ return `required: [${compactSchema.required.join(", ")}], properties: ${JSON.stringify(requiredProps)}`;
+}
diff --git a/agent-service/src/agent/util/workflow-utils.ts b/agent-service/src/agent/util/workflow-utils.ts
new file mode 100644
index 00000000000..d723f61eff0
--- /dev/null
+++ b/agent-service/src/agent/util/workflow-utils.ts
@@ -0,0 +1,209 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import Ajv from "ajv";
+import type {
+ OperatorPredicate,
+ PortDescription,
+ OperatorLink,
+ PortSchema,
+ OperatorPortSchemaMap,
+} from "../../types/workflow";
+import type { WorkflowSystemMetadata } from "./workflow-system-metadata";
+import type { WorkflowState } from "../workflow-state";
+
+// Format "{id}_{internal}" must align with the backend port-identity serializer.
+function serializePortIdentity(id: number, internal: boolean = false): string {
+ return `${id}_${internal}`;
+}
+
+function parseLogicalOperatorPortID(portId: string): { portNumber: number; portType: "input" | "output" } | undefined {
+ const match = portId.match(/^(input|output)-(\d+)$/);
+ if (!match) {
+ return undefined;
+ }
+
+ const portType = match[1] as "input" | "output";
+ const portNumber = parseInt(match[2]);
+
+ return { portNumber, portType };
+}
+
+function getInputLinksByOperatorId(operatorId: string, links: OperatorLink[]): OperatorLink[] {
+ return links.filter(link => link.target.operatorID === operatorId);
+}
+
+export function extractOperatorInputPortSchemaMap(
+ operatorId: string,
+ operator: OperatorPredicate,
+ outputSchemas: Record,
+ links: OperatorLink[]
+): OperatorPortSchemaMap | undefined {
+ const inputLinks = getInputLinksByOperatorId(operatorId, links);
+ if (!inputLinks.length) return undefined;
+
+ const inputPortSchemaMap: Record = {};
+
+ operator.inputPorts.forEach((_, portIndex) => {
+ const portId = serializePortIdentity(portIndex, false);
+ inputPortSchemaMap[portId] = undefined;
+
+ const linksToThisPort = inputLinks.filter(link => {
+ const parsedPort = parseLogicalOperatorPortID(link.target.portID);
+ if (!parsedPort) return false;
+ return parsedPort.portNumber === portIndex;
+ });
+
+ if (linksToThisPort.length > 0) {
+ const schemas: (PortSchema | undefined)[] = linksToThisPort.map(link => {
+ const sourcePortSchemaMap = outputSchemas[link.source.operatorID];
+ if (!sourcePortSchemaMap) {
+ return undefined;
+ }
+
+ const outputPort = parseLogicalOperatorPortID(link.source.portID);
+ if (!outputPort) {
+ return undefined;
+ }
+
+ return sourcePortSchemaMap[serializePortIdentity(outputPort.portNumber, false)];
+ });
+
+ // Unlike the frontend, we don't flag mismatched schemas as a compilation
+ // error; we just pick the first defined one.
+ if (schemas.length > 0) {
+ inputPortSchemaMap[portId] = schemas.find(s => s !== undefined);
+ }
+ }
+ });
+
+ const hasAnySchema = Object.values(inputPortSchemaMap).some(s => s !== undefined);
+ return hasAnySchema ? inputPortSchemaMap : undefined;
+}
+
+interface InputPortInfo {
+ displayName?: string;
+ disallowMultiLinks?: boolean;
+ dependencies?: { id: number; internal: boolean }[];
+}
+
+interface OutputPortInfo {
+ displayName?: string;
+}
+
+function inputPortToPortDescription(portID: string, inputPortInfo: InputPortInfo): PortDescription {
+ return {
+ portID,
+ displayName: inputPortInfo.displayName ?? "",
+ disallowMultiInputs: inputPortInfo.disallowMultiLinks ?? false,
+ isDynamicPort: false,
+ dependencies: inputPortInfo.dependencies ?? [],
+ };
+}
+
+function outputPortToPortDescription(portID: string, outputPortInfo: OutputPortInfo): PortDescription {
+ return {
+ portID,
+ displayName: outputPortInfo.displayName ?? "",
+ disallowMultiInputs: false,
+ isDynamicPort: false,
+ };
+}
+
+/**
+ * Builds new `OperatorPredicate` instances from operator metadata.
+ *
+ * Given an operator type, reads the JSON schema and additional metadata from
+ * `WorkflowSystemMetadata`, materializes default properties via Ajv, and
+ * synthesizes input/output port descriptions so the operator is ready to
+ * drop into a `WorkflowState`.
+ */
+export class WorkflowUtilService {
+ private metadataStore: WorkflowSystemMetadata;
+ private workflowState: WorkflowState;
+ private ajv: Ajv;
+
+ constructor(metadataStore: WorkflowSystemMetadata, workflowState: WorkflowState) {
+ this.metadataStore = metadataStore;
+ this.workflowState = workflowState;
+ this.ajv = new Ajv({ useDefaults: true, strict: false });
+ }
+
+ public getNewOperatorPredicate(operatorType: string, customDisplayName?: string): OperatorPredicate {
+ const jsonSchema = this.metadataStore.getSchema(operatorType);
+ const additionalMetadata = this.metadataStore.getAdditionalMetadata(operatorType);
+
+ if (!jsonSchema || !additionalMetadata) {
+ throw new Error(`operatorType ${operatorType} doesn't exist in operator metadata`);
+ }
+
+ const operatorId = this.workflowState.generateOperatorId(operatorType);
+ const operatorProperties: Record = {};
+
+ // Strip $id so Ajv doesn't warn about a duplicate schema registration.
+ const { $id, ...schemaWithoutId } = jsonSchema as any;
+
+ // Calling validate() here populates operatorProperties with a deep clone of
+ // the schema defaults via Ajv's useDefaults option.
+ const validate = this.ajv.compile(schemaWithoutId);
+ validate(operatorProperties);
+
+ const inputPorts: PortDescription[] = [];
+ const outputPorts: PortDescription[] = [];
+
+ const showAdvanced = false;
+
+ const isDisabled = false;
+
+ const displayName = customDisplayName ?? additionalMetadata.userFriendlyName;
+
+ const dynamicInputPorts = additionalMetadata.dynamicInputPorts ?? false;
+ const dynamicOutputPorts = additionalMetadata.dynamicOutputPorts ?? false;
+
+ const inputPortInfos = additionalMetadata.inputPorts || [];
+ for (let i = 0; i < inputPortInfos.length; i++) {
+ const portID = "input-" + i.toString();
+ const portInfo = inputPortInfos[i] as InputPortInfo;
+ inputPorts.push(inputPortToPortDescription(portID, portInfo));
+ }
+
+ const outputPortInfos = additionalMetadata.outputPorts || [];
+ for (let i = 0; i < outputPortInfos.length; i++) {
+ const portID = "output-" + i.toString();
+ const portInfo = outputPortInfos[i] as OutputPortInfo;
+ outputPorts.push(outputPortToPortDescription(portID, portInfo));
+ }
+
+ const operatorVersion = (additionalMetadata as any).operatorVersion ?? "N/A";
+
+ return {
+ operatorID: operatorId,
+ operatorType,
+ operatorVersion,
+ operatorProperties,
+ inputPorts,
+ outputPorts,
+ showAdvanced,
+ isDisabled,
+ customDisplayName: displayName,
+ dynamicInputPorts,
+ dynamicOutputPorts,
+ };
+ }
+}
diff --git a/agent-service/src/agent/workflow-result-state.test.ts b/agent-service/src/agent/workflow-result-state.test.ts
new file mode 100644
index 00000000000..b2e46fd0d9e
--- /dev/null
+++ b/agent-service/src/agent/workflow-result-state.test.ts
@@ -0,0 +1,95 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { describe, expect, test } from "bun:test";
+import { WorkflowResultState } from "./workflow-result-state";
+import type { OperatorInfo } from "../types/execution";
+
+function makeInfo(outputTuples: number): OperatorInfo {
+ return {
+ state: "Completed",
+ inputTuples: 0,
+ outputTuples,
+ resultMode: "table",
+ };
+}
+
+describe("WorkflowResultState - ancestor walk", () => {
+ test("returns the most recent ancestor entry", () => {
+ let path: string[] = [];
+ const state = new WorkflowResultState(() => path);
+
+ state.set("op1", "step-A", makeInfo(1));
+ state.set("op1", "step-B", makeInfo(2));
+ state.set("op1", "step-C", makeInfo(3));
+
+ path = ["step-A", "step-B", "step-C"];
+ expect(state.get("op1")?.operatorInfo.outputTuples).toBe(3);
+
+ // Rewind to step-B; step-C is no longer an ancestor.
+ path = ["step-A", "step-B"];
+ expect(state.get("op1")?.operatorInfo.outputTuples).toBe(2);
+
+ // Rewind further.
+ path = ["step-A"];
+ expect(state.get("op1")?.operatorInfo.outputTuples).toBe(1);
+ });
+
+ test("returns undefined when no ancestor has a result", () => {
+ const state = new WorkflowResultState(() => ["step-X"]);
+ state.set("op1", "step-A", makeInfo(1));
+ expect(state.get("op1")).toBeUndefined();
+ });
+
+ test("returns undefined for unknown operator", () => {
+ const state = new WorkflowResultState(() => ["step-A"]);
+ expect(state.get("missing")).toBeUndefined();
+ });
+
+ test("getAllVisible returns one entry per operator on the current branch", () => {
+ let path: string[] = [];
+ const state = new WorkflowResultState(() => path);
+
+ // op1 has results on step-A and step-C; the branch only goes through A and B.
+ state.set("op1", "step-A", makeInfo(1));
+ state.set("op1", "step-C", makeInfo(99));
+ state.set("op2", "step-B", makeInfo(7));
+
+ path = ["step-A", "step-B"];
+ const visible = state.getAllVisible();
+ expect(visible.size).toBe(2);
+ expect(visible.get("op1")?.operatorInfo.outputTuples).toBe(1);
+ expect(visible.get("op2")?.operatorInfo.outputTuples).toBe(7);
+ });
+
+ test("clear drops all stored results", () => {
+ const state = new WorkflowResultState(() => ["step-A"]);
+ state.set("op1", "step-A", makeInfo(1));
+ state.clear();
+ expect(state.get("op1")).toBeUndefined();
+ expect(state.getAllVisible().size).toBe(0);
+ });
+
+ test("set on the same step overwrites", () => {
+ const state = new WorkflowResultState(() => ["step-A"]);
+ state.set("op1", "step-A", makeInfo(1));
+ state.set("op1", "step-A", makeInfo(42));
+ expect(state.get("op1")?.operatorInfo.outputTuples).toBe(42);
+ });
+});
diff --git a/agent-service/src/agent/workflow-result-state.ts b/agent-service/src/agent/workflow-result-state.ts
new file mode 100644
index 00000000000..e6f13c2301a
--- /dev/null
+++ b/agent-service/src/agent/workflow-result-state.ts
@@ -0,0 +1,83 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import type { OperatorInfo } from "../types/execution";
+
+interface ResultEntry {
+ operatorInfo: OperatorInfo;
+ stepId: string;
+}
+
+/**
+ * Versioned per-operator execution results keyed by step id.
+ *
+ * Each operator can have multiple result snapshots (one per step that
+ * executed it). Lookups walk the current ancestor path from HEAD and return
+ * the most recent result visible on that branch, so checking out an earlier
+ * step exposes the results that were live at that point.
+ */
+export class WorkflowResultState {
+ private results = new Map>();
+
+ constructor(private getAncestorPath: () => string[]) {}
+
+ set(operatorId: string, stepId: string, operatorInfo: OperatorInfo): void {
+ let versions = this.results.get(operatorId);
+ if (!versions) {
+ versions = new Map();
+ this.results.set(operatorId, versions);
+ }
+ versions.set(stepId, { operatorInfo, stepId });
+ }
+
+ get(operatorId: string): ResultEntry | undefined {
+ const versions = this.results.get(operatorId);
+ if (!versions) return undefined;
+
+ const path = this.getAncestorPath();
+ for (let i = path.length - 1; i >= 0; i--) {
+ const entry = versions.get(path[i]);
+ if (entry) return entry;
+ }
+ return undefined;
+ }
+
+ getOperatorInfo(operatorId: string): OperatorInfo | undefined {
+ return this.get(operatorId)?.operatorInfo;
+ }
+
+ getAllVisible(): Map {
+ const result = new Map();
+ const path = this.getAncestorPath();
+
+ for (const [operatorId, versions] of this.results) {
+ for (let i = path.length - 1; i >= 0; i--) {
+ if (versions.has(path[i])) {
+ result.set(operatorId, versions.get(path[i])!);
+ break;
+ }
+ }
+ }
+ return result;
+ }
+
+ clear(): void {
+ this.results.clear();
+ }
+}
diff --git a/agent-service/src/agent/workflow-state.test.ts b/agent-service/src/agent/workflow-state.test.ts
new file mode 100644
index 00000000000..bffc3769eb7
--- /dev/null
+++ b/agent-service/src/agent/workflow-state.test.ts
@@ -0,0 +1,176 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { describe, expect, test } from "bun:test";
+import { WorkflowState } from "./workflow-state";
+import type { OperatorPredicate, OperatorLink } from "../types/workflow";
+
+function makeOperator(id: string, overrides: Partial = {}): OperatorPredicate {
+ return {
+ operatorID: id,
+ operatorType: "TestOp",
+ operatorVersion: "1.0",
+ operatorProperties: {},
+ inputPorts: [{ portID: "input-0", displayName: "Input 0" }],
+ outputPorts: [{ portID: "output-0", displayName: "Output 0" }],
+ showAdvanced: false,
+ ...overrides,
+ };
+}
+
+function makeLink(linkId: string, sourceId: string, targetId: string): OperatorLink {
+ return {
+ linkID: linkId,
+ source: { operatorID: sourceId, portID: "output-0" },
+ target: { operatorID: targetId, portID: "input-0" },
+ };
+}
+
+describe("WorkflowState - operators", () => {
+ test("add and get operator round-trips", () => {
+ const state = new WorkflowState();
+ const op = makeOperator("op1");
+ state.addOperator(op);
+ expect(state.getOperator("op1")).toEqual(op);
+ expect(state.getAllOperators()).toHaveLength(1);
+ });
+
+ test("delete operator removes connected links", () => {
+ const state = new WorkflowState();
+ state.addOperator(makeOperator("op1"));
+ state.addOperator(makeOperator("op2"));
+ state.addLink(makeLink("l1", "op1", "op2"));
+
+ expect(state.deleteOperator("op1")).toBe(true);
+ expect(state.getOperator("op1")).toBeUndefined();
+ expect(state.getAllLinks()).toHaveLength(0);
+ });
+
+ test("delete on missing operator returns false", () => {
+ const state = new WorkflowState();
+ expect(state.deleteOperator("missing")).toBe(false);
+ });
+
+ test("updateOperatorProperties merges, does not replace", () => {
+ const state = new WorkflowState();
+ state.addOperator(makeOperator("op1", { operatorProperties: { a: 1, b: 2 } }));
+ state.updateOperatorProperties("op1", { b: 99, c: 3 });
+
+ expect(state.getOperator("op1")?.operatorProperties).toEqual({ a: 1, b: 99, c: 3 });
+ });
+
+ test("updateOperatorDisplayName sets customDisplayName", () => {
+ const state = new WorkflowState();
+ state.addOperator(makeOperator("op1"));
+ expect(state.updateOperatorDisplayName("op1", "Filter rows")).toBe(true);
+ expect(state.getOperator("op1")?.customDisplayName).toBe("Filter rows");
+ });
+
+ test("update on missing operator returns false", () => {
+ const state = new WorkflowState();
+ expect(state.updateOperatorProperties("missing", { a: 1 })).toBe(false);
+ expect(state.updateOperatorDisplayName("missing", "x")).toBe(false);
+ });
+});
+
+describe("WorkflowState - links", () => {
+ test("add, get, and delete link", () => {
+ const state = new WorkflowState();
+ state.addOperator(makeOperator("op1"));
+ state.addOperator(makeOperator("op2"));
+ const link = makeLink("l1", "op1", "op2");
+ state.addLink(link);
+
+ expect(state.getLink("l1")).toEqual(link);
+ expect(state.deleteLink("l1")).toBe(true);
+ expect(state.getLink("l1")).toBeUndefined();
+ });
+
+ test("getLinksConnectedToOperator returns both inbound and outbound", () => {
+ const state = new WorkflowState();
+ state.addOperator(makeOperator("op1"));
+ state.addOperator(makeOperator("op2"));
+ state.addOperator(makeOperator("op3"));
+ state.addLink(makeLink("l1", "op1", "op2"));
+ state.addLink(makeLink("l2", "op2", "op3"));
+
+ const connected = state.getLinksConnectedToOperator("op2");
+ expect(connected.map(l => l.linkID).sort()).toEqual(["l1", "l2"]);
+ });
+});
+
+describe("WorkflowState - generated ids", () => {
+ test("generateLinkId is monotonically increasing", () => {
+ const state = new WorkflowState();
+ expect(state.generateLinkId()).toBe("link-1");
+ expect(state.generateLinkId()).toBe("link-2");
+ expect(state.generateLinkId()).toBe("link-3");
+ });
+
+ test("generateOperatorId is namespaced by type", () => {
+ const state = new WorkflowState();
+ expect(state.generateOperatorId("Filter")).toBe("Filter-operator-1");
+ expect(state.generateOperatorId("Filter")).toBe("Filter-operator-2");
+ expect(state.generateOperatorId("Sort")).toBe("Sort-operator-3");
+ });
+});
+
+describe("WorkflowState - getSubDAG", () => {
+ test("walks ancestors of the target operator", () => {
+ // op1 -> op2 -> op4
+ // op3 -> op4
+ // sub-DAG of op4 should include all four.
+ const state = new WorkflowState();
+ state.addOperator(makeOperator("op1"));
+ state.addOperator(makeOperator("op2"));
+ state.addOperator(makeOperator("op3"));
+ state.addOperator(makeOperator("op4"));
+ state.addLink(makeLink("l1", "op1", "op2"));
+ state.addLink(makeLink("l2", "op2", "op4"));
+ state.addLink(makeLink("l3", "op3", "op4"));
+
+ const subDag = state.getSubDAG("op4");
+ expect(subDag.operators.map(o => o.operatorID).sort()).toEqual(["op1", "op2", "op3", "op4"]);
+ expect(subDag.links.map(l => l.linkID).sort()).toEqual(["l1", "l2", "l3"]);
+ });
+
+ test("excludes downstream operators", () => {
+ // op1 -> op2 -> op3
+ // sub-DAG of op2 should include op1 and op2 but not op3.
+ const state = new WorkflowState();
+ state.addOperator(makeOperator("op1"));
+ state.addOperator(makeOperator("op2"));
+ state.addOperator(makeOperator("op3"));
+ state.addLink(makeLink("l1", "op1", "op2"));
+ state.addLink(makeLink("l2", "op2", "op3"));
+
+ const subDag = state.getSubDAG("op2");
+ expect(subDag.operators.map(o => o.operatorID).sort()).toEqual(["op1", "op2"]);
+ });
+
+ test("disabled upstream operators are skipped", () => {
+ const state = new WorkflowState();
+ state.addOperator(makeOperator("op1", { isDisabled: true }));
+ state.addOperator(makeOperator("op2"));
+ state.addLink(makeLink("l1", "op1", "op2"));
+
+ const subDag = state.getSubDAG("op2");
+ expect(subDag.operators.map(o => o.operatorID)).toEqual(["op2"]);
+ });
+});
diff --git a/agent-service/src/agent/workflow-state.ts b/agent-service/src/agent/workflow-state.ts
new file mode 100644
index 00000000000..04ad2b0e4e8
--- /dev/null
+++ b/agent-service/src/agent/workflow-state.ts
@@ -0,0 +1,488 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { Subject, Observable, merge, Subscription } from "rxjs";
+import type {
+ OperatorPredicate,
+ OperatorLink,
+ WorkflowContent,
+ LogicalPlan,
+ LogicalOperator,
+ LogicalLink,
+ Point,
+ CommentBox,
+ WorkflowSettings,
+ ValidationError,
+} from "../types/workflow";
+
+export type { ValidationError, Validation } from "../types/workflow";
+
+interface ValidationOutput {
+ errors: Record;
+ workflowEmpty: boolean;
+}
+
+const DEFAULT_WORKFLOW_SETTINGS: WorkflowSettings = {
+ dataTransferBatchSize: 400,
+};
+
+/**
+ * In-memory logical plan the agent edits across a conversation.
+ *
+ * Holds operators, links, positions, comment boxes, and workflow settings,
+ * and emits change events via RxJS so subscribers (auto-persist, websocket
+ * broadcast) can react. Converts to/from the backend `WorkflowContent` wire
+ * format and to the `LogicalPlan` shape used for compilation and execution.
+ */
+export class WorkflowState {
+ private operators: Map = new Map();
+ private links: Map = new Map();
+ private operatorPositions: Map = new Map();
+ private commentBoxes: CommentBox[] = [];
+ private settings: WorkflowSettings = { ...DEFAULT_WORKFLOW_SETTINGS };
+ private operatorsToViewResult: Set = new Set();
+
+ private operatorIdCounter: number = 0;
+ private linkIdCounter: number = 0;
+
+ private readonly operatorAddSubject = new Subject();
+ private readonly operatorDeleteSubject = new Subject<{ deletedOperatorID: string }>();
+ private readonly operatorPropertyChangeSubject = new Subject<{ operator: OperatorPredicate }>();
+ private readonly linkAddSubject = new Subject();
+ private readonly linkDeleteSubject = new Subject<{ deletedLink: OperatorLink }>();
+ private readonly disabledOperatorChangedSubject = new Subject<{
+ newDisabled: string[];
+ newEnabled: string[];
+ }>();
+ private readonly viewResultOperatorChangedSubject = new Subject<{
+ newViewResultOps: string[];
+ newUnviewResultOps: string[];
+ }>();
+
+ private validationErrors: Record = {};
+ private workflowEmpty: boolean = true;
+
+ private readonly validationChangedSubject = new Subject();
+
+ private subscriptions: Subscription[] = [];
+
+ getWorkflowChangedStream(): Observable {
+ return merge(
+ this.operatorAddSubject,
+ this.operatorDeleteSubject,
+ this.operatorPropertyChangeSubject,
+ this.linkAddSubject,
+ this.linkDeleteSubject,
+ this.disabledOperatorChangedSubject
+ );
+ }
+
+ generateOperatorId(operatorType: string): string {
+ return `${operatorType}-operator-${++this.operatorIdCounter}`;
+ }
+
+ generateLinkId(): string {
+ return `link-${++this.linkIdCounter}`;
+ }
+
+ addOperator(operator: OperatorPredicate, position?: Point): void {
+ this.operators.set(operator.operatorID, operator);
+ const defaultPosition: Point = position || {
+ x: 100 + (this.operators.size - 1) * 200,
+ y: 100 + (this.operators.size - 1) * 100,
+ };
+ this.operatorPositions.set(operator.operatorID, defaultPosition);
+ this.operatorAddSubject.next(operator);
+ }
+
+ getOperator(operatorId: string): OperatorPredicate | undefined {
+ return this.operators.get(operatorId);
+ }
+
+ getAllOperators(): OperatorPredicate[] {
+ return Array.from(this.operators.values());
+ }
+
+ getAllEnabledOperators(): OperatorPredicate[] {
+ return this.getAllOperators();
+ }
+
+ deleteOperator(operatorId: string): boolean {
+ const operator = this.operators.get(operatorId);
+ if (!operator) return false;
+
+ const linksToDelete = this.getLinksConnectedToOperator(operatorId);
+ for (const link of linksToDelete) {
+ this.links.delete(link.linkID);
+ this.linkDeleteSubject.next({ deletedLink: link });
+ }
+
+ this.operatorsToViewResult.delete(operatorId);
+ this.operatorPositions.delete(operatorId);
+ const deleted = this.operators.delete(operatorId);
+
+ if (deleted) {
+ this.operatorDeleteSubject.next({ deletedOperatorID: operatorId });
+ }
+
+ return deleted;
+ }
+
+ updateOperatorProperties(operatorId: string, properties: Record): boolean {
+ const operator = this.operators.get(operatorId);
+ if (!operator) return false;
+
+ const updatedOperator: OperatorPredicate = {
+ ...operator,
+ operatorProperties: { ...operator.operatorProperties, ...properties },
+ };
+ this.operators.set(operatorId, updatedOperator);
+ this.operatorPropertyChangeSubject.next({ operator: updatedOperator });
+ return true;
+ }
+
+ updateOperatorDisplayName(operatorId: string, displayName: string): boolean {
+ const operator = this.operators.get(operatorId);
+ if (!operator) return false;
+
+ const updatedOperator: OperatorPredicate = {
+ ...operator,
+ customDisplayName: displayName,
+ };
+ this.operators.set(operatorId, updatedOperator);
+ this.operatorPropertyChangeSubject.next({ operator: updatedOperator });
+ return true;
+ }
+
+ updateOperatorInputPorts(operatorId: string, numInputPorts: number): boolean {
+ const operator = this.operators.get(operatorId);
+ if (!operator) return false;
+
+ const newInputPorts: import("../types/workflow").PortDescription[] = [];
+ for (let i = 0; i < numInputPorts; i++) {
+ newInputPorts.push({
+ portID: `input-${i}`,
+ displayName: `Input ${i}`,
+ disallowMultiInputs: true,
+ isDynamicPort: i > 0,
+ });
+ }
+
+ const updatedOperator: OperatorPredicate = {
+ ...operator,
+ inputPorts: newInputPorts,
+ };
+ this.operators.set(operatorId, updatedOperator);
+ this.operatorPropertyChangeSubject.next({ operator: updatedOperator });
+ return true;
+ }
+
+ updateOperatorPosition(operatorId: string, position: Point): boolean {
+ if (!this.operators.has(operatorId)) {
+ return false;
+ }
+ this.operatorPositions.set(operatorId, position);
+ return true;
+ }
+
+ getOperatorPosition(operatorId: string): Point | undefined {
+ return this.operatorPositions.get(operatorId);
+ }
+
+ addLink(link: OperatorLink): void {
+ this.links.set(link.linkID, link);
+ this.linkAddSubject.next(link);
+ }
+
+ getLink(linkId: string): OperatorLink | undefined {
+ return this.links.get(linkId);
+ }
+
+ getAllLinks(): OperatorLink[] {
+ return Array.from(this.links.values());
+ }
+
+ deleteLink(linkId: string): boolean {
+ const link = this.links.get(linkId);
+ if (!link) return false;
+
+ const deleted = this.links.delete(linkId);
+ if (deleted) {
+ this.linkDeleteSubject.next({ deletedLink: link });
+ }
+ return deleted;
+ }
+
+ getLinksConnectedToOperator(operatorId: string): OperatorLink[] {
+ return this.getAllLinks().filter(
+ link => link.source.operatorID === operatorId || link.target.operatorID === operatorId
+ );
+ }
+
+ getSubDAG(targetOperatorId: string): { operators: OperatorPredicate[]; links: OperatorLink[] } {
+ const visited = new Set();
+ const subDagOperators: OperatorPredicate[] = [];
+ const subDagLinks: OperatorLink[] = [];
+
+ const dfs = (currentOperatorId: string) => {
+ if (visited.has(currentOperatorId)) {
+ return;
+ }
+
+ visited.add(currentOperatorId);
+
+ const currentOperator = this.getOperator(currentOperatorId);
+ if (currentOperator && !currentOperator.isDisabled) {
+ subDagOperators.push(currentOperator);
+
+ const connectedLinks = this.getAllLinks().filter(
+ link => link.target.operatorID === currentOperatorId && !this.getOperator(link.source.operatorID)?.isDisabled
+ );
+
+ connectedLinks.forEach(link => {
+ subDagLinks.push(link);
+ dfs(link.source.operatorID);
+ });
+ }
+ };
+
+ dfs(targetOperatorId);
+
+ return { operators: subDagOperators, links: subDagLinks };
+ }
+
+ getFrontierOperators(depth: number): string[] {
+ const allOperators = this.getAllOperators();
+ if (allOperators.length === 0) return [];
+
+ const sourceOperatorIds = new Set();
+ for (const link of this.getAllLinks()) {
+ sourceOperatorIds.add(link.source.operatorID);
+ }
+
+ const leaves = allOperators.filter(op => !sourceOperatorIds.has(op.operatorID)).map(op => op.operatorID);
+
+ if (leaves.length === 0) {
+ return allOperators.map(op => op.operatorID);
+ }
+
+ const frontier = new Set(leaves);
+ let currentLevel = new Set(leaves);
+
+ for (let d = 1; d < depth; d++) {
+ const nextLevel = new Set();
+ for (const opId of currentLevel) {
+ for (const link of this.getAllLinks()) {
+ if (link.target.operatorID === opId && !frontier.has(link.source.operatorID)) {
+ nextLevel.add(link.source.operatorID);
+ frontier.add(link.source.operatorID);
+ }
+ }
+ }
+ if (nextLevel.size === 0) break;
+ currentLevel = nextLevel;
+ }
+
+ const frontierArray = Array.from(frontier);
+ const inDegree = new Map();
+ const children = new Map();
+ for (const opId of frontierArray) {
+ inDegree.set(opId, 0);
+ children.set(opId, []);
+ }
+ for (const link of this.getAllLinks()) {
+ if (frontier.has(link.source.operatorID) && frontier.has(link.target.operatorID)) {
+ children.get(link.source.operatorID)!.push(link.target.operatorID);
+ inDegree.set(link.target.operatorID, (inDegree.get(link.target.operatorID) ?? 0) + 1);
+ }
+ }
+
+ const queue: string[] = frontierArray.filter(opId => (inDegree.get(opId) ?? 0) === 0);
+ const sorted: string[] = [];
+ while (queue.length > 0) {
+ const node = queue.shift()!;
+ sorted.push(node);
+ for (const child of children.get(node) ?? []) {
+ const newDeg = (inDegree.get(child) ?? 1) - 1;
+ inDegree.set(child, newDeg);
+ if (newDeg === 0) queue.push(child);
+ }
+ }
+
+ if (sorted.length < frontierArray.length) {
+ for (const opId of frontierArray) {
+ if (!sorted.includes(opId)) sorted.push(opId);
+ }
+ }
+
+ return sorted;
+ }
+
+ getValidationChangedStream(): Observable {
+ return this.validationChangedSubject.asObservable();
+ }
+
+ getValidationOutput(): ValidationOutput {
+ return {
+ errors: { ...this.validationErrors },
+ workflowEmpty: this.workflowEmpty,
+ };
+ }
+
+ setValidationError(operatorId: string, error: ValidationError): void {
+ this.validationErrors[operatorId] = error;
+ this.emitValidationChanged();
+ }
+
+ clearValidationError(operatorId: string): void {
+ delete this.validationErrors[operatorId];
+ this.emitValidationChanged();
+ }
+
+ setAllValidationErrors(errors: Record): void {
+ this.validationErrors = { ...errors };
+ this.updateWorkflowEmptyState();
+ this.emitValidationChanged();
+ }
+
+ private updateWorkflowEmptyState(): void {
+ const operators = this.getAllOperators();
+ this.workflowEmpty = operators.length === 0;
+
+ if (!this.workflowEmpty) {
+ this.workflowEmpty = operators.every(op => op.isDisabled);
+ }
+ }
+
+ private emitValidationChanged(): void {
+ this.validationChangedSubject.next({
+ errors: { ...this.validationErrors },
+ workflowEmpty: this.workflowEmpty,
+ });
+ }
+
+ getWorkflowContent(): WorkflowContent {
+ const positionsObj: { [key: string]: Point } = {};
+ for (const [id, pos] of this.operatorPositions) {
+ positionsObj[id] = pos;
+ }
+
+ return {
+ operators: this.getAllOperators(),
+ operatorPositions: positionsObj,
+ links: this.getAllLinks(),
+ commentBoxes: [...this.commentBoxes],
+ settings: { ...this.settings },
+ };
+ }
+
+ setWorkflowContent(content: WorkflowContent): void {
+ this.operators.clear();
+ this.links.clear();
+ this.operatorPositions.clear();
+
+ for (const op of content.operators) {
+ this.operators.set(op.operatorID, op);
+ }
+ for (const link of content.links) {
+ this.links.set(link.linkID, link);
+ }
+
+ if (content.operatorPositions) {
+ for (const [id, pos] of Object.entries(content.operatorPositions)) {
+ this.operatorPositions.set(id, pos);
+ }
+ }
+
+ this.commentBoxes = content.commentBoxes ? [...content.commentBoxes] : [];
+
+ this.settings = content.settings ? { ...content.settings } : { ...DEFAULT_WORKFLOW_SETTINGS };
+ }
+
+ toLogicalPlan(targetOperatorId?: string): LogicalPlan {
+ const enabledOperators = this.getAllEnabledOperators();
+
+ const operators: LogicalOperator[] = enabledOperators.map(op => ({
+ operatorID: op.operatorID,
+ operatorType: op.operatorType,
+ ...op.operatorProperties,
+ inputPorts: op.inputPorts,
+ outputPorts: op.outputPorts,
+ }));
+
+ const operatorIds = new Set(operators.map(op => op.operatorID));
+
+ const links: LogicalLink[] = this.getAllLinks()
+ .filter(link => operatorIds.has(link.source.operatorID) && operatorIds.has(link.target.operatorID))
+ .map(link => {
+ const sourceOp = this.getOperator(link.source.operatorID)!;
+ const targetOp = this.getOperator(link.target.operatorID)!;
+
+ const fromPortIdx = sourceOp.outputPorts.findIndex(p => p.portID === link.source.portID);
+ const toPortIdx = targetOp.inputPorts.findIndex(p => p.portID === link.target.portID);
+
+ return {
+ fromOpId: link.source.operatorID,
+ fromPortId: { id: fromPortIdx >= 0 ? fromPortIdx : 0, internal: false },
+ toOpId: link.target.operatorID,
+ toPortId: { id: toPortIdx >= 0 ? toPortIdx : 0, internal: false },
+ };
+ });
+
+ return {
+ operators,
+ links,
+ opsToViewResult: Array.from(this.operatorsToViewResult).filter(id => operatorIds.has(id)),
+ opsToReuseResult: [],
+ };
+ }
+
+ addSubscription(subscription: Subscription): void {
+ this.subscriptions.push(subscription);
+ }
+
+ reset(): void {
+ this.operators.clear();
+ this.links.clear();
+ this.operatorPositions.clear();
+ this.commentBoxes = [];
+ this.settings = { ...DEFAULT_WORKFLOW_SETTINGS };
+ this.operatorsToViewResult.clear();
+ this.validationErrors = {};
+ this.workflowEmpty = true;
+ }
+
+ destroy(): void {
+ for (const sub of this.subscriptions) {
+ sub.unsubscribe();
+ }
+ this.subscriptions = [];
+
+ this.operatorAddSubject.complete();
+ this.operatorDeleteSubject.complete();
+ this.operatorPropertyChangeSubject.complete();
+ this.linkAddSubject.complete();
+ this.linkDeleteSubject.complete();
+ this.disabledOperatorChangedSubject.complete();
+ this.viewResultOperatorChangedSubject.complete();
+ this.validationChangedSubject.complete();
+
+ this.reset();
+ }
+}
diff --git a/agent-service/src/api/auth-api.ts b/agent-service/src/api/auth-api.ts
new file mode 100644
index 00000000000..087f93ac46f
--- /dev/null
+++ b/agent-service/src/api/auth-api.ts
@@ -0,0 +1,65 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import type { UserInfo } from "../types/agent";
+
+export type { UserInfo } from "../types/agent";
+
+function decodeJWT(token: string): any {
+ try {
+ const parts = token.split(".");
+ if (parts.length !== 3) {
+ throw new Error("Invalid JWT format");
+ }
+ return JSON.parse(Buffer.from(parts[1], "base64").toString("utf-8"));
+ } catch (error) {
+ throw new Error(`Failed to decode JWT: ${error}`);
+ }
+}
+
+export function extractUserFromToken(token: string): UserInfo {
+ const payload = decodeJWT(token);
+ return {
+ uid: payload.userId,
+ name: payload.sub,
+ email: payload.email || "",
+ role: payload.role || "REGULAR",
+ };
+}
+
+function isTokenExpired(token: string): boolean {
+ try {
+ const payload = decodeJWT(token);
+ if (!payload.exp) return false;
+ return Date.now() >= payload.exp * 1000;
+ } catch {
+ return true;
+ }
+}
+
+export function validateToken(token: string): boolean {
+ return !isTokenExpired(token);
+}
+
+export function createAuthHeaders(token: string): Record {
+ return {
+ Authorization: `Bearer ${token}`,
+ "Content-Type": "application/json",
+ };
+}
diff --git a/agent-service/src/api/backend-api.ts b/agent-service/src/api/backend-api.ts
new file mode 100644
index 00000000000..ffd2c59433f
--- /dev/null
+++ b/agent-service/src/api/backend-api.ts
@@ -0,0 +1,88 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { env } from "../config/env";
+
+interface BackendConfig {
+ apiEndpoint: string;
+ modelsEndpoint: string;
+ compileEndpoint: string;
+ executionEndpoint: string;
+}
+
+const currentConfig: BackendConfig = {
+ apiEndpoint: env.TEXERA_DASHBOARD_SERVICE_ENDPOINT,
+ modelsEndpoint: env.LLM_ENDPOINT,
+ compileEndpoint: env.WORKFLOW_COMPILING_SERVICE_ENDPOINT,
+ executionEndpoint: env.WORKFLOW_EXECUTION_SERVICE_ENDPOINT,
+};
+
+export function getBackendConfig(): BackendConfig {
+ return { ...currentConfig };
+}
+
+export interface InputPortInfo {
+ displayName?: string;
+ disallowMultiLinks?: boolean;
+ dependencies?: { id: number; internal: boolean }[];
+}
+
+export interface OutputPortInfo {
+ displayName?: string;
+}
+
+interface OperatorAdditionalMetadata {
+ userFriendlyName: string;
+ operatorGroupName: string;
+ operatorDescription?: string;
+ inputPorts: InputPortInfo[];
+ outputPorts: OutputPortInfo[];
+ dynamicInputPorts?: boolean;
+ dynamicOutputPorts?: boolean;
+ supportReconfiguration?: boolean;
+ allowPortCustomization?: boolean;
+}
+
+export interface OperatorSchema {
+ operatorType: string;
+ jsonSchema: any;
+ additionalMetadata: OperatorAdditionalMetadata;
+ operatorVersion: string;
+}
+
+interface GroupInfo {
+ groupName: string;
+ children?: GroupInfo[] | null;
+}
+
+export interface OperatorMetadata {
+ operators: OperatorSchema[];
+ groups: GroupInfo[];
+}
+
+export async function fetchOperatorMetadata(): Promise {
+ const url = `${currentConfig.apiEndpoint}/api/resources/operator-metadata`;
+ const response = await fetch(url);
+
+ if (!response.ok) {
+ throw new Error(`Failed to fetch operator metadata: ${response.status} ${response.statusText}`);
+ }
+
+ return (await response.json()) as OperatorMetadata;
+}
diff --git a/agent-service/src/api/compile-api.ts b/agent-service/src/api/compile-api.ts
new file mode 100644
index 00000000000..8ffd27fd52c
--- /dev/null
+++ b/agent-service/src/api/compile-api.ts
@@ -0,0 +1,74 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { getBackendConfig } from "./backend-api";
+import type { LogicalPlan, OperatorPortSchemaMap } from "../types/workflow";
+import { createLogger } from "../logger";
+
+const log = createLogger("CompileAPI");
+
+export interface SchemaAttribute {
+ attributeName: string;
+ attributeType: "string" | "integer" | "double" | "boolean" | "long" | "timestamp" | "binary";
+}
+
+export type PortSchema = ReadonlyArray;
+
+export interface WorkflowFatalError {
+ type: string;
+ message: string;
+ operatorId?: string;
+}
+
+export interface WorkflowCompilationResponse {
+ physicalPlan?: any;
+ operatorOutputSchemas: Record;
+ operatorErrors: Record;
+}
+
+export async function compileWorkflowAsync(logicalPlan: LogicalPlan): Promise {
+ const config = getBackendConfig();
+ const url = `${config.compileEndpoint}/api/compile`;
+
+ const body = {
+ operators: logicalPlan.operators,
+ links: logicalPlan.links,
+ opsToReuseResult: [],
+ opsToViewResult: [],
+ };
+
+ try {
+ const response = await fetch(url, {
+ method: "POST",
+ headers: { "Content-Type": "application/json" },
+ body: JSON.stringify(body),
+ });
+
+ if (!response.ok) {
+ const errorText = await response.text();
+ log.warn({ status: response.status, statusText: response.statusText, body: errorText }, "compilation failed");
+ return null;
+ }
+
+ return (await response.json()) as WorkflowCompilationResponse;
+ } catch (error) {
+ log.warn({ err: error }, "compile workflow API error");
+ return null;
+ }
+}
diff --git a/agent-service/src/api/execution-api.ts b/agent-service/src/api/execution-api.ts
new file mode 100644
index 00000000000..4692a61d754
--- /dev/null
+++ b/agent-service/src/api/execution-api.ts
@@ -0,0 +1,38 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export interface LogicalLink {
+ fromOpId: string;
+ fromPortId: { id: number; internal: boolean };
+ toOpId: string;
+ toPortId: { id: number; internal: boolean };
+}
+
+interface LogicalOperator {
+ operatorID: string;
+ operatorType: string;
+ [key: string]: any;
+}
+
+export interface LogicalPlan {
+ operators: LogicalOperator[];
+ links: LogicalLink[];
+ opsToViewResult?: string[];
+ opsToReuseResult?: string[];
+}
diff --git a/agent-service/src/api/index.ts b/agent-service/src/api/index.ts
new file mode 100644
index 00000000000..eca292d7ffe
--- /dev/null
+++ b/agent-service/src/api/index.ts
@@ -0,0 +1,24 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export * from "./backend-api";
+export * from "./execution-api";
+export * from "./workflow-api";
+export * from "./auth-api";
+export * from "./compile-api";
diff --git a/agent-service/src/api/workflow-api.ts b/agent-service/src/api/workflow-api.ts
new file mode 100644
index 00000000000..7a96f979a1c
--- /dev/null
+++ b/agent-service/src/api/workflow-api.ts
@@ -0,0 +1,97 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { getBackendConfig } from "./backend-api";
+import { createAuthHeaders } from "./auth-api";
+import type { WorkflowContent } from "../types/workflow";
+
+export interface Workflow {
+ wid: number;
+ name: string;
+ description?: string;
+ content: WorkflowContent;
+ creationTime?: number;
+ lastModifiedTime?: number;
+ isPublished?: boolean;
+}
+
+interface WorkflowPersistRequest {
+ wid?: number;
+ name: string;
+ description?: string;
+ content: string;
+ isPublic?: boolean;
+}
+
+const WORKFLOW_BASE_URL = "workflow";
+
+export async function persistWorkflow(
+ token: string,
+ wid: number,
+ name: string,
+ content: WorkflowContent,
+ description?: string
+): Promise {
+ const config = getBackendConfig();
+ const url = `${config.apiEndpoint}/api/${WORKFLOW_BASE_URL}/persist`;
+
+ const response = await fetch(url, {
+ method: "POST",
+ headers: createAuthHeaders(token),
+ body: JSON.stringify({
+ wid,
+ name,
+ description: description || "",
+ content: JSON.stringify(content),
+ isPublic: false,
+ } as WorkflowPersistRequest),
+ });
+
+ if (!response.ok) {
+ const errorText = await response.text();
+ throw new Error(`Failed to persist workflow: ${response.status} ${response.statusText} - ${errorText}`);
+ }
+
+ const data = (await response.json()) as Workflow;
+ if (typeof data.content === "string") {
+ data.content = JSON.parse(data.content as unknown as string);
+ }
+ return data;
+}
+
+export async function retrieveWorkflow(token: string, wid: number): Promise {
+ const config = getBackendConfig();
+ const url = `${config.apiEndpoint}/api/${WORKFLOW_BASE_URL}/${wid}`;
+
+ const response = await fetch(url, {
+ method: "GET",
+ headers: createAuthHeaders(token),
+ });
+
+ if (!response.ok) {
+ const errorText = await response.text();
+ throw new Error(`Failed to retrieve workflow: ${response.status} ${response.statusText} - ${errorText}`);
+ }
+
+ const data = (await response.json()) as Workflow;
+ if (typeof data.content === "string") {
+ data.content = JSON.parse(data.content as unknown as string);
+ }
+ return data;
+}
diff --git a/agent-service/src/config/env.ts b/agent-service/src/config/env.ts
new file mode 100644
index 00000000000..16a25b9be77
--- /dev/null
+++ b/agent-service/src/config/env.ts
@@ -0,0 +1,39 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { z } from "zod";
+
+const EnvSchema = z.object({
+ PORT: z.coerce.number().default(3001),
+ API_PREFIX: z.string().default("/api"),
+ LLM_API_KEY: z.string().default("dummy"),
+ TEXERA_SERVICE_LOG_LEVEL: z
+ .enum(["ERROR", "WARN", "INFO", "DEBUG"])
+ .transform(v => v.toLowerCase() as "error" | "warn" | "info" | "debug")
+ .default("INFO"),
+ LOG_PRETTY: z.coerce.boolean().default(false),
+
+ TEXERA_DASHBOARD_SERVICE_ENDPOINT: z.string().url().default("http://localhost:8080"),
+ LLM_ENDPOINT: z.string().url().default("http://localhost:9096"),
+ WORKFLOW_COMPILING_SERVICE_ENDPOINT: z.string().url().default("http://localhost:9090"),
+ WORKFLOW_EXECUTION_SERVICE_ENDPOINT: z.string().url().default("http://localhost:8085"),
+ EXECUTION_ENDPOINT_TEMPLATE: z.string().optional(),
+});
+
+export const env = EnvSchema.parse(process.env);
diff --git a/agent-service/src/index.ts b/agent-service/src/index.ts
new file mode 100644
index 00000000000..152a2d21703
--- /dev/null
+++ b/agent-service/src/index.ts
@@ -0,0 +1,26 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export * from "./types";
+export { WorkflowState } from "./agent/workflow-state";
+export { WorkflowResultState } from "./agent/workflow-result-state";
+export { WorkflowSystemMetadata } from "./agent/util/workflow-system-metadata";
+export * from "./agent/tools";
+export { TexeraAgent, type TexeraAgentConfig, type AgentMessageResult } from "./agent/texera-agent";
+export { buildSystemPrompt } from "./agent/prompts";
diff --git a/agent-service/src/logger.ts b/agent-service/src/logger.ts
new file mode 100644
index 00000000000..5f1537370f0
--- /dev/null
+++ b/agent-service/src/logger.ts
@@ -0,0 +1,47 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import pino, { type Logger } from "pino";
+import { env } from "./config/env";
+
+const rootLogger: Logger = pino({
+ level: env.TEXERA_SERVICE_LOG_LEVEL,
+ base: undefined,
+ ...(env.LOG_PRETTY
+ ? {
+ transport: {
+ target: "pino-pretty",
+ options: {
+ colorize: true,
+ translateTime: "HH:MM:ss.l",
+ ignore: "pid,hostname",
+ },
+ },
+ }
+ : {}),
+});
+
+// Prefer child loggers over manual `[Module agentId]` prefixes: `module` and
+// `agent` become structured fields in JSON output and render as a prefix in
+// pretty mode.
+export function createLogger(module: string, bindings: Record = {}): Logger {
+ return rootLogger.child({ module, ...bindings });
+}
+
+export const logger = rootLogger;
diff --git a/agent-service/src/server.test.ts b/agent-service/src/server.test.ts
new file mode 100644
index 00000000000..0f618e599c2
--- /dev/null
+++ b/agent-service/src/server.test.ts
@@ -0,0 +1,223 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { beforeEach, describe, expect, test } from "bun:test";
+import { buildApp, _resetAgentStoreForTests } from "./server";
+import { env } from "./config/env";
+
+const API = env.API_PREFIX;
+const app = buildApp();
+
+function url(path: string): string {
+ return `http://localhost${path}`;
+}
+
+async function postJson(path: string, body: unknown): Promise {
+ return app.handle(
+ new Request(url(path), {
+ method: "POST",
+ headers: { "Content-Type": "application/json" },
+ body: JSON.stringify(body),
+ })
+ );
+}
+
+async function patchJson(path: string, body: unknown): Promise {
+ return app.handle(
+ new Request(url(path), {
+ method: "PATCH",
+ headers: { "Content-Type": "application/json" },
+ body: JSON.stringify(body),
+ })
+ );
+}
+
+async function getJson(path: string): Promise {
+ return app.handle(new Request(url(path)));
+}
+
+async function del(path: string): Promise {
+ return app.handle(new Request(url(path), { method: "DELETE" }));
+}
+
+async function readJson(res: Response): Promise {
+ return (await res.json()) as T;
+}
+
+beforeEach(() => {
+ _resetAgentStoreForTests();
+});
+
+describe(`GET ${API}/healthcheck`, () => {
+ test("returns 200 with status ok", async () => {
+ const res = await getJson(`${API}/healthcheck`);
+ expect(res.status).toBe(200);
+ const body = await readJson<{ status: string; timestamp: string }>(res);
+ expect(body.status).toBe("ok");
+ expect(typeof body.timestamp).toBe("string");
+ });
+});
+
+describe(`POST ${API}/agents`, () => {
+ test("creates an agent with no delegate", async () => {
+ const res = await postJson(`${API}/agents`, { modelType: "test-model", name: "Tester" });
+ expect(res.status).toBe(200);
+
+ const agent = await readJson<{
+ id: string;
+ name: string;
+ modelType: string;
+ state: string;
+ delegate: unknown;
+ }>(res);
+ expect(agent.id).toMatch(/^agent-\d+$/);
+ expect(agent.name).toBe("Tester");
+ expect(agent.modelType).toBe("test-model");
+ expect(agent.state).toBe("AVAILABLE");
+ expect(agent.delegate).toBeUndefined();
+ });
+
+ test("auto-numbers agent ids monotonically", async () => {
+ const a = await readJson<{ id: string }>(await postJson(`${API}/agents`, { modelType: "m" }));
+ const b = await readJson<{ id: string }>(await postJson(`${API}/agents`, { modelType: "m" }));
+
+ const aNum = Number(a.id.split("-")[1]);
+ const bNum = Number(b.id.split("-")[1]);
+ expect(bNum).toBe(aNum + 1);
+ });
+
+ test("rejects invalid token", async () => {
+ const res = await postJson(`${API}/agents`, {
+ modelType: "m",
+ userToken: "obviously-not-a-jwt",
+ });
+ expect(res.status).toBe(401);
+ const body = await readJson<{ error: string }>(res);
+ expect(body.error).toBe("Invalid or expired token");
+ });
+
+ test("rejects missing modelType", async () => {
+ const res = await postJson(`${API}/agents`, { name: "no-model" });
+ // Body schema violation; the exact status depends on the Elysia version but
+ // it is always a 4xx or 5xx, never a successful 2xx.
+ expect(res.status).toBeGreaterThanOrEqual(400);
+ });
+});
+
+describe(`GET ${API}/agents`, () => {
+ test("empty store returns no agents", async () => {
+ const res = await getJson(`${API}/agents`);
+ expect(res.status).toBe(200);
+ const body = await readJson<{ agents: unknown[] }>(res);
+ expect(body.agents).toEqual([]);
+ });
+
+ test("lists every created agent", async () => {
+ await postJson(`${API}/agents`, { modelType: "m", name: "one" });
+ await postJson(`${API}/agents`, { modelType: "m", name: "two" });
+
+ const res = await getJson(`${API}/agents`);
+ const body = await readJson<{ agents: { name: string }[] }>(res);
+ expect(body.agents).toHaveLength(2);
+ expect(body.agents.map(a => a.name).sort()).toEqual(["one", "two"]);
+ });
+});
+
+describe(`GET ${API}/agents/:id`, () => {
+ test("returns the agent plus its workflow snapshot", async () => {
+ const created = await readJson<{ id: string }>(await postJson(`${API}/agents`, { modelType: "m" }));
+
+ const res = await getJson(`${API}/agents/${created.id}`);
+ expect(res.status).toBe(200);
+ const body = await readJson<{ id: string; workflow: unknown; stepCount: number }>(res);
+ expect(body.id).toBe(created.id);
+ expect(body.workflow).toBeDefined();
+ expect(typeof body.stepCount).toBe("number");
+ });
+
+ test("returns 404 for an unknown id", async () => {
+ const res = await getJson(`${API}/agents/agent-does-not-exist`);
+ expect(res.status).toBe(404);
+ const body = await readJson<{ error: string }>(res);
+ expect(body.error).toBe("Agent not found");
+ });
+});
+
+describe(`DELETE ${API}/agents/:id`, () => {
+ test("destroys the agent and a follow-up GET returns 404", async () => {
+ const created = await readJson<{ id: string }>(await postJson(`${API}/agents`, { modelType: "m" }));
+
+ const delRes = await del(`${API}/agents/${created.id}`);
+ expect(delRes.status).toBe(200);
+ expect(await readJson(delRes)).toEqual({ deleted: true });
+
+ const getRes = await getJson(`${API}/agents/${created.id}`);
+ expect(getRes.status).toBe(404);
+ });
+
+ test("returns 404 when deleting an unknown agent", async () => {
+ const res = await del(`${API}/agents/missing`);
+ expect(res.status).toBe(404);
+ });
+});
+
+describe("Agent control routes", () => {
+ test("POST /:id/stop returns stopping", async () => {
+ const created = await readJson<{ id: string }>(await postJson(`${API}/agents`, { modelType: "m" }));
+ const res = await postJson(`${API}/agents/${created.id}/stop`, {});
+ expect(res.status).toBe(200);
+ expect(await readJson(res)).toEqual({ status: "stopping" });
+ });
+
+ test("POST /:id/clear resets history", async () => {
+ const created = await readJson<{ id: string }>(await postJson(`${API}/agents`, { modelType: "m" }));
+ const res = await postJson(`${API}/agents/${created.id}/clear`, {});
+ expect(res.status).toBe(200);
+ expect(await readJson(res)).toEqual({ status: "cleared" });
+ });
+
+ test("GET /:id/operator-results returns an empty map on the framework build", async () => {
+ const created = await readJson<{ id: string }>(await postJson(`${API}/agents`, { modelType: "m" }));
+ const res = await getJson(`${API}/agents/${created.id}/operator-results`);
+ expect(res.status).toBe(200);
+ expect(await readJson(res)).toEqual({ results: {} });
+ });
+});
+
+describe(`PATCH ${API}/agents/:id/settings`, () => {
+ test("updates settings and returns the new values", async () => {
+ const created = await readJson<{ id: string }>(await postJson(`${API}/agents`, { modelType: "m" }));
+
+ const res = await patchJson(`${API}/agents/${created.id}/settings`, {
+ maxSteps: 7,
+ toolTimeoutSeconds: 30,
+ });
+ expect(res.status).toBe(200);
+ const body = await readJson<{ maxSteps: number; toolTimeoutSeconds: number }>(res);
+ expect(body.maxSteps).toBe(7);
+ expect(body.toolTimeoutSeconds).toBe(30);
+
+ // A follow-up GET reflects the same values.
+ const reread = await readJson<{ maxSteps: number; toolTimeoutSeconds: number }>(
+ await getJson(`${API}/agents/${created.id}/settings`)
+ );
+ expect(reread.maxSteps).toBe(7);
+ expect(reread.toolTimeoutSeconds).toBe(30);
+ });
+});
diff --git a/agent-service/src/server.ts b/agent-service/src/server.ts
new file mode 100644
index 00000000000..a31f9ede115
--- /dev/null
+++ b/agent-service/src/server.ts
@@ -0,0 +1,666 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { Elysia, t } from "elysia";
+import { cors } from "@elysiajs/cors";
+import { createOpenAI } from "@ai-sdk/openai";
+import { TexeraAgent } from "./agent/texera-agent";
+import { getBackendConfig } from "./api/backend-api";
+import { extractUserFromToken, validateToken } from "./api/auth-api";
+import { retrieveWorkflow } from "./api/workflow-api";
+import { WorkflowSystemMetadata } from "./agent/util/workflow-system-metadata";
+import { env } from "./config/env";
+import { createLogger } from "./logger";
+
+const log = createLogger("Server");
+const wsLog = createLogger("WS");
+import type {
+ AgentInfo,
+ AgentDelegateConfig,
+ CreateAgentRequest,
+ UpdateAgentSettingsRequest,
+ AgentSettingsApi,
+ ReActStep,
+} from "./types/agent";
+import { OperatorResultSerializationMode } from "./types/agent";
+
+const agentStore = new Map();
+let agentCounter = 0;
+
+async function createAgentInstance(
+ modelType: string,
+ customName?: string,
+ delegateConfig?: AgentDelegateConfig
+): Promise<{ agentId: string; agent: TexeraAgent }> {
+ const agentId = `agent-${++agentCounter}`;
+ const config = getBackendConfig();
+
+ const openai = createOpenAI({
+ baseURL: `${config.modelsEndpoint}/api`,
+ apiKey: env.LLM_API_KEY,
+ });
+
+ // Reasoning effort variants are configured as separate model entries in litellm-config.yaml
+ // with extra_body to inject reasoning_effort, bypassing LiteLLM's param validation.
+ const agent = new TexeraAgent({
+ model: openai.chat(modelType),
+ modelType,
+ agentId,
+ agentName: customName || "Bob",
+ });
+
+ await agent.initialize();
+
+ if (delegateConfig?.workflowId && delegateConfig.userToken) {
+ try {
+ const workflow = await retrieveWorkflow(delegateConfig.userToken, delegateConfig.workflowId);
+ delegateConfig.workflowName = workflow.name;
+
+ const workflowState = agent.getWorkflowState();
+ workflowState.setWorkflowContent(workflow.content);
+
+ agent.setDelegateConfig({
+ userToken: delegateConfig.userToken,
+ userInfo: delegateConfig.userInfo,
+ workflowId: delegateConfig.workflowId,
+ workflowName: delegateConfig.workflowName,
+ computingUnitId: delegateConfig.computingUnitId,
+ });
+
+ log.info({ agentId, workflowId: delegateConfig.workflowId }, "loaded workflow for agent");
+ } catch (error) {
+ log.warn({ agentId, workflowId: delegateConfig.workflowId, err: error }, "failed to load workflow");
+ }
+ }
+
+ agentStore.set(agentId, agent);
+ log.info({ agentId, delegate: !!delegateConfig }, "created agent");
+
+ return { agentId, agent };
+}
+
+function getAgentInfo(agentId: string, agent: TexeraAgent): AgentInfo {
+ const agentSettings = agent.getSettings();
+ const settingsApi: AgentSettingsApi = {
+ maxOperatorResultCharLimit: agentSettings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: agentSettings.maxOperatorResultCellCharLimit,
+ operatorResultSerializationMode: agentSettings.operatorResultSerializationMode,
+ toolTimeoutSeconds: Math.round(agentSettings.toolTimeoutMs / 1000),
+ executionTimeoutMinutes: Math.round(agentSettings.executionTimeoutMs / 60000),
+ disabledTools: Array.from(agentSettings.disabledTools),
+ maxSteps: agentSettings.maxSteps,
+ allowedOperatorTypes: agentSettings.allowedOperatorTypes,
+ };
+
+ const delegateConfig = agent.getDelegateConfig();
+
+ return {
+ id: agentId,
+ name: agent.agentName,
+ modelType: agent.modelType,
+ state: agent.getState(),
+ createdAt: agent.createdAt,
+ delegate: delegateConfig
+ ? {
+ userToken: "***",
+ userInfo: delegateConfig.userInfo,
+ workflowId: delegateConfig.workflowId,
+ workflowName: delegateConfig.workflowName,
+ computingUnitId: delegateConfig.computingUnitId,
+ }
+ : undefined,
+ settings: settingsApi,
+ };
+}
+
+function getAgent(agentId: string): TexeraAgent {
+ const agent = agentStore.get(agentId);
+ if (!agent) {
+ throw new Error("Agent not found");
+ }
+ return agent;
+}
+
+const agentsRouter = new Elysia({ prefix: "/agents" })
+ // Error handler must live on the same Elysia instance whose routes throw, or
+ // its scope will not see the errors. Elysia 1.x defaults to local scoping for
+ // .onError, so attach here rather than on the outer app.
+ .onError(({ error, set }) => {
+ log.error({ err: error }, "request error");
+ const errorMessage = error instanceof Error ? error.message : String(error);
+ if (errorMessage === "Agent not found") {
+ set.status = 404;
+ return { error: "Agent not found" };
+ }
+ if (errorMessage === "Invalid or expired token") {
+ set.status = 401;
+ return { error: "Invalid or expired token" };
+ }
+ if (errorMessage === "modelType is required") {
+ set.status = 400;
+ return { error: "modelType is required" };
+ }
+ set.status = 500;
+ return { error: errorMessage || "Internal server error" };
+ })
+ .get("/", () => {
+ const agentList = Array.from(agentStore.entries()).map(([id, agent]) => getAgentInfo(id, agent));
+ return { agents: agentList };
+ })
+
+ .post(
+ "/",
+ async ({ body }) => {
+ const { modelType, name, userToken, workflowId, computingUnitId, settings } = body as CreateAgentRequest;
+
+ if (!modelType) {
+ throw new Error("modelType is required");
+ }
+
+ let delegateConfig: AgentDelegateConfig | undefined;
+ if (userToken) {
+ if (!validateToken(userToken)) {
+ throw new Error("Invalid or expired token");
+ }
+
+ const userInfo = extractUserFromToken(userToken);
+ delegateConfig = {
+ userToken,
+ userInfo,
+ workflowId,
+ computingUnitId,
+ };
+ }
+
+ const { agentId, agent } = await createAgentInstance(modelType, name, delegateConfig);
+
+ if (settings) {
+ log.info(
+ {
+ agentId,
+ maxOperatorResultCharLimit: settings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: settings.maxOperatorResultCellCharLimit,
+ },
+ "applying initial agent settings"
+ );
+ agent.updateSettings({
+ maxOperatorResultCharLimit: settings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: settings.maxOperatorResultCellCharLimit,
+ operatorResultSerializationMode: settings.operatorResultSerializationMode
+ ? (settings.operatorResultSerializationMode as OperatorResultSerializationMode)
+ : undefined,
+ toolTimeoutMs: settings.toolTimeoutSeconds ? settings.toolTimeoutSeconds * 1000 : undefined,
+ executionTimeoutMs: settings.executionTimeoutMinutes ? settings.executionTimeoutMinutes * 60000 : undefined,
+ disabledTools: settings.disabledTools ? new Set(settings.disabledTools) : undefined,
+ maxSteps: settings.maxSteps,
+ allowedOperatorTypes: settings.allowedOperatorTypes,
+ });
+ }
+
+ return getAgentInfo(agentId, agent);
+ },
+ {
+ body: t.Object({
+ modelType: t.String(),
+ name: t.Optional(t.String()),
+ userToken: t.Optional(t.String()),
+ workflowId: t.Optional(t.Number()),
+ computingUnitId: t.Optional(t.Number()),
+ settings: t.Optional(
+ t.Object({
+ maxOperatorResultCharLimit: t.Optional(t.Number()),
+ maxOperatorResultCellCharLimit: t.Optional(t.Number()),
+ operatorResultSerializationMode: t.Optional(t.Literal("tsv")),
+ toolTimeoutSeconds: t.Optional(t.Number()),
+ executionTimeoutMinutes: t.Optional(t.Number()),
+ disabledTools: t.Optional(t.Array(t.String())),
+ maxSteps: t.Optional(t.Number()),
+ allowedOperatorTypes: t.Optional(t.Array(t.String())),
+ })
+ ),
+ }),
+ }
+ )
+
+ .get("/:id", ({ params: { id } }) => {
+ const agent = getAgent(id);
+ return {
+ ...getAgentInfo(id, agent),
+ workflow: agent.getWorkflowState().getWorkflowContent(),
+ stepCount: agent.getReActSteps().length,
+ };
+ })
+
+ .delete("/:id", ({ params: { id }, set }) => {
+ const agent = agentStore.get(id);
+ if (!agent) {
+ set.status = 404;
+ return { error: "Agent not found" };
+ }
+
+ agent.destroy();
+ agentStore.delete(id);
+ return { deleted: true };
+ })
+
+ .get("/:id/react-steps", ({ params: { id } }) => {
+ const agent = getAgent(id);
+ return { steps: agent.getReActSteps(), state: agent.getState() };
+ })
+
+ .get("/:id/operator-results", ({ params: { id } }) => {
+ const agent = getAgent(id);
+ return { results: getOperatorResultSummaries(agent) };
+ })
+
+ .post(
+ "/:id/steps-by-operators",
+ ({ params: { id }, body }) => {
+ const agent = getAgent(id);
+ const { operatorIds } = body;
+ return { steps: agent.getReActStepsByOperatorIds(operatorIds || []) };
+ },
+ {
+ body: t.Object({
+ operatorIds: t.Array(t.String()),
+ }),
+ }
+ )
+
+ .get("/:id/system-info", ({ params: { id } }) => {
+ const agent = getAgent(id);
+ return agent.getSystemInfo();
+ })
+
+ .post("/:id/stop", ({ params: { id } }) => {
+ const agent = getAgent(id);
+ agent.stop();
+ return { status: "stopping" };
+ })
+
+ .post("/:id/clear", ({ params: { id } }) => {
+ const agent = getAgent(id);
+ agent.clearHistory();
+ return { status: "cleared" };
+ })
+
+ .post("/:id/checkout", ({ params: { id }, body }) => {
+ const agent = getAgent(id);
+ const { stepId } = body as { stepId: string };
+ if (!stepId) throw new Error("stepId is required");
+
+ const success = agent.checkout(stepId);
+ if (!success) throw new Error(`Step ${stepId} not found or checkout failed`);
+
+ const allSteps = agent.getAllSteps();
+ const workflowContent = agent.getWorkflowState().getWorkflowContent();
+
+ broadcastToAgent(id, {
+ type: "headChange",
+ headId: stepId,
+ steps: allSteps,
+ workflowContent,
+ operatorResults: getOperatorResultSummaries(agent),
+ });
+
+ return {
+ status: "checked out",
+ headId: stepId,
+ };
+ })
+
+ .get("/:id/operator-types", ({ params: { id } }) => {
+ const agent = getAgent(id);
+ const metadataStore = agent.getMetadataStore();
+ const allTypes = metadataStore.getAllOperatorTypes();
+ return Object.entries(allTypes).map(([type, description]) => ({ type, description }));
+ })
+
+ .get("/:id/settings", ({ params: { id } }) => {
+ const agent = getAgent(id);
+ const agentSettings = agent.getSettings();
+ return {
+ maxOperatorResultCharLimit: agentSettings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: agentSettings.maxOperatorResultCellCharLimit,
+ operatorResultSerializationMode: agentSettings.operatorResultSerializationMode,
+ toolTimeoutSeconds: Math.round(agentSettings.toolTimeoutMs / 1000),
+ executionTimeoutMinutes: Math.round(agentSettings.executionTimeoutMs / 60000),
+ disabledTools: Array.from(agentSettings.disabledTools),
+ maxSteps: agentSettings.maxSteps,
+ allowedOperatorTypes: agentSettings.allowedOperatorTypes,
+ };
+ })
+
+ .patch(
+ "/:id/settings",
+ ({ params: { id }, body }) => {
+ const agent = getAgent(id);
+ const settings = body as UpdateAgentSettingsRequest;
+
+ log.info(
+ {
+ agentId: id,
+ maxOperatorResultCharLimit: settings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: settings.maxOperatorResultCellCharLimit,
+ },
+ "updating agent settings"
+ );
+
+ agent.updateSettings({
+ maxOperatorResultCharLimit: settings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: settings.maxOperatorResultCellCharLimit,
+ operatorResultSerializationMode: settings.operatorResultSerializationMode
+ ? (settings.operatorResultSerializationMode as OperatorResultSerializationMode)
+ : undefined,
+ toolTimeoutMs: settings.toolTimeoutSeconds !== undefined ? settings.toolTimeoutSeconds * 1000 : undefined,
+ executionTimeoutMs:
+ settings.executionTimeoutMinutes !== undefined ? settings.executionTimeoutMinutes * 60000 : undefined,
+ disabledTools: settings.disabledTools ? new Set(settings.disabledTools) : undefined,
+ maxSteps: settings.maxSteps,
+ allowedOperatorTypes: settings.allowedOperatorTypes,
+ });
+
+ const agentSettings = agent.getSettings();
+ return {
+ maxOperatorResultCharLimit: agentSettings.maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit: agentSettings.maxOperatorResultCellCharLimit,
+ operatorResultSerializationMode: agentSettings.operatorResultSerializationMode,
+ toolTimeoutSeconds: Math.round(agentSettings.toolTimeoutMs / 1000),
+ executionTimeoutMinutes: Math.round(agentSettings.executionTimeoutMs / 60000),
+ disabledTools: Array.from(agentSettings.disabledTools),
+ maxSteps: agentSettings.maxSteps,
+ allowedOperatorTypes: agentSettings.allowedOperatorTypes,
+ };
+ },
+ {
+ body: t.Object({
+ maxOperatorResultCharLimit: t.Optional(t.Number()),
+ maxOperatorResultCellCharLimit: t.Optional(t.Number()),
+ operatorResultSerializationMode: t.Optional(t.Literal("tsv")),
+ toolTimeoutSeconds: t.Optional(t.Number()),
+ executionTimeoutMinutes: t.Optional(t.Number()),
+ maxSteps: t.Optional(t.Number()),
+ disabledTools: t.Optional(t.Array(t.String())),
+ allowedOperatorTypes: t.Optional(t.Array(t.String())),
+ }),
+ }
+ );
+
+interface WsMessage {
+ type: "message" | "stop";
+ content?: string;
+ messageSource?: "chat" | "feedback";
+}
+
+interface OperatorResultSummaryWs {
+ state: string;
+ inputTuples: number;
+ outputTuples: number;
+ inputPortShapes?: { portIndex: number; rows: number; columns: number }[];
+ outputColumns?: number;
+ error?: string;
+ warnings?: string[];
+ consoleLogCount?: number;
+ totalRowCount?: number;
+ sampleRecords?: Record[];
+ resultStatistics?: Record;
+}
+
+interface WsOutgoingMessage {
+ type: "step" | "state" | "error" | "complete" | "init" | "headChange";
+ step?: ReActStep;
+ state?: string;
+ error?: string;
+ steps?: ReActStep[];
+ headId?: string;
+ operatorResults?: Record;
+ workflowContent?: any;
+}
+
+function getOperatorResultSummaries(agent: TexeraAgent): Record {
+ const resultState = agent.getWorkflowResultState();
+ const visible = resultState.getAllVisible();
+ const results: Record = {};
+ for (const [opId, entry] of visible) {
+ const info = entry.operatorInfo;
+ results[opId] = {
+ state: info.state,
+ inputTuples: info.inputTuples,
+ outputTuples: info.outputTuples,
+ inputPortShapes: info.inputPortShapes,
+ outputColumns:
+ info.result && info.result.length > 0
+ ? Object.keys(info.result[0]).filter(k => k !== "__row_index__").length
+ : undefined,
+ error: info.error,
+ warnings: info.warnings,
+ consoleLogCount: info.consoleLogs?.length,
+ totalRowCount: info.totalRowCount,
+ sampleRecords: info.result,
+ resultStatistics: info.resultStatistics,
+ };
+ }
+ return results;
+}
+
+function broadcastToAgent(agentId: string, message: WsOutgoingMessage): void {
+ const agent = agentStore.get(agentId);
+ if (!agent) return;
+
+ const jsonMessage = JSON.stringify(message);
+ for (const ws of agent.getWebsockets()) {
+ try {
+ ws.send(jsonMessage);
+ } catch (error) {
+ wsLog.error({ agentId, err: error }, "failed to send message to client");
+ agent.removeWebsocket(ws);
+ }
+ }
+}
+
+export function buildApp() {
+ return new Elysia()
+ .use(cors())
+ .group(env.API_PREFIX, app =>
+ app
+ .get("/healthcheck", () => ({
+ status: "ok",
+ timestamp: new Date().toISOString(),
+ }))
+ .use(agentsRouter)
+ )
+ .ws(`${env.API_PREFIX}/agents/:id/react`, {
+ open(ws) {
+ const agentId = (ws.data as any).params?.id;
+ wsLog.info({ agentId }, "client connected");
+
+ const agent = agentStore.get(agentId);
+ if (!agent) {
+ ws.send(JSON.stringify({ type: "error", error: "Agent not found" }));
+ ws.close();
+ return;
+ }
+
+ agent.addWebsocket(ws);
+
+ const initMessage: WsOutgoingMessage = {
+ type: "init",
+ state: agent.getState(),
+ steps: agent.getAllSteps(),
+ headId: agent.getHead(),
+ operatorResults: getOperatorResultSummaries(agent),
+ };
+ ws.send(JSON.stringify(initMessage));
+ },
+
+ async message(ws, messageData) {
+ const agentId = (ws.data as any).params?.id;
+ const agent = agentStore.get(agentId);
+
+ if (!agent) {
+ ws.send(JSON.stringify({ type: "error", error: "Agent not found" }));
+ return;
+ }
+
+ let msg: WsMessage;
+ try {
+ msg = typeof messageData === "string" ? JSON.parse(messageData) : (messageData as WsMessage);
+ } catch {
+ ws.send(JSON.stringify({ type: "error", error: "Invalid message format" }));
+ return;
+ }
+
+ if (msg.type === "stop") {
+ agent.stop();
+ broadcastToAgent(agentId, { type: "state", state: "STOPPING" });
+ return;
+ }
+
+ if (msg.type === "message") {
+ if (!msg.content || typeof msg.content !== "string") {
+ ws.send(JSON.stringify({ type: "error", error: "Message content is required" }));
+ return;
+ }
+
+ wsLog.info({ agentId, preview: msg.content.substring(0, 50) }, "received message");
+
+ agent.setStepCallback((step: ReActStep) => {
+ const hasToolCalls = step.toolCalls && step.toolCalls.length > 0;
+ broadcastToAgent(agentId, {
+ type: "step",
+ step,
+ ...(hasToolCalls ? { operatorResults: getOperatorResultSummaries(agent) } : {}),
+ });
+ });
+
+ broadcastToAgent(agentId, { type: "state", state: "GENERATING" });
+
+ try {
+ const result = await agent.sendMessage(msg.content, msg.messageSource);
+
+ agent.setStepCallback(null);
+
+ const allSteps = agent.getReActSteps();
+ const lastStep = allSteps[allSteps.length - 1];
+ if (lastStep && lastStep.isEnd) {
+ broadcastToAgent(agentId, { type: "step", step: lastStep });
+ }
+
+ broadcastToAgent(agentId, {
+ type: "complete",
+ state: agent.getState(),
+ operatorResults: getOperatorResultSummaries(agent),
+ });
+
+ wsLog.info({ agentId, steps: result.messages.length }, "agent run complete");
+ } catch (error: any) {
+ agent.setStepCallback(null);
+ broadcastToAgent(agentId, { type: "error", error: error.message });
+ }
+ }
+ },
+
+ close(ws) {
+ const agentId = (ws.data as any).params?.id;
+ wsLog.info({ agentId }, "client disconnected");
+
+ const agent = agentStore.get(agentId);
+ if (agent) {
+ agent.removeWebsocket(ws);
+ }
+ },
+ })
+ .onError(({ error, set }) => {
+ // Catch-all for non-router routes such as /api/healthcheck and the websocket route.
+ log.error({ err: error }, "request error");
+ set.status = 500;
+ return { error: error instanceof Error ? error.message : String(error) };
+ });
+}
+
+// Reset module-level state. Used by tests to start each case from a clean store.
+export function _resetAgentStoreForTests(): void {
+ agentStore.clear();
+ agentCounter = 0;
+}
+
+function printStartupMessage(app: ReturnType) {
+ const LINE = "=".repeat(60);
+ console.log(LINE);
+ console.log("Texera Agent Service (Elysia.js + RxJS)");
+ console.log(LINE);
+ console.log(`Server running at http://localhost:${env.PORT}`);
+ console.log("");
+
+ console.log("Registered Routes:");
+ const routes = app.routes;
+
+ const httpRoutes = routes.filter(r => r.method !== "WS");
+ const wsRoutes = routes.filter(r => r.method === "WS");
+
+ for (const route of httpRoutes) {
+ const method = route.method.padEnd(6);
+ console.log(` ${method} ${route.path}`);
+ }
+
+ if (wsRoutes.length > 0) {
+ console.log("");
+ console.log("WebSocket Endpoints:");
+ for (const route of wsRoutes) {
+ console.log(` WS ${route.path}`);
+ }
+ console.log(" Send: { type: 'message', content: '...' }");
+ console.log(" Send: { type: 'stop' }");
+ console.log(" Recv: { type: 'step' | 'state' | 'complete' | 'error' | 'init', ... }");
+ }
+
+ console.log("");
+ console.log("Environment:");
+ console.log(` LLM_API_KEY: ${env.LLM_API_KEY === "dummy" ? "dummy (default)" : "set"}`);
+ console.log(` LLM_ENDPOINT: ${getBackendConfig().modelsEndpoint}`);
+ console.log(` WORKFLOW_COMPILING_SERVICE_ENDPOINT: ${getBackendConfig().compileEndpoint}`);
+ console.log(` TEXERA_DASHBOARD_SERVICE_ENDPOINT: ${getBackendConfig().apiEndpoint}`);
+ console.log("");
+ console.log("Features:");
+ console.log(" - Auto-persistence with debounce (500ms)");
+ console.log(LINE);
+}
+
+async function initializeServices() {
+ try {
+ log.info("initializing global workflow system metadata");
+ const metadata = await WorkflowSystemMetadata.initializeGlobal();
+ log.info({ operatorCount: metadata.getOperatorCount() }, "loaded operators into global metadata");
+ } catch (error) {
+ log.warn({ err: error }, "failed to initialize global metadata; agents will initialize individually");
+ }
+}
+
+export async function start() {
+ await initializeServices();
+ const app = buildApp().listen(env.PORT);
+ printStartupMessage(app);
+ return app;
+}
+
+// Run the server only when this file is the entry point, not when it is
+// imported by tests or other modules.
+if (import.meta.main) {
+ start();
+}
diff --git a/agent-service/src/types/agent.ts b/agent-service/src/types/agent.ts
new file mode 100644
index 00000000000..765f5a7cb46
--- /dev/null
+++ b/agent-service/src/types/agent.ts
@@ -0,0 +1,165 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import type { WorkflowContent } from "./workflow";
+
+export enum AgentState {
+ UNAVAILABLE = "UNAVAILABLE",
+ AVAILABLE = "AVAILABLE",
+ GENERATING = "GENERATING",
+ STOPPING = "STOPPING",
+}
+
+export interface TokenUsage {
+ inputTokens?: number;
+ outputTokens?: number;
+ totalTokens?: number;
+ cachedInputTokens?: number;
+}
+
+export const INITIAL_STEP_ID = "step-initial";
+
+export interface ReActStep {
+ id: string;
+ parentId?: string;
+ messageId: string;
+ stepId: number;
+ timestamp: number;
+ role: "user" | "agent";
+ content: string;
+ isBegin: boolean;
+ isEnd: boolean;
+ toolCalls?: Array<{
+ toolName: string;
+ toolCallId: string;
+ input: any;
+ }>;
+ toolResults?: Array<{
+ toolCallId: string;
+ output: any;
+ isError?: boolean;
+ }>;
+ usage?: TokenUsage;
+ inputMessages?: any[];
+ messageSource?: "chat" | "feedback";
+ beforeWorkflowContent?: WorkflowContent;
+ afterWorkflowContent?: WorkflowContent;
+}
+
+export enum OperatorResultSerializationMode {
+ TSV = "tsv",
+}
+
+export interface AgentSettings {
+ systemPrompt: string;
+ disabledTools: Set;
+ maxOperatorResultCharLimit: number;
+ maxOperatorResultCellCharLimit: number;
+ operatorResultSerializationMode: OperatorResultSerializationMode;
+ toolTimeoutMs: number;
+ executionTimeoutMs: number;
+ maxSteps: number;
+ allowedOperatorTypes: string[];
+}
+
+export const DEFAULT_AGENT_SETTINGS: Omit = {
+ disabledTools: new Set(),
+ maxOperatorResultCharLimit: 2000,
+ maxOperatorResultCellCharLimit: 2000,
+ operatorResultSerializationMode: OperatorResultSerializationMode.TSV,
+ toolTimeoutMs: 240000,
+ executionTimeoutMs: 240000,
+ maxSteps: 100,
+ allowedOperatorTypes: [
+ "CSVFileScan",
+ "Filter",
+ "Projection",
+ "TypeCasting",
+ "Sort",
+ "Limit",
+ "Distinct",
+ "Union",
+ "KeywordSearch",
+ "HashJoin",
+ "Aggregate",
+ "LineChart",
+ "BarChart",
+ "PieChart",
+ "Histogram",
+ "Scatterplot",
+ "WordCloud",
+ "PythonUDFV2",
+ ],
+};
+
+export interface UserInfo {
+ uid: number;
+ name: string;
+ email: string;
+ role: string;
+}
+
+export interface AgentDelegateConfig {
+ userToken: string;
+ userInfo?: UserInfo;
+ workflowId?: number;
+ workflowName?: string;
+ computingUnitId?: number;
+}
+
+export interface AgentSettingsApi {
+ maxOperatorResultCharLimit?: number;
+ maxOperatorResultCellCharLimit?: number;
+ operatorResultSerializationMode?: "tsv";
+ toolTimeoutSeconds?: number;
+ executionTimeoutMinutes?: number;
+ disabledTools?: string[];
+ maxSteps?: number;
+ allowedOperatorTypes?: string[];
+}
+
+export interface AgentInfo {
+ id: string;
+ name: string;
+ modelType: string;
+ state: AgentState;
+ createdAt: Date;
+ delegate?: AgentDelegateConfig;
+ settings?: AgentSettingsApi;
+}
+
+export interface CreateAgentRequest {
+ modelType: string;
+ name?: string;
+ userToken?: string;
+ workflowId?: number;
+ computingUnitId?: number;
+ settings?: AgentSettingsApi;
+}
+
+export interface UpdateAgentSettingsRequest {
+ maxOperatorResultCharLimit?: number;
+ maxOperatorResultCellCharLimit?: number;
+ operatorResultSerializationMode?: "tsv";
+ toolTimeoutSeconds?: number;
+ executionTimeoutMinutes?: number;
+ disabledTools?: string[];
+ maxSteps?: number;
+ allowedOperatorTypes?: string[];
+}
diff --git a/agent-service/src/types/execution.ts b/agent-service/src/types/execution.ts
new file mode 100644
index 00000000000..f93be5c583e
--- /dev/null
+++ b/agent-service/src/types/execution.ts
@@ -0,0 +1,53 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+interface ConsoleMessage {
+ msgType: string;
+ message: string;
+}
+
+interface PortShape {
+ portIndex: number;
+ rows: number;
+ columns: number;
+}
+
+export interface OperatorInfo {
+ state: string;
+ inputTuples: number;
+ outputTuples: number;
+ inputPortShapes?: PortShape[];
+ resultMode: string;
+ result?: Record[];
+ totalRowCount?: number;
+ displayedRows?: number;
+ truncated?: boolean;
+ consoleLogs?: ConsoleMessage[];
+ error?: string;
+ warnings?: string[];
+ resultStatistics?: Record;
+}
+
+export interface SyncExecutionResult {
+ success: boolean;
+ state: string;
+ operators: Record;
+ compilationErrors?: Record;
+ errors?: string[];
+}
diff --git a/agent-service/src/types/index.ts b/agent-service/src/types/index.ts
new file mode 100644
index 00000000000..c6d7291e51d
--- /dev/null
+++ b/agent-service/src/types/index.ts
@@ -0,0 +1,22 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export * from "./workflow";
+export * from "./execution";
+export * from "./agent";
diff --git a/agent-service/src/types/workflow.ts b/agent-service/src/types/workflow.ts
new file mode 100644
index 00000000000..52c6493cf5f
--- /dev/null
+++ b/agent-service/src/types/workflow.ts
@@ -0,0 +1,144 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+interface LogicalPort {
+ readonly operatorID: string;
+ readonly portID: string;
+}
+
+interface PortIdentity {
+ readonly id: number;
+ readonly internal: boolean;
+}
+
+type PartitionInfo =
+ | { readonly type: "hash"; readonly hashAttributeNames: string[] }
+ | {
+ readonly type: "range";
+ readonly rangeAttributeNames: string[];
+ readonly rangeMin: number;
+ readonly rangeMax: number;
+ }
+ | { readonly type: "single" }
+ | { readonly type: "broadcast" }
+ | { readonly type: "none" };
+
+export interface PortDescription {
+ readonly portID: string;
+ readonly displayName?: string;
+ readonly disallowMultiInputs?: boolean;
+ readonly isDynamicPort?: boolean;
+ readonly partitionRequirement?: PartitionInfo;
+ readonly dependencies?: { id: number; internal: boolean }[];
+}
+
+export interface OperatorPredicate {
+ readonly operatorID: string;
+ readonly operatorType: string;
+ readonly operatorVersion: string;
+ readonly operatorProperties: Record;
+ readonly inputPorts: PortDescription[];
+ readonly outputPorts: PortDescription[];
+ readonly dynamicInputPorts?: boolean;
+ readonly dynamicOutputPorts?: boolean;
+ readonly showAdvanced: boolean;
+ readonly isDisabled?: boolean;
+ readonly viewResult?: boolean;
+ readonly markedForReuse?: boolean;
+ readonly customDisplayName?: string;
+}
+
+export interface LogicalOperator {
+ readonly operatorID: string;
+ readonly operatorType: string;
+ readonly [key: string]: any;
+}
+
+export interface OperatorLink {
+ readonly linkID: string;
+ readonly source: LogicalPort;
+ readonly target: LogicalPort;
+}
+
+export interface LogicalLink {
+ readonly fromOpId: string;
+ readonly fromPortId: PortIdentity;
+ readonly toOpId: string;
+ readonly toPortId: PortIdentity;
+}
+
+export interface LogicalPlan {
+ readonly operators: LogicalOperator[];
+ readonly links: LogicalLink[];
+ readonly opsToViewResult?: string[];
+ readonly opsToReuseResult?: string[];
+}
+
+export interface Point {
+ readonly x: number;
+ readonly y: number;
+}
+
+export interface CommentBox {
+ readonly commentBoxID: string;
+ readonly comments: string;
+ readonly x: number;
+ readonly y: number;
+ readonly width: number;
+ readonly height: number;
+}
+
+export interface WorkflowSettings {
+ readonly dataTransferBatchSize: number;
+}
+
+export interface WorkflowContent {
+ readonly operators: OperatorPredicate[];
+ readonly operatorPositions: { [key: string]: Point };
+ readonly links: OperatorLink[];
+ readonly commentBoxes: CommentBox[];
+ readonly settings: WorkflowSettings;
+}
+
+type AttributeType = "string" | "integer" | "double" | "boolean" | "long" | "timestamp" | "binary";
+
+export interface SchemaAttribute {
+ readonly attributeName: string;
+ readonly attributeType: AttributeType;
+}
+
+export type PortSchema = readonly SchemaAttribute[];
+
+export type OperatorPortSchemaMap = Record;
+
+export interface OperatorDetail {
+ operatorId: string;
+ operatorType: string;
+ customDisplayName?: string;
+ operatorProperties: Record;
+ inputPorts: PortDescription[];
+ outputPorts: PortDescription[];
+}
+
+export type ValidationError = {
+ isValid: false;
+ messages: Record;
+};
+
+export type Validation = { isValid: true } | ValidationError;
diff --git a/agent-service/tsconfig.json b/agent-service/tsconfig.json
new file mode 100644
index 00000000000..5f0da1e5bee
--- /dev/null
+++ b/agent-service/tsconfig.json
@@ -0,0 +1,23 @@
+{
+ "compilerOptions": {
+ "target": "ESNext",
+ "module": "Preserve",
+ "moduleResolution": "bundler",
+ "lib": ["ESNext"],
+ "types": ["bun-types"],
+
+ "strict": true,
+ "noImplicitOverride": true,
+
+ "esModuleInterop": true,
+ "forceConsistentCasingInFileNames": true,
+ "skipLibCheck": true,
+ "resolveJsonModule": true,
+ "isolatedModules": true,
+ "verbatimModuleSyntax": true,
+ "allowImportingTsExtensions": true,
+
+ "noEmit": true
+ },
+ "include": ["src/**/*"]
+}
diff --git a/amber/.scalafix.conf b/amber/.scalafix.conf
new file mode 100644
index 00000000000..238028c0ce4
--- /dev/null
+++ b/amber/.scalafix.conf
@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+rules = [
+ ProcedureSyntax,
+ RemoveUnused,
+]
+RemoveUnused.imports = true
+RemoveUnused.privates = true
+RemoveUnused.locals = false
+RemoveUnused.patternvars = false
+RemoveUnused.params = false
\ No newline at end of file
diff --git a/amber/.scalafmt.conf b/amber/.scalafmt.conf
new file mode 100644
index 00000000000..3c5f3019cf2
--- /dev/null
+++ b/amber/.scalafmt.conf
@@ -0,0 +1,24 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+version=2.6.4
+maxColumn = 100
+
+project.excludeFilters = [
+ "src/main/scala/com/kjetland/.*",
+ "src/main/scalapb/.*"
+]
diff --git a/amber/DESCRIPTION b/amber/DESCRIPTION
new file mode 100644
index 00000000000..59723d32ea5
--- /dev/null
+++ b/amber/DESCRIPTION
@@ -0,0 +1,16 @@
+Package: Texera-R-UDF
+Title: Required Libraries for R UDF
+Version: 1.0.0
+Authors@R: person("Texera", "Team", role = c("aut","cre"))
+Description: Below are the required libraries that should be installed to your R installation
+ before you begin to use/develop R UDF. Additionally, the version of R that you should be using
+ is also listed below. This package should also be used by the GitHub Actions Workflow files.
+URL: https://github.com/Texera/texera/
+BugReports: https://github.com/Texera/texera/
+Depends: R (>= 4.3.3)
+Imports:
+ arrow (>= 21.0.0),
+ coro (>= 1.0.4),
+ dplyr,
+ reticulate
+License: GPL (>= 2)
\ No newline at end of file
diff --git a/amber/README.md b/amber/README.md
new file mode 100644
index 00000000000..8b137891791
--- /dev/null
+++ b/amber/README.md
@@ -0,0 +1 @@
+
diff --git a/amber/build.sbt b/amber/build.sbt
new file mode 100644
index 00000000000..8b35af74497
--- /dev/null
+++ b/amber/build.sbt
@@ -0,0 +1,285 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+name := "amber"
+
+
+enablePlugins(JavaAppPackaging)
+
+// Ship LICENSE-binary, NOTICE-binary, DISCLAIMER-WIP, and the licenses/
+// directory at the top of the Universal dist zip.
+// See project/AddMetaInfLicenseFiles.scala.
+Universal / mappings := AddMetaInfLicenseFiles.distMappings(
+ (Universal / mappings).value,
+ (ThisBuild / baseDirectory).value
+)
+
+semanticdbEnabled := true
+semanticdbVersion := scalafixSemanticdb.revision
+
+// to turn on, use: INFO
+// to turn off, use: WARNING
+scalacOptions ++= Seq("-Xelide-below", "WARNING")
+
+// to check feature warnings
+scalacOptions += "-feature"
+// to check deprecation warnings
+scalacOptions += "-deprecation"
+// to check unused imports
+scalacOptions += "-Ywarn-unused:imports"
+
+conflictManager := ConflictManager.latestRevision
+
+// ensuring no parallel execution of multiple tasks
+concurrentRestrictions in Global += Tags.limit(Tags.Test, 1)
+
+// add python as an additional source
+Compile / unmanagedSourceDirectories += baseDirectory.value / "src" / "main" / "python"
+
+// Excluding some proto files:
+PB.generate / excludeFilter := "scalapb.proto"
+
+/////////////////////////////////////////////////////////////////////////////
+// Pekko related
+val pekkoVersion = "1.2.1"
+val pekkoDependencies = Seq(
+ "org.apache.pekko" %% "pekko-actor" % pekkoVersion,
+ "org.apache.pekko" %% "pekko-remote" % pekkoVersion,
+ "org.apache.pekko" %% "pekko-cluster" % pekkoVersion,
+ "org.apache.pekko" %% "pekko-cluster-metrics" % pekkoVersion,
+ "org.apache.pekko" %% "pekko-cluster-tools" % pekkoVersion,
+ "org.apache.pekko" %% "pekko-multi-node-testkit" % pekkoVersion % Test,
+ "org.apache.pekko" %% "pekko-testkit" % pekkoVersion % Test,
+ "org.apache.pekko" %% "pekko-persistence" % pekkoVersion,
+ "io.kamon" % "sigar-loader" % "1.6.6-rev002",
+ "com.softwaremill.macwire" %% "macros" % "2.6.7" % Provided,
+ "com.softwaremill.macwire" %% "macrospekko" % "2.6.7" % Provided,
+ "com.softwaremill.macwire" %% "util" % "2.6.7",
+ "com.softwaremill.macwire" %% "proxy" % "2.6.7",
+ "org.apache.pekko" %% "pekko-slf4j" % pekkoVersion,
+ "ch.qos.logback" % "logback-classic" % "1.2.13" % Test
+)
+
+// dropwizard web framework
+
+/////////////////////////////////////////////////////////////////////////////
+// DropWizard server related
+val dropwizardVersion = "1.3.23"
+
+val dropwizardDependencies = Seq(
+ "io.dropwizard" % "dropwizard-core" % dropwizardVersion,
+ "io.dropwizard" % "dropwizard-client" % dropwizardVersion,
+ "io.dropwizard" % "dropwizard-auth" % dropwizardVersion,
+ // https://mvnrepository.com/artifact/com.github.toastshaman/dropwizard-auth-jwt
+ "com.github.toastshaman" % "dropwizard-auth-jwt" % "1.1.2-0",
+ "com.github.dirkraft.dropwizard" % "dropwizard-file-assets" % "0.0.2",
+ "io.dropwizard-bundles" % "dropwizard-redirect-bundle" % "1.0.5",
+ "com.liveperson" % "dropwizard-websockets" % "1.3.14",
+ // https://mvnrepository.com/artifact/commons-io/commons-io
+ "commons-io" % "commons-io" % "2.15.1"
+)
+
+
+val jacksonVersion = "2.18.6"
+val mbknorJacksonJsonSchemaDependencies = Seq(
+ "com.fasterxml.jackson.core" % "jackson-databind" % jacksonVersion,
+ "javax.validation" % "validation-api" % "2.0.1.Final",
+ "org.slf4j" % "slf4j-api" % "1.7.26",
+ "io.github.classgraph" % "classgraph" % "4.8.157",
+ "ch.qos.logback" % "logback-classic" % "1.2.13" % "test",
+ "com.github.java-json-tools" % "json-schema-validator" % "2.2.14" % "test",
+ "com.fasterxml.jackson.module" % "jackson-module-kotlin" % jacksonVersion % "test",
+ "com.fasterxml.jackson.datatype" % "jackson-datatype-jdk8" % jacksonVersion % "test",
+ "com.fasterxml.jackson.datatype" % "jackson-datatype-jsr310" % jacksonVersion % "test",
+ "joda-time" % "joda-time" % "2.12.5" % "test",
+ "com.fasterxml.jackson.datatype" % "jackson-datatype-joda" % jacksonVersion % "test",
+ "com.fasterxml.jackson.module" % "jackson-module-jsonSchema" % jacksonVersion,
+ "com.fasterxml.jackson.module" %% "jackson-module-scala" % jacksonVersion,
+ // https://mvnrepository.com/artifact/com.fasterxml.jackson.module/jackson-module-no-ctor-deser
+ "com.fasterxml.jackson.module" % "jackson-module-no-ctor-deser" % jacksonVersion,
+)
+
+/////////////////////////////////////////////////////////////////////////////
+// Lucene related
+val luceneVersion = "8.7.0"
+val luceneDependencies = Seq(
+ "org.apache.lucene" % "lucene-core" % luceneVersion,
+ "org.apache.lucene" % "lucene-queryparser" % luceneVersion,
+ "org.apache.lucene" % "lucene-queries" % luceneVersion,
+ "org.apache.lucene" % "lucene-memory" % luceneVersion
+)
+
+/////////////////////////////////////////////////////////////////////////////
+// Hadoop related
+val hadoopVersion = "3.3.3"
+val excludeHadoopJersey = ExclusionRule(organization = "com.sun.jersey")
+val excludeHadoopSlf4j = ExclusionRule(organization = "org.slf4j")
+val excludeHadoopJetty = ExclusionRule(organization = "org.eclipse.jetty")
+val excludeHadoopJsp = ExclusionRule(organization = "javax.servlet.jsp")
+val hadoopDependencies = Seq(
+ "org.apache.hadoop" % "hadoop-common" % hadoopVersion excludeAll(excludeHadoopJersey, excludeHadoopSlf4j, excludeHadoopJsp, excludeHadoopJetty)
+)
+
+/////////////////////////////////////////////////////////////////////////////
+// Google Service related
+val googleServiceDependencies = Seq(
+ "com.google.oauth-client" % "google-oauth-client-jetty" % "1.34.1" exclude("com.google.guava", "guava"),
+ "com.google.api-client" % "google-api-client" % "2.2.0" exclude("com.google.guava", "guava"),
+ "com.sun.mail" % "javax.mail" % "1.6.2"
+)
+
+libraryDependencies ++= pekkoDependencies
+libraryDependencies ++= luceneDependencies
+libraryDependencies ++= dropwizardDependencies
+libraryDependencies ++= mbknorJacksonJsonSchemaDependencies
+libraryDependencies ++= googleServiceDependencies
+libraryDependencies ++= hadoopDependencies
+
+/////////////////////////////////////////////////////////////////////////////
+// protobuf related
+// run the following with sbt to have protobuf codegen
+
+PB.protocVersion := "3.19.4"
+
+enablePlugins(Fs2Grpc)
+
+Compile / PB.targets := Seq(
+ scalapb.gen(
+ singleLineToProtoString = true
+ ) -> (Compile / sourceManaged).value,
+ // let fs2 compile grpc-related proto, skip other protos in fs2 compilation pipeline.
+ scalapbCodeGenerators.value(1)
+)
+
+libraryDependencies ++= Seq(
+ "com.thesamet.scalapb" %% "scalapb-runtime" % scalapb.compiler.Version.scalapbVersion % "protobuf"
+)
+// For ScalaPB 0.11.x:
+libraryDependencies += "com.thesamet.scalapb" %% "scalapb-json4s" % "0.12.0"
+
+// enable protobuf compilation in Test
+Test / PB.protoSources += PB.externalSourcePath.value
+
+/////////////////////////////////////////////////////////////////////////////
+// Test related
+// https://mvnrepository.com/artifact/org.scalamock/scalamock
+libraryDependencies += "org.scalamock" %% "scalamock" % "5.2.0" % Test
+// https://mvnrepository.com/artifact/ch.vorburger.mariaDB4j/mariaDB4j
+libraryDependencies += "ch.vorburger.mariaDB4j" % "mariaDB4j" % "2.4.0" % Test
+// https://www.scalatest.org/getting_started_with_fun_suite
+libraryDependencies += "org.scalatest" %% "scalatest" % "3.2.15" % Test
+// JUnit related dependencies
+libraryDependencies += "junit" % "junit" % "4.13.2" % Test // JUnit dependency for Java tests
+libraryDependencies += "com.novocode" % "junit-interface" % "0.11" % Test // SBT interface for JUnit
+
+/////////////////////////////////////////////////////////////////////////////
+// Workflow version control related
+// https://mvnrepository.com/artifact/com.flipkart.zjsonpatch/zjsonpatch
+libraryDependencies += "com.flipkart.zjsonpatch" % "zjsonpatch" % "0.4.13"
+
+/////////////////////////////////////////////////////////////////////////////
+// Uncategorized
+
+// https://mvnrepository.com/artifact/io.reactivex.rxjava3/rxjava
+libraryDependencies += "io.reactivex.rxjava3" % "rxjava" % "3.1.6"
+
+// https://mvnrepository.com/artifact/org.postgresql/postgresql
+libraryDependencies += "org.postgresql" % "postgresql" % "42.5.4"
+
+// https://mvnrepository.com/artifact/com.typesafe.scala-logging/scala-logging
+libraryDependencies += "com.typesafe.scala-logging" %% "scala-logging" % "3.9.5"
+
+// https://mvnrepository.com/artifact/org.scalactic/scalactic
+libraryDependencies += "org.scalactic" %% "scalactic" % "3.2.15"
+
+// https://mvnrepository.com/artifact/com.github.tototoshi/scala-csv
+libraryDependencies += "com.github.tototoshi" %% "scala-csv" % "1.3.10"
+
+// https://mvnrepository.com/artifact/com.univocity/univocity-parsers
+libraryDependencies += "com.univocity" % "univocity-parsers" % "2.9.1"
+
+// https://mvnrepository.com/artifact/com.konghq/unirest-java
+libraryDependencies += "com.konghq" % "unirest-java" % "3.14.2"
+
+// https://mvnrepository.com/artifact/com.github.marianobarrios/lbmq
+libraryDependencies += "com.github.marianobarrios" % "lbmq" % "0.6.0"
+
+// https://mvnrepository.com/artifact/org.jooq/jooq
+libraryDependencies += "org.jooq" % "jooq" % "3.14.16"
+
+// https://mvnrepository.com/artifact/org.jgrapht/jgrapht-core
+libraryDependencies += "org.jgrapht" % "jgrapht-core" % "1.4.0"
+
+// https://mvnrepository.com/artifact/com.esotericsoftware/kryo
+libraryDependencies += "com.esotericsoftware" % "kryo" % "5.6.2"
+libraryDependencies += "com.esotericsoftware" % "kryo5" % "5.6.0"
+
+// https://mvnrepository.com/artifact/io.altoo/pekko-kryo-serialization
+libraryDependencies += "io.altoo" %% "pekko-kryo-serialization" % "1.3.0"
+
+// https://mvnrepository.com/artifact/io.altoo/scala-kryo-serialization
+libraryDependencies += "io.altoo" %% "scala-kryo-serialization" % "1.3.0"
+
+// https://mvnrepository.com/artifact/com.twitter/util-core
+libraryDependencies += "com.twitter" %% "util-core" % "22.12.0"
+
+// https://mvnrepository.com/artifact/com.typesafe.play/play-json
+libraryDependencies += "com.typesafe.play" %% "play-json" % "2.9.4"
+
+// https://mvnrepository.com/artifact/org.fusesource.leveldbjni/leveldbjni-all
+libraryDependencies += "org.fusesource.leveldbjni" % "leveldbjni-all" % "1.8"
+
+// https://mvnrepository.com/artifact/com.github.nscala-time/nscala-time
+libraryDependencies += "com.github.nscala-time" %% "nscala-time" % "2.32.0"
+
+// https://mvnrepository.com/artifact/com.google.guava/guava
+libraryDependencies += "com.google.guava" % "guava" % "29.0-jre"
+
+// https://mvnrepository.com/artifact/org.tukaani/xz
+libraryDependencies += "org.tukaani" % "xz" % "1.9"
+
+// https://mvnrepository.com/artifact/org.jasypt/jasypt
+libraryDependencies += "org.jasypt" % "jasypt" % "1.9.3"
+
+// Jgit library for tracking operator version
+// https://mvnrepository.com/artifact/org.eclipse.jgit/org.eclipse.jgit
+libraryDependencies += "org.eclipse.jgit" % "org.eclipse.jgit" % "5.13.0.202109080827-r"
+
+// https://mvnrepository.com/artifact/org.ehcache/sizeof
+libraryDependencies += "org.ehcache" % "sizeof" % "0.4.3"
+
+// https://mvnrepository.com/artifact/org.mindrot/jbcrypt
+libraryDependencies += "org.mindrot" % "jbcrypt" % "0.4"
+
+// https://mvnrepository.com/artifact/com.github.sisyphsu/dateparser
+libraryDependencies += "com.github.sisyphsu" % "dateparser" % "1.0.11"
+
+// https://mvnrepository.com/artifact/org.apache.commons/commons-vfs2
+libraryDependencies += "org.apache.commons" % "commons-vfs2" % "2.9.0"
+
+// https://mvnrepository.com/artifact/org.apache.commons/commons-jcs3-core
+libraryDependencies += "org.apache.commons" % "commons-jcs3-core" % "3.2"
+
+// For supporting MultiDict
+// https://mvnrepository.com/artifact/org.scala-lang.modules/scala-collection-contrib
+libraryDependencies += "org.scala-lang.modules" %% "scala-collection-contrib" % "0.3.0"
+
+// For supporting deepcopy
+// https://mvnrepository.com/artifact/io.github.kostaskougios/cloning
+libraryDependencies += "io.github.kostaskougios" % "cloning" % "1.10.3"
+
+
diff --git a/amber/operator-requirements.txt b/amber/operator-requirements.txt
new file mode 100644
index 00000000000..a7296394662
--- /dev/null
+++ b/amber/operator-requirements.txt
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+wordcloud==1.9.3
+plotly==5.24.1
+praw==7.6.1
+pillow==12.1.1
+pybase64==1.3.2
+
+# Pin torch to the CPU wheel on Linux x86_64 to avoid the NVIDIA CUDA deps.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.8.0+cpu ; platform_system == "Linux" and platform_machine == "x86_64"
+torch==2.8.0 ; platform_system != "Linux" or platform_machine != "x86_64"
+
+scikit-learn==1.5.0
+transformers==4.57.3
+boto3==1.40.53
+scikit-image==0.25.2
diff --git a/amber/project/build.properties b/amber/project/build.properties
new file mode 100644
index 00000000000..eaefbb622ed
--- /dev/null
+++ b/amber/project/build.properties
@@ -0,0 +1,18 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+sbt.version=1.12.9
\ No newline at end of file
diff --git a/amber/project/plugins.sbt b/amber/project/plugins.sbt
new file mode 100644
index 00000000000..d7742bf2da2
--- /dev/null
+++ b/amber/project/plugins.sbt
@@ -0,0 +1,23 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.4.2")
+addSbtPlugin("ch.epfl.scala" % "sbt-scalafix" % "0.14.6")
+
+addSbtPlugin("com.github.sbt" % "sbt-native-packager" % "1.11.1")
+// for scalapb code gen
+addSbtPlugin("org.typelevel" % "sbt-fs2-grpc" % "2.11.0")
diff --git a/amber/requirements.txt b/amber/requirements.txt
new file mode 100644
index 00000000000..72ac4da75f2
--- /dev/null
+++ b/amber/requirements.txt
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+wheel==0.41.2
+setuptools==80.10.2
+numpy==2.1.0
+pandas==2.2.3
+ruff==0.14.7
+iniconfig==1.1.1
+loguru==0.7.0
+pyarrow==21.0.0
+pytest==7.4.0
+python-dateutil==2.8.2
+pytest-timeout==2.2.0
+protobuf==4.25.8
+betterproto==2.0.0b7
+pampy==0.3.0
+overrides==7.4.0
+typing_extensions==4.14.1
+pytest-reraise==2.1.2
+Deprecated==1.2.14
+fs==2.4.16
+praw==7.6.1
+bidict==0.22.0
+cached_property==1.5.2
+psutil==5.9.0
+tzlocal==2.1
+s3fs==2025.9.0
+aiobotocore==2.25.1
+botocore==1.40.53
+pyiceberg==0.11.1
+readerwriterlock==1.0.9
+tenacity==8.5.0
+SQLAlchemy==2.0.37
+pg8000==1.31.5
+pympler==1.1
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto
new file mode 100644
index 00000000000..b22c1bdf7c2
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto
@@ -0,0 +1,275 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+package org.apache.texera.amber.engine.architecture.rpc;
+
+import "org/apache/texera/amber/core/virtualidentity.proto";
+import "org/apache/texera/amber/core/workflow.proto";
+import "org/apache/texera/amber/core/executor.proto";
+import "org/apache/texera/amber/engine/architecture/worker/statistics.proto";
+import "org/apache/texera/amber/engine/architecture/sendsemantics/partitionings.proto";
+import "scalapb/scalapb.proto";
+import "google/protobuf/timestamp.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+};
+
+message ControlRequest {
+ oneof sealed_value {
+ // request for controller
+ PropagateEmbeddedControlMessageRequest propagateEmbeddedControlMessageRequest = 1;
+ TakeGlobalCheckpointRequest takeGlobalCheckpointRequest = 2;
+ DebugCommandRequest debugCommandRequest = 3;
+ EvaluatePythonExpressionRequest evaluatePythonExpressionRequest = 4;
+ RetryWorkflowRequest retryWorkflowRequest = 5;
+ ConsoleMessageTriggeredRequest consoleMessageTriggeredRequest = 6;
+ PortCompletedRequest portCompletedRequest = 7;
+ WorkerStateUpdatedRequest workerStateUpdatedRequest = 8;
+ LinkWorkersRequest linkWorkersRequest = 9;
+ WorkflowReconfigureRequest workflowReconfigureRequest = 10;
+
+ // request for worker
+ AddInputChannelRequest addInputChannelRequest = 50;
+ AddPartitioningRequest addPartitioningRequest = 51;
+ AssignPortRequest assignPortRequest = 52;
+ FinalizeCheckpointRequest finalizeCheckpointRequest = 53;
+ InitializeExecutorRequest initializeExecutorRequest = 54;
+ UpdateExecutorRequest updateExecutorRequest = 55;
+ EmptyRequest emptyRequest = 56;
+ PrepareCheckpointRequest prepareCheckpointRequest = 57;
+ QueryStatisticsRequest queryStatisticsRequest = 58;
+
+ // request for testing
+ Ping ping = 100;
+ Pong pong = 101;
+ Nested nested = 102;
+ Pass pass = 103;
+ ErrorCommand errorCommand = 104;
+ Recursion recursion = 105;
+ Collect collect = 106;
+ GenerateNumber generateNumber = 107;
+ MultiCall multiCall = 108;
+ Chain chain = 109;
+ }
+}
+
+message EmptyRequest{}
+
+message AsyncRPCContext {
+ option (scalapb.message).no_box = true;
+ core.ActorVirtualIdentity sender = 1 [(scalapb.field).no_box = true];
+ core.ActorVirtualIdentity receiver = 2 [(scalapb.field).no_box = true];
+}
+
+message ControlInvocation {
+ option (scalapb.message).extends = "org.apache.texera.amber.engine.common.ambermessage.DirectControlMessagePayload";
+ string methodName = 1;
+ ControlRequest command = 2 [(scalapb.field).no_box = true];
+ AsyncRPCContext context = 3;
+ int64 commandId = 4;
+}
+
+enum EmbeddedControlMessageType {
+ ALL_ALIGNMENT = 0;
+ NO_ALIGNMENT = 1;
+ PORT_ALIGNMENT = 2;
+}
+
+message EmbeddedControlMessage {
+ option (scalapb.message).extends = "org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessagePayload";
+ core.EmbeddedControlMessageIdentity id = 1 [(scalapb.field).no_box = true];
+ EmbeddedControlMessageType ecm_type = 2;
+ repeated core.ChannelIdentity scope = 3;
+ map commandMapping = 4;
+}
+
+message PropagateEmbeddedControlMessageRequest {
+ repeated core.PhysicalOpIdentity sourceOpToStartProp = 1;
+ core.EmbeddedControlMessageIdentity id = 2 [(scalapb.field).no_box = true];
+ EmbeddedControlMessageType ecm_type = 3;
+ repeated core.PhysicalOpIdentity scope = 4;
+ repeated core.PhysicalOpIdentity targetOps = 5;
+ ControlRequest command = 6;
+ string methodName = 7;
+}
+
+message TakeGlobalCheckpointRequest {
+ bool estimationOnly = 1;
+ core.EmbeddedControlMessageIdentity checkpointId = 2 [(scalapb.field).no_box = true];
+ string destination = 3;
+}
+
+message WorkflowReconfigureRequest{
+ repeated UpdateExecutorRequest reconfiguration = 1;
+ string reconfigurationId = 2;
+}
+
+
+message DebugCommandRequest {
+ string workerId = 1;
+ string cmd = 2;
+}
+
+message EvaluatePythonExpressionRequest {
+ string expression = 1;
+ string operatorId = 2;
+}
+
+message RetryWorkflowRequest {
+ repeated core.ActorVirtualIdentity workers = 1;
+}
+
+enum ConsoleMessageType{
+ PRINT = 0;
+ ERROR = 1;
+ COMMAND = 2;
+ DEBUGGER = 3;
+}
+
+message ConsoleMessage {
+ option (scalapb.message).extends = "org.apache.texera.amber.engine.architecture.controller.ClientEvent";
+ string worker_id = 1;
+ google.protobuf.Timestamp timestamp = 2 [(scalapb.field).no_box = true];
+ ConsoleMessageType msg_type = 3;
+ string source = 4;
+ string title = 5;
+ string message = 6;
+}
+
+message ConsoleMessageTriggeredRequest {
+ ConsoleMessage consoleMessage = 1 [(scalapb.field).no_box = true];
+}
+
+message PortCompletedRequest {
+ core.PortIdentity portId = 1 [(scalapb.field).no_box = true];
+ bool input = 2;
+}
+
+message WorkerStateUpdatedRequest {
+ worker.WorkerState state = 1 [(scalapb.field).no_box = true];
+}
+
+message LinkWorkersRequest {
+ core.PhysicalLink link = 1 [(scalapb.field).no_box = true];
+}
+
+// Ping message
+message Ping {
+ int32 i = 1;
+ int32 end = 2;
+ core.ActorVirtualIdentity to = 3 [(scalapb.field).no_box = true];
+}
+
+// Pong message
+message Pong {
+ int32 i = 1;
+ int32 end = 2;
+ core.ActorVirtualIdentity to = 3 [(scalapb.field).no_box = true];
+}
+
+// Pass message
+message Pass {
+ string value = 1;
+}
+
+// Nested message
+message Nested {
+ int32 k = 1;
+}
+
+// MultiCall message
+message MultiCall {
+ repeated core.ActorVirtualIdentity seq = 1;
+}
+
+// ErrorCommand message
+message ErrorCommand {
+}
+
+// Collect message
+message Collect {
+ repeated core.ActorVirtualIdentity workers = 1;
+}
+
+// GenerateNumber message
+message GenerateNumber {
+}
+
+// Chain message
+message Chain {
+ repeated core.ActorVirtualIdentity nexts = 1;
+}
+
+// Recursion message
+message Recursion {
+ int32 i = 1;
+}
+
+// Messages for the commands
+message AddInputChannelRequest {
+ core.ChannelIdentity channelId = 1 [(scalapb.field).no_box = true];
+ core.PortIdentity portId = 2 [(scalapb.field).no_box = true];
+}
+
+message AddPartitioningRequest {
+ core.PhysicalLink tag = 1 [(scalapb.field).no_box = true];
+ sendsemantics.Partitioning partitioning = 2 [(scalapb.field).no_box = true];
+}
+
+message AssignPortRequest {
+ core.PortIdentity portId = 1 [(scalapb.field).no_box = true];
+ bool input = 2;
+ map schema = 3;
+ repeated string storageUris = 4;
+ repeated sendsemantics.Partitioning partitionings = 5;
+}
+
+message FinalizeCheckpointRequest {
+ core.EmbeddedControlMessageIdentity checkpointId = 1 [(scalapb.field).no_box = true];
+ string writeTo = 2;
+}
+
+message InitializeExecutorRequest {
+ int32 totalWorkerCount = 1;
+ core.OpExecInitInfo opExecInitInfo = 2;
+ bool isSource = 3;
+}
+
+message UpdateExecutorRequest {
+ core.PhysicalOpIdentity targetOpId = 1 [(scalapb.field).no_box = true];
+ core.OpExecInitInfo newExecInitInfo = 2;
+}
+
+message PrepareCheckpointRequest{
+ core.EmbeddedControlMessageIdentity checkpointId = 1 [(scalapb.field).no_box = true];
+ bool estimationOnly = 2;
+}
+
+enum StatisticsUpdateTarget {
+ BOTH_UI_AND_PERSISTENCE = 0;
+ UI_ONLY = 1;
+ PERSISTENCE_ONLY = 2;
+}
+
+message QueryStatisticsRequest{
+ repeated core.ActorVirtualIdentity filterByWorkers = 1;
+ StatisticsUpdateTarget updateTarget = 2;
+}
\ No newline at end of file
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto
new file mode 100644
index 00000000000..27b4727ee98
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto
@@ -0,0 +1,49 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+package org.apache.texera.amber.engine.architecture.rpc;
+
+import "org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto";
+import "org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto";
+import "scalapb/scalapb.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+};
+
+
+service ControllerService {
+ rpc RetrieveWorkflowState(EmptyRequest) returns (RetrieveWorkflowStateResponse);
+ rpc PropagateEmbeddedControlMessage(PropagateEmbeddedControlMessageRequest) returns (PropagateEmbeddedControlMessageResponse);
+ rpc TakeGlobalCheckpoint(TakeGlobalCheckpointRequest) returns (TakeGlobalCheckpointResponse);
+ rpc DebugCommand(DebugCommandRequest) returns (EmptyReturn);
+ rpc EvaluatePythonExpression(EvaluatePythonExpressionRequest) returns (EvaluatePythonExpressionResponse);
+ rpc ConsoleMessageTriggered(ConsoleMessageTriggeredRequest) returns (EmptyReturn);
+ rpc PortCompleted(PortCompletedRequest) returns (EmptyReturn);
+ rpc StartWorkflow(EmptyRequest) returns (StartWorkflowResponse);
+ rpc ResumeWorkflow(EmptyRequest) returns (EmptyReturn);
+ rpc PauseWorkflow(EmptyRequest) returns (EmptyReturn);
+ rpc WorkerStateUpdated(WorkerStateUpdatedRequest) returns (EmptyReturn);
+ rpc WorkerExecutionCompleted(EmptyRequest) returns (EmptyReturn);
+ rpc LinkWorkers(LinkWorkersRequest) returns (EmptyReturn);
+ rpc ControllerInitiateQueryStatistics(QueryStatisticsRequest) returns (EmptyReturn);
+ rpc RetryWorkflow(RetryWorkflowRequest) returns (EmptyReturn);
+ rpc ReconfigureWorkflow(WorkflowReconfigureRequest) returns (EmptyReturn);
+}
\ No newline at end of file
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto
new file mode 100644
index 00000000000..43613b5cfdc
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto
@@ -0,0 +1,141 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+package org.apache.texera.amber.engine.architecture.rpc;
+
+import "org/apache/texera/amber/engine/architecture/worker/statistics.proto";
+import "scalapb/scalapb.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+};
+
+
+// The generic return message
+message ControlReturn {
+ // Oneof block for various return types
+ oneof sealed_value {
+ // controller responses
+ RetrieveWorkflowStateResponse retrieveWorkflowStateResponse = 1;
+ PropagateEmbeddedControlMessageResponse propagateEmbeddedControlMessageResponse = 2;
+ TakeGlobalCheckpointResponse takeGlobalCheckpointResponse = 3;
+ EvaluatePythonExpressionResponse evaluatePythonExpressionResponse = 4;
+ StartWorkflowResponse startWorkflowResponse = 5;
+
+ // worker responses
+ WorkerStateResponse workerStateResponse = 50;
+ WorkerMetricsResponse workerMetricsResponse = 51;
+ FinalizeCheckpointResponse finalizeCheckpointResponse = 52;
+
+ // common responses
+ ControlError controlError = 101;
+ EmptyReturn emptyReturn = 102;
+ StringResponse stringResponse = 103;
+ IntResponse intResponse = 104;
+ }
+}
+
+message EmptyReturn {}
+
+enum ErrorLanguage {
+ PYTHON = 0;
+ SCALA = 1;
+}
+
+message ControlError {
+ string errorMessage = 1;
+ string errorDetails = 2;
+ string stackTrace = 3;
+ ErrorLanguage language = 4;
+}
+
+message ReturnInvocation {
+ option (scalapb.message).extends = "org.apache.texera.amber.engine.common.ambermessage.DirectControlMessagePayload";
+ int64 commandId = 1;
+ ControlReturn returnValue = 2 [(scalapb.field).no_box = true];
+}
+
+
+message StringResponse {
+ string value = 1;
+}
+
+message IntResponse {
+ int32 value = 1;
+}
+
+message RetrieveWorkflowStateResponse {
+ map state = 1;
+}
+
+message FinalizeCheckpointResponse {
+ int64 size = 1;
+}
+
+message PropagateEmbeddedControlMessageResponse {
+ map returns = 1;
+}
+
+message TakeGlobalCheckpointResponse {
+ int64 totalSize = 1;
+}
+
+message TypedValue {
+ string expression = 1;
+ string value_ref = 2;
+ string value_str = 3;
+ string value_type = 4;
+ bool expandable = 5;
+}
+
+message EvaluatedValue {
+ TypedValue value = 1;
+ repeated TypedValue attributes = 2;
+}
+
+message EvaluatePythonExpressionResponse {
+ repeated EvaluatedValue values = 1;
+}
+
+enum WorkflowAggregatedState {
+ UNINITIALIZED = 0;
+ READY = 1;
+ RUNNING = 2;
+ PAUSING = 3;
+ PAUSED = 4;
+ RESUMING = 5;
+ COMPLETED = 6;
+ FAILED = 7;
+ UNKNOWN = 8;
+ KILLED = 9;
+ TERMINATED = 10;
+}
+
+message StartWorkflowResponse {
+ WorkflowAggregatedState workflowState = 1 [(scalapb.field).no_box = true];
+}
+
+message WorkerStateResponse {
+ worker.WorkerState state = 1 [(scalapb.field).no_box = true];
+}
+
+message WorkerMetricsResponse {
+ worker.WorkerMetrics metrics = 1 [(scalapb.field).no_box = true];
+}
\ No newline at end of file
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/testerservice.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/testerservice.proto
new file mode 100644
index 00000000000..229b03f689c
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/testerservice.proto
@@ -0,0 +1,43 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+package org.apache.texera.amber.engine.architecture.rpc;
+
+import "org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto";
+import "org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto";
+import "scalapb/scalapb.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+};
+
+
+service RPCTester {
+ rpc SendPing(Ping) returns (IntResponse){}
+ rpc SendPong(Pong) returns (IntResponse){}
+ rpc SendNested(Nested) returns (StringResponse){}
+ rpc SendPass(Pass) returns (StringResponse){}
+ rpc SendErrorCommand(ErrorCommand) returns (StringResponse) {}
+ rpc SendRecursion(Recursion) returns (StringResponse) {}
+ rpc SendCollect(Collect) returns (StringResponse) {}
+ rpc SendGenerateNumber(GenerateNumber) returns (IntResponse) {}
+ rpc SendMultiCall(MultiCall) returns (StringResponse) {}
+ rpc SendChain(Chain) returns (StringResponse) {}
+}
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/workerservice.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/workerservice.proto
new file mode 100644
index 00000000000..21944ffefc6
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/workerservice.proto
@@ -0,0 +1,54 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+package org.apache.texera.amber.engine.architecture.rpc;
+
+import "org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto";
+import "org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto";
+import "scalapb/scalapb.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+};
+
+// RPC Service
+service WorkerService {
+ rpc AddInputChannel(AddInputChannelRequest) returns (EmptyReturn);
+ rpc AddPartitioning(AddPartitioningRequest) returns (EmptyReturn);
+ rpc AssignPort(AssignPortRequest) returns (EmptyReturn);
+ rpc FinalizeCheckpoint(FinalizeCheckpointRequest) returns (FinalizeCheckpointResponse);
+ rpc FlushNetworkBuffer(EmptyRequest) returns (EmptyReturn);
+ rpc InitializeExecutor(InitializeExecutorRequest) returns (EmptyReturn);
+ rpc OpenExecutor(EmptyRequest) returns (EmptyReturn);
+ rpc PauseWorker(EmptyRequest) returns (WorkerStateResponse);
+ rpc PrepareCheckpoint(PrepareCheckpointRequest) returns (EmptyReturn);
+ rpc QueryStatistics(EmptyRequest) returns (WorkerMetricsResponse);
+ rpc ResumeWorker(EmptyRequest) returns (WorkerStateResponse);
+ rpc RetrieveState(EmptyRequest) returns (EmptyReturn);
+ rpc RetryCurrentTuple(EmptyRequest) returns (EmptyReturn);
+ rpc StartWorker(EmptyRequest) returns (WorkerStateResponse);
+ rpc EndWorker(EmptyRequest) returns (EmptyReturn);
+ rpc StartChannel(EmptyRequest) returns (EmptyReturn);
+ rpc EndChannel(EmptyRequest) returns (EmptyReturn);
+ rpc DebugCommand(DebugCommandRequest) returns (EmptyReturn);
+ rpc EvaluatePythonExpression(EvaluatePythonExpressionRequest) returns (EvaluatedValue);
+ rpc NoOperation(EmptyRequest) returns (EmptyReturn);
+ rpc UpdateExecutor(UpdateExecutorRequest) returns (EmptyReturn);
+}
\ No newline at end of file
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/sendsemantics/partitionings.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/sendsemantics/partitionings.proto
new file mode 100644
index 00000000000..813a4041b31
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/sendsemantics/partitionings.proto
@@ -0,0 +1,68 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+
+package org.apache.texera.amber.engine.architecture.sendsemantics;
+
+import "org/apache/texera/amber/core/virtualidentity.proto";
+import "scalapb/scalapb.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+};
+
+message Partitioning{
+ oneof sealed_value{
+ OneToOnePartitioning oneToOnePartitioning = 1;
+ RoundRobinPartitioning roundRobinPartitioning = 2;
+ HashBasedShufflePartitioning hashBasedShufflePartitioning = 3;
+ RangeBasedShufflePartitioning rangeBasedShufflePartitioning = 4;
+ BroadcastPartitioning broadcastPartitioning = 5;
+ }
+}
+
+message OneToOnePartitioning{
+ int32 batchSize = 1;
+ repeated core.ChannelIdentity channels = 2;
+}
+
+message RoundRobinPartitioning{
+ int32 batchSize = 1;
+ repeated core.ChannelIdentity channels = 2;
+}
+
+message HashBasedShufflePartitioning{
+ int32 batchSize = 1;
+ repeated core.ChannelIdentity channels = 2;
+ repeated string hashAttributeNames = 3;
+}
+
+message RangeBasedShufflePartitioning {
+ int32 batchSize = 1;
+ repeated core.ChannelIdentity channels = 2;
+ repeated string rangeAttributeNames = 3;
+ int64 rangeMin = 4;
+ int64 rangeMax = 5;
+}
+
+message BroadcastPartitioning{
+ int32 batchSize = 1;
+ repeated core.ChannelIdentity channels = 2;
+}
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/worker/statistics.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/worker/statistics.proto
new file mode 100644
index 00000000000..85d1fcf4aaa
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/worker/statistics.proto
@@ -0,0 +1,63 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+
+package org.apache.texera.amber.engine.architecture.worker;
+
+import "org/apache/texera/amber/core/workflow.proto";
+import "scalapb/scalapb.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+
+};
+
+enum WorkerState {
+ UNINITIALIZED = 0;
+ READY = 1;
+ RUNNING = 2;
+ PAUSED = 3;
+ COMPLETED = 4;
+ TERMINATED = 5;
+
+}
+
+message PortTupleMetricsMapping {
+ core.PortIdentity port_id = 1 [(scalapb.field).no_box = true];
+ TupleMetrics tuple_metrics = 2 [(scalapb.field).no_box = true];
+}
+
+message TupleMetrics {
+ int64 count = 1;
+ int64 size = 2;
+}
+
+message WorkerStatistics {
+ repeated PortTupleMetricsMapping input_tuple_metrics = 1;
+ repeated PortTupleMetricsMapping output_tuple_metrics = 2;
+ int64 data_processing_time = 3;
+ int64 control_processing_time = 4;
+ int64 idle_time = 5;
+}
+
+message WorkerMetrics {
+ WorkerState worker_state = 1 [(scalapb.field).no_box = true];
+ WorkerStatistics worker_statistics = 2 [(scalapb.field).no_box = true];
+}
\ No newline at end of file
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/common/actormessage.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/common/actormessage.proto
new file mode 100644
index 00000000000..2a548e6f753
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/common/actormessage.proto
@@ -0,0 +1,47 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+
+package org.apache.texera.amber.engine.common;
+
+import "scalapb/scalapb.proto";
+
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+};
+
+message Backpressure {
+ bool enableBackpressure = 1;
+}
+
+message CreditUpdate {
+}
+
+message ActorCommand {
+ oneof sealed_value {
+ Backpressure backpressure = 1;
+ CreditUpdate creditUpdate = 2;
+ }
+}
+
+message PythonActorMessage {
+ ActorCommand payload = 1 [(scalapb.field).no_box = true];
+}
\ No newline at end of file
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/common/ambermessage.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/common/ambermessage.proto
new file mode 100644
index 00000000000..54bd36abfc4
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/common/ambermessage.proto
@@ -0,0 +1,48 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+
+package org.apache.texera.amber.engine.common;
+
+import "org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto";
+import "org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto";
+import "org/apache/texera/amber/core/virtualidentity.proto";
+import "scalapb/scalapb.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: true
+};
+
+message DirectControlMessagePayloadV2 {
+ oneof value {
+ architecture.rpc.ControlInvocation control_invocation = 1;
+ architecture.rpc.ReturnInvocation return_invocation = 2;
+ }
+}
+
+message PythonDataHeader {
+ core.ChannelIdentity tag = 1 [(scalapb.field).no_box = true];
+ string payload_type = 2;
+}
+
+message PythonControlMessage {
+ core.ChannelIdentity tag = 1 [(scalapb.field).no_box = true];
+ DirectControlMessagePayloadV2 payload = 2 [(scalapb.field).no_box = true];
+}
diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/common/executionruntimestate.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/common/executionruntimestate.proto
new file mode 100644
index 00000000000..e712b3adc8a
--- /dev/null
+++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/common/executionruntimestate.proto
@@ -0,0 +1,100 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto3";
+
+package org.apache.texera.amber.engine.common;
+
+import "org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto";
+import "org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto";
+import "org/apache/texera/amber/engine/architecture/worker/statistics.proto";
+import "org/apache/texera/amber/core/virtualidentity.proto";
+import "org/apache/texera/amber/core/workflowruntimestate.proto";
+import "scalapb/scalapb.proto";
+
+option (scalapb.options) = {
+ scope: FILE,
+ preserve_unknown_fields: false
+ no_default_values_in_constructor: false
+};
+
+
+message BreakpointFault{
+ message BreakpointTuple{
+ int64 id = 1;
+ bool is_input = 2;
+ repeated string tuple = 3;
+ }
+
+ string worker_name = 1;
+ BreakpointTuple faulted_tuple = 2;
+}
+
+message OperatorBreakpoints{
+ repeated BreakpointFault unresolved_breakpoints = 1;
+}
+
+message ExecutionBreakpointStore{
+ map operator_info = 1;
+}
+
+message EvaluatedValueList{
+ repeated architecture.rpc.EvaluatedValue values = 1;
+}
+
+message OperatorConsole{
+ repeated architecture.rpc.ConsoleMessage console_messages = 1;
+ map evaluate_expr_results = 2;
+}
+
+message ExecutionConsoleStore{
+ map operator_console = 1;
+}
+
+message OperatorWorkerMapping{
+ string operatorId = 1;
+ repeated string workerIds = 2;
+}
+
+message OperatorStatistics{
+ repeated architecture.worker.PortTupleMetricsMapping input_metrics = 1;
+ repeated architecture.worker.PortTupleMetricsMapping output_metrics = 2;
+ int32 num_workers = 3;
+ int64 data_processing_time = 4;
+ int64 control_processing_time = 5;
+ int64 idle_time = 6;
+}
+
+message OperatorMetrics{
+ architecture.rpc.WorkflowAggregatedState operator_state = 1 [(scalapb.field).no_box = true];
+ OperatorStatistics operator_statistics = 2 [(scalapb.field).no_box = true];
+}
+
+message ExecutionStatsStore {
+ int64 startTimeStamp = 1;
+ int64 endTimeStamp = 2;
+ map operator_info = 3;
+ repeated OperatorWorkerMapping operator_worker_mapping = 4;
+}
+
+
+message ExecutionMetadataStore{
+ architecture.rpc.WorkflowAggregatedState state = 1;
+ repeated core.WorkflowFatalError fatal_errors = 2;
+ core.ExecutionIdentity executionId = 3 [(scalapb.field).no_box = true];
+ bool is_recovering = 4;
+}
\ No newline at end of file
diff --git a/amber/src/main/python/core/__init__.py b/amber/src/main/python/core/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/architecture/__init__.py b/amber/src/main/python/core/architecture/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/architecture/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/architecture/handlers/__init__.py b/amber/src/main/python/core/architecture/handlers/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/architecture/handlers/actorcommand/actor_handler_base.py b/amber/src/main/python/core/architecture/handlers/actorcommand/actor_handler_base.py
new file mode 100644
index 00000000000..59a68312cf8
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/actorcommand/actor_handler_base.py
@@ -0,0 +1,30 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import ABC
+
+from core.models import InternalQueue
+from proto.org.apache.texera.amber.engine.common import ActorCommand
+
+
+class ActorCommandHandler(ABC):
+ cmd: ActorCommand = None
+
+ def __call__(
+ self, command: ActorCommand, input_queue: InternalQueue, *args, **kwargs
+ ) -> None:
+ pass
diff --git a/amber/src/main/python/core/architecture/handlers/actorcommand/backpressure_handler.py b/amber/src/main/python/core/architecture/handlers/actorcommand/backpressure_handler.py
new file mode 100644
index 00000000000..764348b1d13
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/actorcommand/backpressure_handler.py
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.actorcommand.actor_handler_base import (
+ ActorCommandHandler,
+)
+from core.models.internal_queue import DCMElement, InternalQueue
+from core.util import set_one_of
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ControlInvocation,
+ ControlRequest,
+ EmptyRequest,
+ AsyncRpcContext,
+)
+from proto.org.apache.texera.amber.engine.common import (
+ Backpressure,
+ DirectControlMessagePayloadV2,
+)
+
+
+class BackpressureHandler(ActorCommandHandler):
+ cmd = Backpressure
+
+ def __call__(
+ self, command: Backpressure, input_queue: InternalQueue, *args, **kwargs
+ ):
+ if command.enable_backpressure:
+ input_queue.disable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE)
+ else:
+ input_queue.enable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE)
+ input_queue.put(
+ DCMElement(
+ tag=ChannelIdentity(
+ ActorVirtualIdentity("self"), ActorVirtualIdentity("self"), True
+ ),
+ payload=set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ "NoOperation",
+ set_one_of(ControlRequest, EmptyRequest()),
+ AsyncRpcContext(),
+ -1,
+ ),
+ ),
+ )
+ )
+
+ return None
diff --git a/amber/src/main/python/core/architecture/handlers/actorcommand/credit_update_handler.py b/amber/src/main/python/core/architecture/handlers/actorcommand/credit_update_handler.py
new file mode 100644
index 00000000000..b7430b0fefb
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/actorcommand/credit_update_handler.py
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.actorcommand.actor_handler_base import (
+ ActorCommandHandler,
+)
+from core.models import InternalQueue
+from proto.org.apache.texera.amber.engine.common import CreditUpdate
+
+
+class CreditUpdateHandler(ActorCommandHandler):
+ cmd = CreditUpdate
+
+ def __call__(
+ self, command: CreditUpdate, input_queue: InternalQueue, *args, **kwargs
+ ):
+ # do nothing
+ return None
diff --git a/amber/src/main/python/core/architecture/handlers/control/__init__.py b/amber/src/main/python/core/architecture/handlers/control/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/architecture/handlers/control/add_input_channel_handler.py b/amber/src/main/python/core/architecture/handlers/control/add_input_channel_handler.py
new file mode 100644
index 00000000000..9d398210b52
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/add_input_channel_handler.py
@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ AddInputChannelRequest,
+)
+
+
+class AddInputChannelHandler(ControlHandler):
+ async def add_input_channel(self, req: AddInputChannelRequest) -> EmptyReturn:
+ if not req.channel_id.is_control:
+ # Explicitly set is_control to trigger lazy computation.
+ # If not set, it may be computed at different times,
+ # causing hash inconsistencies.
+ req.channel_id.is_control = False
+ self.context.input_manager.register_input(req.channel_id, req.port_id)
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/add_partitioning_handler.py b/amber/src/main/python/core/architecture/handlers/control/add_partitioning_handler.py
new file mode 100644
index 00000000000..b4a88d0fe12
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/add_partitioning_handler.py
@@ -0,0 +1,28 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ AddPartitioningRequest,
+)
+
+
+class AddPartitioningHandler(ControlHandler):
+ async def add_partitioning(self, req: AddPartitioningRequest) -> EmptyReturn:
+ self.context.output_manager.add_partitioning(req.tag, req.partitioning)
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/assign_port_handler.py b/amber/src/main/python/core/architecture/handlers/control/assign_port_handler.py
new file mode 100644
index 00000000000..73ebad26b31
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/assign_port_handler.py
@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.models import Schema
+from core.util.virtual_identity import get_from_actor_id_for_input_port_storage
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ AssignPortRequest,
+)
+
+
+class AssignPortHandler(ControlHandler):
+ async def assign_port(self, req: AssignPortRequest) -> EmptyReturn:
+ if req.input:
+ self.context.input_manager.add_input_port(
+ req.port_id,
+ Schema(raw_schema=req.schema),
+ req.storage_uris,
+ req.partitionings,
+ )
+ for uri in req.storage_uris:
+ to_actor_id = ActorVirtualIdentity(self.context.worker_id)
+ from_actor_id = get_from_actor_id_for_input_port_storage(
+ uri, to_actor_id
+ )
+ channel_id = ChannelIdentity(from_actor_id, to_actor_id, False)
+ self.context.input_manager.register_input(
+ channel_id=channel_id, port_id=req.port_id
+ )
+ else:
+ storage_uri = None
+ if len(req.storage_uris) > 0 and req.storage_uris[0]:
+ storage_uri = req.storage_uris[0]
+ self.context.output_manager.add_output_port(
+ req.port_id, Schema(raw_schema=req.schema), storage_uri
+ )
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/control_handler_base.py b/amber/src/main/python/core/architecture/handlers/control/control_handler_base.py
new file mode 100644
index 00000000000..9378a5e561b
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/control_handler_base.py
@@ -0,0 +1,23 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from proto.org.apache.texera.amber.engine.architecture.rpc import WorkerServiceBase
+
+
+class ControlHandler(WorkerServiceBase):
+ def __init__(self, context):
+ self.context = context
diff --git a/amber/src/main/python/core/architecture/handlers/control/debug_command_handler.py b/amber/src/main/python/core/architecture/handlers/control/debug_command_handler.py
new file mode 100644
index 00000000000..cf041b47ea6
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/debug_command_handler.py
@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.architecture.managers.context import Context
+from core.architecture.managers.pause_manager import PauseType
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ DebugCommandRequest,
+)
+
+
+class WorkerDebugCommandHandler(ControlHandler):
+ async def debug_command(self, req: DebugCommandRequest) -> EmptyReturn:
+ # translate the command with the context.
+ translated_command = self.translate_debug_command(req.cmd, self.context)
+
+ # send the translated command to debugger to consume later.
+ self.context.debug_manager.put_debug_command(translated_command)
+
+ # allow MainLoop to switch into DataProcessor.
+ self.context.pause_manager.resume(PauseType.USER_PAUSE)
+ self.context.pause_manager.resume(PauseType.EXCEPTION_PAUSE)
+ self.context.pause_manager.resume(PauseType.DEBUG_PAUSE)
+ return EmptyReturn()
+
+ @staticmethod
+ def translate_debug_command(command: str, context: Context) -> str:
+ """
+ Cleans up and translates a debug command into one pdb can consume.
+
+ For `b`/`break` with a numeric line target, the operator's UDF module
+ name is prepended so the breakpoint lands inside the user's code:
+ ``b 5`` becomes ``b my_udf:5``.
+
+ Three forms are passed through unchanged because pdb already accepts
+ them and the module rewrite would corrupt them:
+
+ - bare ``b`` / ``break`` with no args
+ - ``b `` (pdb resolves the symbol itself)
+ - ``b :`` (the user already specified a file)
+
+ :raises ValueError: if the command is empty/whitespace-only, or if a
+ ``b``/``break`` with a numeric target is issued before the
+ operator module has been initialized.
+ """
+ parts = command.strip().split()
+ if not parts:
+ raise ValueError("debug command cannot be empty")
+ debug_command, *debug_args = parts
+
+ is_break_with_lineno = (
+ debug_command in ("b", "break") and debug_args and debug_args[0].isdigit()
+ )
+ if is_break_with_lineno:
+ module_name = context.executor_manager.operator_module_name
+ if module_name is None:
+ raise ValueError(
+ "executor module not initialized; cannot set breakpoint"
+ )
+ translated = (
+ f"{debug_command} {module_name}:{debug_args[0]} "
+ f"{' '.join(debug_args[1:])}"
+ )
+ else:
+ translated = f"{debug_command} {' '.join(debug_args)}"
+
+ return translated.strip()
diff --git a/amber/src/main/python/core/architecture/handlers/control/end_channel_handler.py b/amber/src/main/python/core/architecture/handlers/control/end_channel_handler.py
new file mode 100644
index 00000000000..d60b3874754
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/end_channel_handler.py
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.models.internal_marker import EndChannel
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ EmptyRequest,
+)
+
+
+class EndChannelHandler(ControlHandler):
+ async def end_channel(self, req: EmptyRequest) -> EmptyReturn:
+ self.context.input_manager.complete_current_port(
+ self.context.current_input_channel_id
+ )
+ self.context.tuple_processing_manager.current_internal_marker = EndChannel()
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py b/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py
new file mode 100644
index 00000000000..434225a5f0a
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py
@@ -0,0 +1,49 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from loguru import logger
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.util import IQueue
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ EmptyRequest,
+)
+
+
+class EndWorkerHandler(ControlHandler):
+ """
+ The EndWorker control messages is needed to ensure all the other
+ control messages in a worker are processed before worker termination.
+ """
+
+ async def end_worker(self, req: EmptyRequest) -> EmptyReturn:
+ """
+ The response of EndWorker to the controller indicates that this worker
+ has finished not only the data processing logic, but also the processing
+ of all the control messages.
+ """
+ # Ensure this is really the last message.
+ input_queue: IQueue = self.context.input_queue
+ if not input_queue.is_empty():
+ logger.warning(
+ f"Received EndHandler before all messages are "
+ f"processed. Unprocessed messages: {input_queue.get()}"
+ )
+ assert input_queue.is_empty()
+ # Now we can safely acknowledge that this worker can be terminated.
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/evaluate_expression_handler.py b/amber/src/main/python/core/architecture/handlers/control/evaluate_expression_handler.py
new file mode 100644
index 00000000000..b5c63acc55c
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/evaluate_expression_handler.py
@@ -0,0 +1,40 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.util.expression_evaluator import ExpressionEvaluator
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EvaluatedValue,
+ EvaluatePythonExpressionRequest,
+)
+
+
+class EvaluateExpressionHandler(ControlHandler):
+ async def evaluate_python_expression(
+ self, req: EvaluatePythonExpressionRequest
+ ) -> EvaluatedValue:
+ runtime_context = {
+ r"self": self.context.executor_manager.executor,
+ r"tuple_": self.context.tuple_processing_manager.current_input_tuple,
+ r"input_": self.context.tuple_processing_manager.current_input_port_id,
+ }
+
+ evaluated_value: EvaluatedValue = ExpressionEvaluator.evaluate(
+ req.expression, runtime_context
+ )
+
+ return evaluated_value
diff --git a/amber/src/main/python/core/architecture/handlers/control/initialize_executor_handler.py b/amber/src/main/python/core/architecture/handlers/control/initialize_executor_handler.py
new file mode 100644
index 00000000000..2c2dc1ad3c7
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/initialize_executor_handler.py
@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.util import get_one_of
+from proto.org.apache.texera.amber.core import OpExecWithCode
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ InitializeExecutorRequest,
+)
+
+
+class InitializeExecutorHandler(ControlHandler):
+ async def initialize_executor(self, req: InitializeExecutorRequest) -> EmptyReturn:
+ op_exec_with_code: OpExecWithCode = get_one_of(req.op_exec_init_info)
+ self.context.executor_manager.initialize_executor(
+ op_exec_with_code.code, req.is_source, op_exec_with_code.language
+ )
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/no_operation_handler.py b/amber/src/main/python/core/architecture/handlers/control/no_operation_handler.py
new file mode 100644
index 00000000000..926ceb61533
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/no_operation_handler.py
@@ -0,0 +1,27 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ EmptyRequest,
+)
+
+
+class NoOperationHandler(ControlHandler):
+ async def no_operation(self, req: EmptyRequest) -> EmptyReturn:
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/open_executor_handler.py b/amber/src/main/python/core/architecture/handlers/control/open_executor_handler.py
new file mode 100644
index 00000000000..2b178da5569
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/open_executor_handler.py
@@ -0,0 +1,28 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ EmptyRequest,
+)
+
+
+class OpenExecutorHandler(ControlHandler):
+ async def open_executor(self, req: EmptyRequest) -> EmptyReturn:
+ self.context.executor_manager.executor.open()
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/pause_worker_handler.py b/amber/src/main/python/core/architecture/handlers/control/pause_worker_handler.py
new file mode 100644
index 00000000000..ef9188914e0
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/pause_worker_handler.py
@@ -0,0 +1,30 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.architecture.managers.pause_manager import PauseType
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ WorkerStateResponse,
+ EmptyRequest,
+)
+
+
+class PauseWorkerHandler(ControlHandler):
+ async def pause_worker(self, req: EmptyRequest) -> WorkerStateResponse:
+ self.context.pause_manager.pause(PauseType.USER_PAUSE)
+ state = self.context.state_manager.get_current_state()
+ return WorkerStateResponse(state)
diff --git a/amber/src/main/python/core/architecture/handlers/control/query_statistics_handler.py b/amber/src/main/python/core/architecture/handlers/control/query_statistics_handler.py
new file mode 100644
index 00000000000..b636249d714
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/query_statistics_handler.py
@@ -0,0 +1,34 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ WorkerMetricsResponse,
+ EmptyRequest,
+)
+from proto.org.apache.texera.amber.engine.architecture.worker import (
+ WorkerMetrics,
+)
+
+
+class QueryStatisticsHandler(ControlHandler):
+ async def query_statistics(self, req: EmptyRequest) -> WorkerMetricsResponse:
+ metrics = WorkerMetrics(
+ worker_state=self.context.state_manager.get_current_state(),
+ worker_statistics=self.context.statistics_manager.get_statistics(),
+ )
+ return WorkerMetricsResponse(metrics)
diff --git a/amber/src/main/python/core/architecture/handlers/control/replay_current_tuple_handler.py b/amber/src/main/python/core/architecture/handlers/control/replay_current_tuple_handler.py
new file mode 100644
index 00000000000..5a916c4094a
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/replay_current_tuple_handler.py
@@ -0,0 +1,44 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import itertools
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.architecture.managers.pause_manager import PauseType
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ EmptyRequest,
+)
+from proto.org.apache.texera.amber.engine.architecture.worker import (
+ WorkerState,
+)
+
+
+class RetryCurrentTupleHandler(ControlHandler):
+ async def retry_current_tuple(self, req: EmptyRequest) -> EmptyReturn:
+ if not self.context.state_manager.confirm_state(WorkerState.COMPLETED):
+ # chain the current input tuple back on top of the current iterator to
+ # be processed once more
+ self.context.tuple_processing_manager.current_input_tuple_iter = (
+ itertools.chain(
+ [self.context.tuple_processing_manager.current_input_tuple],
+ self.context.tuple_processing_manager.current_input_tuple_iter,
+ )
+ )
+ self.context.pause_manager.resume(PauseType.USER_PAUSE)
+ self.context.pause_manager.resume(PauseType.EXCEPTION_PAUSE)
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/resume_worker_handler.py b/amber/src/main/python/core/architecture/handlers/control/resume_worker_handler.py
new file mode 100644
index 00000000000..3ebaadb6611
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/resume_worker_handler.py
@@ -0,0 +1,30 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.architecture.managers.pause_manager import PauseType
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ WorkerStateResponse,
+ EmptyRequest,
+)
+
+
+class ResumeWorkerHandler(ControlHandler):
+ async def resume_worker(self, req: EmptyRequest) -> WorkerStateResponse:
+ self.context.pause_manager.resume(PauseType.USER_PAUSE)
+ state = self.context.state_manager.get_current_state()
+ return WorkerStateResponse(state)
diff --git a/amber/src/main/python/core/architecture/handlers/control/start_channel_handler.py b/amber/src/main/python/core/architecture/handlers/control/start_channel_handler.py
new file mode 100644
index 00000000000..36747351e1e
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/start_channel_handler.py
@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.models.internal_marker import StartChannel
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ EmptyRequest,
+)
+
+
+class StartChannelHandler(ControlHandler):
+ async def start_channel(self, req: EmptyRequest) -> EmptyReturn:
+ self.context.tuple_processing_manager.current_internal_marker = StartChannel()
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/handlers/control/start_worker_handler.py b/amber/src/main/python/core/architecture/handlers/control/start_worker_handler.py
new file mode 100644
index 00000000000..bfa1556722f
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/start_worker_handler.py
@@ -0,0 +1,100 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from core.architecture.packaging.input_manager import InputManager
+from core.models import Schema
+from core.models.internal_queue import ECMElement
+from proto.org.apache.texera.amber.core import (
+ ChannelIdentity,
+ ActorVirtualIdentity,
+ PortIdentity,
+ EmbeddedControlMessageIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ WorkerStateResponse,
+ ControlInvocation,
+ EmptyRequest,
+ EmbeddedControlMessage,
+ AsyncRpcContext,
+ ControlRequest,
+ EmbeddedControlMessageType,
+)
+from proto.org.apache.texera.amber.engine.architecture.worker import (
+ WorkerState,
+)
+
+
+class StartWorkerHandler(ControlHandler):
+ async def start_worker(self, req: EmptyRequest) -> WorkerStateResponse:
+ if self.context.executor_manager.executor.is_source:
+ self.context.state_manager.transit_to(WorkerState.RUNNING)
+ input_channel_id = ChannelIdentity(
+ InputManager.SOURCE_STARTER,
+ ActorVirtualIdentity(self.context.worker_id),
+ False,
+ )
+ port_id = PortIdentity(0, False)
+ self.context.input_manager.add_input_port(
+ port_id=port_id, schema=Schema(), storage_uris=[], partitionings=[]
+ )
+ self.context.input_manager.register_input(input_channel_id, port_id)
+ self.context.current_input_channel_id = input_channel_id
+ self.context.input_queue.put(
+ ECMElement(
+ tag=input_channel_id,
+ payload=EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity("StartChannel"),
+ EmbeddedControlMessageType.NO_ALIGNMENT,
+ [],
+ {
+ input_channel_id.to_worker_id.name: ControlInvocation(
+ "StartChannel",
+ ControlRequest(empty_request=EmptyRequest()),
+ AsyncRpcContext(
+ ActorVirtualIdentity(), ActorVirtualIdentity()
+ ),
+ -1,
+ )
+ },
+ ),
+ )
+ )
+ self.context.input_queue.put(
+ ECMElement(
+ tag=input_channel_id,
+ payload=EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity("EndChannel"),
+ EmbeddedControlMessageType.PORT_ALIGNMENT,
+ [],
+ {
+ input_channel_id.to_worker_id.name: ControlInvocation(
+ "EndChannel",
+ ControlRequest(empty_request=EmptyRequest()),
+ AsyncRpcContext(
+ ActorVirtualIdentity(), ActorVirtualIdentity()
+ ),
+ -1,
+ )
+ },
+ ),
+ )
+ )
+ elif self.context.input_manager.get_input_port_mat_reader_threads():
+ self.context.input_manager.start_input_port_mat_reader_threads()
+
+ return WorkerStateResponse(self.context.state_manager.get_current_state())
diff --git a/amber/src/main/python/core/architecture/handlers/control/test_debug_command_handler.py b/amber/src/main/python/core/architecture/handlers/control/test_debug_command_handler.py
new file mode 100644
index 00000000000..382d30412d9
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/test_debug_command_handler.py
@@ -0,0 +1,195 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from core.architecture.handlers.control.debug_command_handler import (
+ WorkerDebugCommandHandler,
+)
+from core.architecture.managers.pause_manager import PauseType
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ DebugCommandRequest,
+ EmptyReturn,
+)
+
+
+class TestTranslateDebugCommand:
+ @pytest.fixture
+ def context(self):
+ return SimpleNamespace(
+ executor_manager=SimpleNamespace(operator_module_name="my_udf")
+ )
+
+ def test_break_with_lineno_prepends_module(self, context):
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("b 5", context)
+ == "b my_udf:5"
+ )
+
+ def test_long_break_with_lineno_prepends_module(self, context):
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("break 12", context)
+ == "break my_udf:12"
+ )
+
+ def test_break_preserves_condition_arg(self, context):
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("b 7 x > 0", context)
+ == "b my_udf:7 x > 0"
+ )
+
+ def test_break_with_no_args_passes_through(self, context):
+ # No args → falls through to the else branch (no module rewriting).
+ assert WorkerDebugCommandHandler.translate_debug_command("b", context) == "b"
+
+ def test_non_break_command_passes_through(self, context):
+ assert WorkerDebugCommandHandler.translate_debug_command("n", context) == "n"
+
+ def test_non_break_command_with_args_is_rejoined(self, context):
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("p some_var", context)
+ == "p some_var"
+ )
+
+ def test_leading_and_trailing_whitespace_is_stripped(self, context):
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command(" c ", context) == "c"
+ )
+
+ def test_internal_whitespace_is_collapsed_to_single_space(self, context):
+ # split() with no args collapses any run of whitespace, so the rejoined
+ # form has single spaces regardless of how many the user typed.
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("p foo bar", context)
+ == "p foo bar"
+ )
+
+ def test_break_with_only_lineno_has_no_trailing_space(self, context):
+ # The implementation joins the (empty) tail with " "; the final strip()
+ # must remove the trailing whitespace so the command stays valid pdb.
+ result = WorkerDebugCommandHandler.translate_debug_command("b 5", context)
+ assert result == "b my_udf:5"
+ assert not result.endswith(" ")
+
+ # ----- edge cases / invalid input -----
+
+ def test_empty_command_raises_descriptive_error(self, context):
+ with pytest.raises(ValueError, match="cannot be empty"):
+ WorkerDebugCommandHandler.translate_debug_command("", context)
+
+ def test_whitespace_only_command_raises_descriptive_error(self, context):
+ with pytest.raises(ValueError, match="cannot be empty"):
+ WorkerDebugCommandHandler.translate_debug_command(" \t ", context)
+
+ def test_uppercase_break_is_not_recognized(self, context):
+ # The match list is case-sensitive: ("b", "break"). "BREAK" / "B" fall
+ # through to the pass-through branch and won't get the module prefix.
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("BREAK 5", context)
+ == "BREAK 5"
+ )
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("B 5", context) == "B 5"
+ )
+
+ def test_break_with_function_name_passes_through(self, context):
+ # pdb's `b` accepts a bare function name and resolves it itself; the
+ # `module:funcname` form is invalid (pdb expects a lineno after a
+ # filename prefix). So we leave function-name args unchanged.
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("b my_func", context)
+ == "b my_func"
+ )
+
+ def test_break_with_explicit_filename_passes_through(self, context):
+ # The user already typed `filename:lineno` — don't double-prefix.
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("b foo.py:5", context)
+ == "b foo.py:5"
+ )
+
+ def test_break_with_lineno_before_module_init_raises(self, context):
+ # Without an initialized executor module we cannot construct
+ # `module:lineno`, so refuse instead of emitting `b None:5`.
+ context.executor_manager.operator_module_name = None
+ with pytest.raises(ValueError, match="executor module not initialized"):
+ WorkerDebugCommandHandler.translate_debug_command("b 5", context)
+
+ def test_break_with_function_name_before_module_init_passes_through(self, context):
+ # Function-name and filename:lineno forms don't need the module name,
+ # so they should still work even before the executor is initialized.
+ context.executor_manager.operator_module_name = None
+ assert (
+ WorkerDebugCommandHandler.translate_debug_command("b my_func", context)
+ == "b my_func"
+ )
+
+
+class TestDebugCommandAsyncFlow:
+ @pytest.fixture
+ def handler(self):
+ # ControlHandler.__init__ just stashes context; bypass the protobuf
+ # base class' __init__ by constructing via __new__.
+ instance = WorkerDebugCommandHandler.__new__(WorkerDebugCommandHandler)
+ instance.context = SimpleNamespace(
+ executor_manager=SimpleNamespace(operator_module_name="my_udf"),
+ debug_manager=MagicMock(),
+ pause_manager=MagicMock(),
+ )
+ return instance
+
+ def test_translates_then_forwards_to_debug_manager(self, handler):
+ asyncio.run(handler.debug_command(DebugCommandRequest(cmd="b 5")))
+ handler.context.debug_manager.put_debug_command.assert_called_once_with(
+ "b my_udf:5"
+ )
+
+ def test_resumes_all_three_pause_types(self, handler):
+ asyncio.run(handler.debug_command(DebugCommandRequest(cmd="c")))
+ actual = [
+ call.args[0] for call in handler.context.pause_manager.resume.call_args_list
+ ]
+ assert actual == [
+ PauseType.USER_PAUSE,
+ PauseType.EXCEPTION_PAUSE,
+ PauseType.DEBUG_PAUSE,
+ ]
+
+ def test_returns_empty_return(self, handler):
+ result = asyncio.run(handler.debug_command(DebugCommandRequest(cmd="n")))
+ assert isinstance(result, EmptyReturn)
+
+ def test_passes_through_non_break_command_unchanged(self, handler):
+ asyncio.run(handler.debug_command(DebugCommandRequest(cmd="p x")))
+ handler.context.debug_manager.put_debug_command.assert_called_once_with("p x")
+
+ def test_empty_cmd_propagates_value_error(self, handler):
+ # An empty cmd hits the ValueError in translate_debug_command. The
+ # handler does not catch it — the RPC layer will surface the failure
+ # back to the caller. Pin this so silent swallowing doesn't sneak in.
+ with pytest.raises(ValueError):
+ asyncio.run(handler.debug_command(DebugCommandRequest(cmd="")))
+
+ def test_translation_failure_skips_put_and_resume(self, handler):
+ with pytest.raises(ValueError):
+ asyncio.run(handler.debug_command(DebugCommandRequest(cmd="")))
+ handler.context.debug_manager.put_debug_command.assert_not_called()
+ handler.context.pause_manager.resume.assert_not_called()
diff --git a/amber/src/main/python/core/architecture/handlers/control/test_evaluate_expression_handler.py b/amber/src/main/python/core/architecture/handlers/control/test_evaluate_expression_handler.py
new file mode 100644
index 00000000000..a72c1f82631
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/test_evaluate_expression_handler.py
@@ -0,0 +1,156 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from core.architecture.handlers.control.evaluate_expression_handler import (
+ EvaluateExpressionHandler,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EvaluatedValue,
+ EvaluatePythonExpressionRequest,
+ TypedValue,
+)
+
+
+class TestEvaluateExpressionHandler:
+ @pytest.fixture
+ def executor(self):
+ # A stand-in for the user's UDF instance — anything addressable as
+ # `self` from the evaluated expression will do.
+ return SimpleNamespace(state="alive")
+
+ @pytest.fixture
+ def handler(self, executor):
+ instance = EvaluateExpressionHandler.__new__(EvaluateExpressionHandler)
+ instance.context = SimpleNamespace(
+ executor_manager=SimpleNamespace(executor=executor),
+ tuple_processing_manager=SimpleNamespace(
+ current_input_tuple={"col": 42},
+ current_input_port_id="port-0",
+ ),
+ )
+ return instance
+
+ def test_returns_what_the_evaluator_returns(self, handler):
+ sentinel = EvaluatedValue(
+ value=TypedValue(expression="1+1", value_ref="2", value_type="int")
+ )
+ with patch(
+ "core.architecture.handlers.control.evaluate_expression_handler"
+ ".ExpressionEvaluator.evaluate",
+ return_value=sentinel,
+ ) as evaluate:
+ result = asyncio.run(
+ handler.evaluate_python_expression(
+ EvaluatePythonExpressionRequest(expression="1+1")
+ )
+ )
+
+ assert result is sentinel
+ evaluate.assert_called_once()
+
+ def test_runtime_context_exposes_self_tuple_input(self, handler, executor):
+ with patch(
+ "core.architecture.handlers.control.evaluate_expression_handler"
+ ".ExpressionEvaluator.evaluate",
+ return_value=EvaluatedValue(),
+ ) as evaluate:
+ asyncio.run(
+ handler.evaluate_python_expression(
+ EvaluatePythonExpressionRequest(expression="self.state")
+ )
+ )
+
+ expression, runtime_context = evaluate.call_args.args
+ assert expression == "self.state"
+ assert runtime_context["self"] is executor
+ assert runtime_context["tuple_"] == {"col": 42}
+ assert runtime_context["input_"] == "port-0"
+
+ def test_runtime_context_reflects_current_tuple_at_call_time(
+ self, handler, executor
+ ):
+ # The handler must read the *current* tuple/port out of the context on
+ # each call — not snapshot them at construction. Drive two calls with
+ # different intermediate state.
+ captured: list = []
+
+ def capture(_expression, runtime_context):
+ captured.append((runtime_context["tuple_"], runtime_context["input_"]))
+ return EvaluatedValue()
+
+ with patch(
+ "core.architecture.handlers.control.evaluate_expression_handler"
+ ".ExpressionEvaluator.evaluate",
+ side_effect=capture,
+ ):
+ asyncio.run(
+ handler.evaluate_python_expression(
+ EvaluatePythonExpressionRequest(expression="x")
+ )
+ )
+ handler.context.tuple_processing_manager.current_input_tuple = {"col": 99}
+ handler.context.tuple_processing_manager.current_input_port_id = "port-1"
+ asyncio.run(
+ handler.evaluate_python_expression(
+ EvaluatePythonExpressionRequest(expression="x")
+ )
+ )
+
+ assert captured == [({"col": 42}, "port-0"), ({"col": 99}, "port-1")]
+
+ def test_handles_none_input_tuple_and_port(self, handler):
+ # Before the worker has received any input, current_input_tuple and
+ # current_input_port_id are None. The handler must still build a
+ # context (the user might be evaluating `self.foo`).
+ handler.context.tuple_processing_manager.current_input_tuple = None
+ handler.context.tuple_processing_manager.current_input_port_id = None
+ with patch(
+ "core.architecture.handlers.control.evaluate_expression_handler"
+ ".ExpressionEvaluator.evaluate",
+ return_value=EvaluatedValue(),
+ ) as evaluate:
+ asyncio.run(
+ handler.evaluate_python_expression(
+ EvaluatePythonExpressionRequest(expression="self.state")
+ )
+ )
+
+ _expression, runtime_context = evaluate.call_args.args
+ assert runtime_context["tuple_"] is None
+ assert runtime_context["input_"] is None
+
+ def test_evaluator_exception_propagates(self, handler):
+ # If the evaluator raises (bad syntax, attribute error in the user's
+ # expression, etc.), the handler must not swallow it — the RPC layer
+ # is responsible for surfacing the failure to the frontend.
+ with patch(
+ "core.architecture.handlers.control.evaluate_expression_handler"
+ ".ExpressionEvaluator.evaluate",
+ side_effect=AttributeError("no such attribute"),
+ ):
+ with pytest.raises(AttributeError, match="no such attribute"):
+ asyncio.run(
+ handler.evaluate_python_expression(
+ EvaluatePythonExpressionRequest(expression="self.missing")
+ )
+ )
diff --git a/amber/src/main/python/core/architecture/handlers/control/test_replay_current_tuple_handler.py b/amber/src/main/python/core/architecture/handlers/control/test_replay_current_tuple_handler.py
new file mode 100644
index 00000000000..2ba9b921310
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/test_replay_current_tuple_handler.py
@@ -0,0 +1,139 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from core.architecture.handlers.control.replay_current_tuple_handler import (
+ RetryCurrentTupleHandler,
+)
+from core.architecture.managers.pause_manager import PauseType
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyRequest,
+ EmptyReturn,
+)
+from proto.org.apache.texera.amber.engine.architecture.worker import WorkerState
+
+
+def _build_handler(state: WorkerState, current_tuple, remaining_iter):
+ instance = RetryCurrentTupleHandler.__new__(RetryCurrentTupleHandler)
+ state_manager = MagicMock()
+ state_manager.confirm_state.side_effect = lambda *states: state in states
+ instance.context = SimpleNamespace(
+ state_manager=state_manager,
+ tuple_processing_manager=SimpleNamespace(
+ current_input_tuple=current_tuple,
+ current_input_tuple_iter=iter(remaining_iter),
+ ),
+ pause_manager=MagicMock(),
+ )
+ return instance
+
+
+class TestRetryCurrentTupleHandler:
+ @pytest.fixture
+ def running_handler(self):
+ return _build_handler(
+ WorkerState.RUNNING,
+ current_tuple={"col": "current"},
+ remaining_iter=[{"col": "next"}],
+ )
+
+ def test_returns_empty_return(self, running_handler):
+ result = asyncio.run(running_handler.retry_current_tuple(EmptyRequest()))
+ assert isinstance(result, EmptyReturn)
+
+ def test_chains_current_tuple_back_onto_iterator(self, running_handler):
+ asyncio.run(running_handler.retry_current_tuple(EmptyRequest()))
+ # The iterator must now yield the current tuple first, then the
+ # tuples that were already queued.
+ chained = list(
+ running_handler.context.tuple_processing_manager.current_input_tuple_iter
+ )
+ assert chained == [{"col": "current"}, {"col": "next"}]
+
+ def test_resumes_user_and_exception_pause_in_order(self, running_handler):
+ asyncio.run(running_handler.retry_current_tuple(EmptyRequest()))
+ actual = [
+ call.args[0]
+ for call in running_handler.context.pause_manager.resume.call_args_list
+ ]
+ assert actual == [PauseType.USER_PAUSE, PauseType.EXCEPTION_PAUSE]
+
+ def test_does_not_resume_debug_pause(self, running_handler):
+ # Unlike WorkerDebugCommandHandler, retry only releases USER and
+ # EXCEPTION pauses — DEBUG_PAUSE must remain in effect so an active
+ # debugging session is not silently dropped.
+ asyncio.run(running_handler.retry_current_tuple(EmptyRequest()))
+ resumed = {
+ call.args[0]
+ for call in running_handler.context.pause_manager.resume.call_args_list
+ }
+ assert PauseType.DEBUG_PAUSE not in resumed
+
+ def test_no_op_when_state_is_completed(self):
+ completed_handler = _build_handler(
+ WorkerState.COMPLETED,
+ current_tuple={"col": "current"},
+ remaining_iter=[{"col": "next"}],
+ )
+ result = asyncio.run(completed_handler.retry_current_tuple(EmptyRequest()))
+
+ # Iterator must be untouched (no chaining), and no pause type is
+ # resumed — replaying a tuple after completion is meaningless.
+ remaining = list(
+ completed_handler.context.tuple_processing_manager.current_input_tuple_iter
+ )
+ assert remaining == [{"col": "next"}]
+ completed_handler.context.pause_manager.resume.assert_not_called()
+ assert isinstance(result, EmptyReturn)
+
+ def test_chains_even_when_remaining_iter_is_exhausted(self):
+ handler = _build_handler(
+ WorkerState.RUNNING,
+ current_tuple={"col": "lone"},
+ remaining_iter=[],
+ )
+ asyncio.run(handler.retry_current_tuple(EmptyRequest()))
+ chained = list(
+ handler.context.tuple_processing_manager.current_input_tuple_iter
+ )
+ assert chained == [{"col": "lone"}]
+
+ def test_paused_state_still_chains_and_resumes(self):
+ # The completion guard is `if not confirm_state(COMPLETED)`, so every
+ # other state — RUNNING, READY, PAUSED, UNINITIALIZED — must take the
+ # chain+resume path. PAUSED is the most likely real-world entry point
+ # (the user hits "retry" while the worker is paused on an exception).
+ handler = _build_handler(
+ WorkerState.PAUSED,
+ current_tuple={"col": "current"},
+ remaining_iter=[{"col": "next"}],
+ )
+ asyncio.run(handler.retry_current_tuple(EmptyRequest()))
+
+ chained = list(
+ handler.context.tuple_processing_manager.current_input_tuple_iter
+ )
+ assert chained == [{"col": "current"}, {"col": "next"}]
+ resumed = [
+ call.args[0] for call in handler.context.pause_manager.resume.call_args_list
+ ]
+ assert resumed == [PauseType.USER_PAUSE, PauseType.EXCEPTION_PAUSE]
diff --git a/amber/src/main/python/core/architecture/handlers/control/update_executor_handler.py b/amber/src/main/python/core/architecture/handlers/control/update_executor_handler.py
new file mode 100644
index 00000000000..f2b6d16d46a
--- /dev/null
+++ b/amber/src/main/python/core/architecture/handlers/control/update_executor_handler.py
@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.control_handler_base import ControlHandler
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmptyReturn,
+ UpdateExecutorRequest,
+)
+from core.util import get_one_of
+from proto.org.apache.texera.amber.core import OpExecWithCode
+
+
+class UpdateExecutorHandler(ControlHandler):
+ async def update_executor(self, req: UpdateExecutorRequest) -> EmptyReturn:
+ op_exec_with_code: OpExecWithCode = get_one_of(req.new_exec_init_info)
+ self.context.executor_manager.update_executor(
+ op_exec_with_code.code, self.context.executor_manager.executor.is_source
+ )
+ return EmptyReturn()
diff --git a/amber/src/main/python/core/architecture/managers/__init__.py b/amber/src/main/python/core/architecture/managers/__init__.py
new file mode 100644
index 00000000000..dae74d4f915
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/__init__.py
@@ -0,0 +1,23 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .context import Context
+from .pause_manager import PauseManager
+from .state_manager import StateManager
+from .statistics_manager import StatisticsManager
+
+__all__ = ["Context", "PauseManager", "StateManager", "StatisticsManager"]
diff --git a/amber/src/main/python/core/architecture/managers/console_message_manager.py b/amber/src/main/python/core/architecture/managers/console_message_manager.py
new file mode 100644
index 00000000000..f75c74cda0f
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/console_message_manager.py
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Iterator
+
+from core.util.buffer.timed_buffer import TimedBuffer
+from proto.org.apache.texera.amber.engine.architecture.rpc import ConsoleMessage
+
+
+class ConsoleMessageManager:
+ def __init__(self):
+ self.print_buf = TimedBuffer()
+
+ def get_messages(self, force_flush: bool = False) -> Iterator[ConsoleMessage]:
+ return self.print_buf.get(force_flush)
+
+ def put_message(self, msg: ConsoleMessage) -> None:
+ self.print_buf.put(msg)
diff --git a/amber/src/main/python/core/architecture/managers/context.py b/amber/src/main/python/core/architecture/managers/context.py
new file mode 100644
index 00000000000..3629a435a79
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/context.py
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Optional
+
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.worker import WorkerState
+from .console_message_manager import ConsoleMessageManager
+from .debug_manager import DebugManager
+from .embedded_control_message_manager import EmbeddedControlMessageManager
+from .exception_manager import ExceptionManager
+from .executor_manager import ExecutorManager
+from .pause_manager import PauseManager
+from .state_manager import StateManager
+from .state_processing_manager import StateProcessingManager
+from .statistics_manager import StatisticsManager
+from .tuple_processing_manager import TupleProcessingManager
+from ..packaging.input_manager import InputManager
+from ..packaging.output_manager import OutputManager
+from ...models import InternalQueue
+
+
+class Context:
+ """
+ Manages context of command handlers. Many of those attributes belongs to the DP
+ thread, they are managed here to show a clean interface what handlers can or
+ should access.
+
+ Context class can be viewed as a friend of DataProcessor.
+ """
+
+ def __init__(self, worker_id, input_queue):
+ self.worker_id = worker_id
+ self.input_queue: InternalQueue = input_queue
+ self.executor_manager = ExecutorManager()
+ self.current_input_channel_id: Optional[ChannelIdentity] = None
+ self.tuple_processing_manager = TupleProcessingManager()
+ self.state_processing_manager = StateProcessingManager()
+ self.exception_manager = ExceptionManager()
+ self.state_manager = StateManager(
+ {
+ WorkerState.UNINITIALIZED: {WorkerState.READY},
+ WorkerState.READY: {WorkerState.PAUSED, WorkerState.RUNNING},
+ WorkerState.RUNNING: {WorkerState.PAUSED, WorkerState.COMPLETED},
+ WorkerState.PAUSED: {WorkerState.RUNNING},
+ WorkerState.COMPLETED: set(),
+ },
+ WorkerState.UNINITIALIZED,
+ )
+
+ self.statistics_manager = StatisticsManager()
+ self.pause_manager = PauseManager(
+ self.input_queue, state_manager=self.state_manager
+ )
+ self.output_manager = OutputManager(worker_id)
+ self.input_manager = InputManager(worker_id, self.input_queue)
+ self.ecm_manager = EmbeddedControlMessageManager(
+ ActorVirtualIdentity(worker_id), self.input_manager
+ )
+ self.console_message_manager = ConsoleMessageManager()
+ self.debug_manager = DebugManager(
+ self.tuple_processing_manager.context_switch_condition
+ )
+
+ def close(self):
+ self.executor_manager.close()
diff --git a/amber/src/main/python/core/architecture/managers/debug_manager.py b/amber/src/main/python/core/architecture/managers/debug_manager.py
new file mode 100644
index 00000000000..99a5aa7ee85
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/debug_manager.py
@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from pdb import Pdb
+from threading import Condition
+
+from core.models.single_blocking_io import SingleBlockingIO
+
+
+class DebugManager:
+ def __init__(self, condition: Condition):
+ self._debug_in = SingleBlockingIO(condition)
+ self._debug_out = SingleBlockingIO(condition)
+ self.debugger = Pdb(stdin=self._debug_in, stdout=self._debug_out, nosigint=True)
+
+ # Customized prompt, we can design our prompt for the debugger.
+ self.debugger.prompt = ""
+
+ def has_debug_command(self) -> bool:
+ return self._debug_in.value is not None
+
+ def has_debug_event(self) -> bool:
+ return self._debug_out.value is not None
+
+ def get_debug_event(self) -> str:
+ """
+ Blocking gets for the next debug event.
+ :return str: the fetched event, in string format.
+ """
+ return self._debug_out.readline()
+
+ def put_debug_command(self, command: str) -> None:
+ """
+ Puts a debug command.
+ :param command: the command to be put, in string format.
+ :return:
+ """
+ self._debug_in.write(command)
+ self._debug_in.flush()
diff --git a/amber/src/main/python/core/architecture/managers/embedded_control_message_manager.py b/amber/src/main/python/core/architecture/managers/embedded_control_message_manager.py
new file mode 100644
index 00000000000..8ba80f28cef
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/embedded_control_message_manager.py
@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from collections import defaultdict
+from typing import Set, Dict
+
+from core.architecture.packaging.input_manager import Channel
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmbeddedControlMessage,
+ EmbeddedControlMessageType,
+)
+
+
+class EmbeddedControlMessageManager:
+ def __init__(self, actor_id: ActorVirtualIdentity, input_gateway):
+ self.actor_id = actor_id
+ self.input_gateway = input_gateway
+ self.ecm_received: Dict[str, Set[ChannelIdentity]] = defaultdict(set)
+
+ def is_ecm_aligned(
+ self, from_channel: ChannelIdentity, ecm: EmbeddedControlMessage
+ ) -> bool:
+ """
+ Checks whether an ECM has been received from all expected
+ input channels, determining whether further processing can proceed.
+
+ Args:
+ from_channel (ChannelIdentity): The channel from which the ECM was received.
+ ecm (EmbeddedControlMessage): The ECM payload containing its type and scope.
+
+ Returns:
+ bool: True if the ECM is considered aligned and processing can
+ continue, False otherwise.
+ """
+
+ self.ecm_received[ecm.id].add(from_channel)
+ ecm_received_from_all_channels = self.get_channels_within_scope(ecm).issubset(
+ self.ecm_received[ecm.id]
+ )
+
+ if ecm.ecm_type == EmbeddedControlMessageType.ALL_ALIGNMENT:
+ ecm_completed = ecm_received_from_all_channels
+ elif ecm.ecm_type == EmbeddedControlMessageType.PORT_ALIGNMENT:
+ port_id = self.input_gateway.get_port_id(from_channel)
+ ecm_completed = (
+ self.input_gateway.get_port(port_id)
+ .get_channels()
+ .issubset(self.ecm_received[ecm.id])
+ )
+ elif ecm.ecm_type == EmbeddedControlMessageType.NO_ALIGNMENT:
+ ecm_completed = (
+ len(self.ecm_received[ecm.id]) == 1
+ ) # Only the first ECM triggers
+ else:
+ raise ValueError(f"Unsupported ECM type: {ecm.ecm_type}")
+
+ if ecm_received_from_all_channels:
+ del self.ecm_received[ecm.id] # Clean up if all ECMs are received
+
+ return ecm_completed
+
+ def get_channels_within_scope(self, ecm: EmbeddedControlMessage) -> Dict[
+ "ChannelIdentity", "Channel"
+ ].keys:
+ if ecm.scope:
+ upstreams = {
+ channel_id
+ for channel_id in ecm.scope
+ if channel_id.to_worker_id == self.actor_id
+ }
+ return self.input_gateway.get_all_channel_ids() & upstreams
+ return self.input_gateway.get_all_data_channel_ids()
diff --git a/amber/src/main/python/core/architecture/managers/exception_manager.py b/amber/src/main/python/core/architecture/managers/exception_manager.py
new file mode 100644
index 00000000000..f3ea4dd4bf1
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/exception_manager.py
@@ -0,0 +1,38 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Optional, List
+
+from core.models import ExceptionInfo
+
+
+class ExceptionManager:
+ def __init__(self):
+ self.exc_info: Optional[ExceptionInfo] = None
+ self.exc_info_history: List[ExceptionInfo] = list()
+
+ def set_exception_info(self, exc_info: ExceptionInfo) -> None:
+ self.exc_info = exc_info
+ self.exc_info_history.append(exc_info)
+
+ def has_exception(self) -> bool:
+ return self.exc_info is not None
+
+ def get_exc_info(self) -> ExceptionInfo:
+ exc_info = self.exc_info
+ self.exc_info = None
+ return exc_info
diff --git a/amber/src/main/python/core/architecture/managers/executor_manager.py b/amber/src/main/python/core/architecture/managers/executor_manager.py
new file mode 100644
index 00000000000..eb1363d0a68
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/executor_manager.py
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import fs
+import importlib
+import inspect
+import sys
+from cached_property import cached_property
+from fs.base import FS
+from loguru import logger
+from pathlib import Path
+from typing import Tuple, Optional
+
+from core.models import Operator, SourceOperator
+
+
+class ExecutorManager:
+ def __init__(self):
+ self.executor: Optional[Operator] = None
+ self.operator_module_name: Optional[str] = None
+ self.executor_version: int = 0 # incremental only
+
+ @cached_property
+ def fs(self) -> FS:
+ """
+ Creates a tmp fs for storing source code, which will be removed when the
+ workflow is completed.
+ :return:
+ """
+ # TODO:
+ # For various reasons when the workflow is not completed successfully,
+ # the tmp fs could not be closed properly. This means it may leave files
+ # in the /var/tmp folder after a partially started or failed execution.
+ # A full-life-cycle management of tmp fs is required to consider all
+ # possible errors happened during execution. However, the full-life-cycle
+ # management could be hard due to errors from JAVA side which causes force
+ # kill on the Python process.
+ # As each python file is usually tiny in size, and the OS can
+ # periodically clean up /var/tmp anyway, the full-life-cycle management is
+ # not a priority to be fixed.
+ temp_fs = fs.open_fs("temp://")
+ root = Path(temp_fs.getsyspath("/"))
+ logger.debug(f"Opening a tmp directory at {root}.")
+ sys.path.append(str(root))
+ return temp_fs
+
+ def gen_module_file_name(self) -> Tuple[str, str]:
+ """
+ Generate a UUID to be used as udf source code file.
+ :return Tuple[str, str]: the pair of module_name and file_name.
+ """
+ self.executor_version += 1
+ module_name = f"udf-v{self.executor_version}"
+ file_name = f"{module_name}.py"
+ return module_name, file_name
+
+ def load_executor_definition(self, code: str) -> type(Operator):
+ """
+ Load the given executor code in string into a class definition
+ :param code: str, python code that defines an Operator, should contain one
+ and only one Executor definition.
+ :return: an Operator sub-class definition
+ """
+ module_name, file_name = self.gen_module_file_name()
+
+ with self.fs.open(file_name, "w") as file:
+ file.write(code)
+ logger.debug(
+ "A tmp py file is written to "
+ f"{Path(self.fs.getsyspath('/')).joinpath(file_name)}."
+ )
+
+ if module_name in sys.modules:
+ executor_module = importlib.import_module(module_name)
+ executor_module.__dict__.clear()
+ executor_module.__dict__["__name__"] = module_name
+ executor_module = importlib.reload(executor_module)
+ else:
+ executor_module = importlib.import_module(module_name)
+ self.operator_module_name = module_name
+
+ executors = list(
+ filter(self.is_concrete_operator, executor_module.__dict__.values())
+ )
+ assert len(executors) == 1, "There should be one and only one Operator defined"
+ return executors[0]
+
+ def close(self) -> None:
+ """
+ Close the tmp fs and release all resources created within it.
+ :return:
+ """
+ self.fs.close()
+ logger.debug(f"Tmp directory {self.fs.getsyspath('/')} is closed and cleared.")
+
+ @staticmethod
+ def is_concrete_operator(cls: type) -> bool:
+ """
+ Check if the class is a non-abstract Operator.
+ :param cls: a target class to be evaluated
+ :return: bool
+ """
+
+ return (
+ inspect.isclass(cls)
+ and issubclass(cls, Operator)
+ and not inspect.isabstract(cls)
+ )
+
+ def initialize_executor(self, code: str, is_source: bool, language: str) -> None:
+ """
+ Initialize the executor with the given code. The output schema is
+ decided by the user.
+
+ :param code: The string version of the code, containing one Operator
+ class declaration.
+ :param is_source: Indicating if the operator is used as a source operator.
+ :param language: The language of the operator code.
+ :return:
+ """
+ if language in ("r-tuple", "r-table"):
+ # R support is provided by an optional plugin (texera-rudf)
+ executor_type = "Tuple" if language == "r-tuple" else "Table"
+ try:
+ import texera_r
+
+ class_suffix = "SourceExecutor" if is_source else "Executor"
+ executor_class = getattr(texera_r, f"R{executor_type}{class_suffix}")
+ except ImportError as e:
+ raise ImportError(
+ "R operators require the texera-rudf package.\n"
+ "Install with: pip install git+https://github.com/Texera/texera-rudf.git\n"
+ f"Import error: {e}"
+ )
+ self.executor = executor_class(code)
+ else:
+ executor: type(Operator) = self.load_executor_definition(code)
+ self.executor = executor()
+ self.executor.is_source = is_source
+ assert isinstance(self.executor, SourceOperator) == self.executor.is_source, (
+ "Please use SourceOperator API for source operators."
+ )
+
+ def update_executor(self, code: str, is_source: bool) -> None:
+ """
+ Update the executor, preserving its state in the __dict__.
+ The user is responsible to make sure the state can be used by the new logic.
+
+ :param code: The string version of python code, containing one Operator
+ class declaration.
+ :param is_source: Indicating if the operator is used as a source operator.
+ :return:
+ """
+ original_internal_state = self.executor.__dict__
+ executor: type(Operator) = self.load_executor_definition(code)
+ self.executor = executor()
+ self.executor.is_source = is_source
+ assert isinstance(self.executor, SourceOperator) == self.executor.is_source, (
+ "Please use SourceOperator API for source operators."
+ )
+ # overwrite the internal state
+ self.executor.__dict__ = original_internal_state
+ # TODO:
+ # it may be an interesting idea to preserve versions of code and versions
+ # of states whenever the operator logic is being updated.
diff --git a/amber/src/main/python/core/architecture/managers/pause_manager.py b/amber/src/main/python/core/architecture/managers/pause_manager.py
new file mode 100644
index 00000000000..11db44b78bd
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/pause_manager.py
@@ -0,0 +1,93 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from collections import defaultdict
+from enum import Enum
+from loguru import logger
+from typing import Set, Dict
+
+from proto.org.apache.texera.amber.core import ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.worker import WorkerState
+from . import state_manager
+from ...models import InternalQueue
+
+
+class PauseType(Enum):
+ NO_PAUSE = 0
+ USER_PAUSE = 1
+ DEBUG_PAUSE = 2
+ EXCEPTION_PAUSE = 3
+ ECM_PAUSE = 4
+
+
+class PauseManager:
+ """
+ Manage pause states.
+ """
+
+ def __init__(
+ self,
+ input_queue: InternalQueue,
+ state_manager: state_manager.StateManager,
+ ):
+ self._input_queue: InternalQueue = input_queue
+ self._global_pauses: Set[PauseType] = set()
+ self._specific_input_pauses: Dict[PauseType, Set[ChannelIdentity]] = (
+ defaultdict(set)
+ )
+ self._state_manager = state_manager
+
+ def pause(self, pause_type: PauseType, change_state=True) -> None:
+ logger.debug("pause by " + str(pause_type))
+ self._global_pauses.add(pause_type)
+ self._input_queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE)
+
+ if change_state and self._state_manager.confirm_state(
+ WorkerState.RUNNING, WorkerState.READY
+ ):
+ self._state_manager.transit_to(WorkerState.PAUSED)
+
+ def pause_input_channel(
+ self, pause_type: PauseType, channel_id: ChannelIdentity
+ ) -> None:
+ self._specific_input_pauses[pause_type].add(channel_id)
+ self._input_queue.disable(channel_id)
+
+ def resume(self, pause_type: PauseType, change_state=True) -> None:
+ if pause_type in self._global_pauses:
+ self._global_pauses.remove(pause_type)
+ if pause_type in self._specific_input_pauses:
+ # need to resume specific input channels
+ for channel_id in self._specific_input_pauses[pause_type]:
+ self._input_queue.enable(channel_id)
+ del self._specific_input_pauses[pause_type]
+
+ # still globally paused no action, don't need to resume anything
+ if self._global_pauses:
+ return
+
+ # global pause is empty, specific input pause is also empty, resume all
+ if not self._specific_input_pauses:
+ self._input_queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE)
+ if change_state and self._state_manager.confirm_state(WorkerState.PAUSED):
+ self._state_manager.transit_to(WorkerState.RUNNING)
+ return
+
+ def is_paused(self) -> bool:
+ return bool(self._global_pauses) and self._state_manager.confirm_state(
+ WorkerState.PAUSED
+ )
diff --git a/amber/src/main/python/core/architecture/managers/state_manager.py b/amber/src/main/python/core/architecture/managers/state_manager.py
new file mode 100644
index 00000000000..e80b6d5fc4f
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/state_manager.py
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Dict, Set, Tuple, Union
+
+from typing_extensions import T
+
+
+class InvalidStateException(Exception):
+ pass
+
+
+class InvalidTransitionException(Exception):
+ pass
+
+
+class StateManager:
+ """
+ A generalized StateManager that provides APIs for state transition, assertion,
+ and confirmation.
+ """
+
+ def __init__(self, state_transition_graph: Dict[T, Set[T]], initial_state: T):
+ self._state_transition_graph = state_transition_graph
+ self._current_state: T = initial_state
+
+ def assert_state(self, state: T) -> None:
+ """
+ Assert the current state to be the expected state, raise exception if otherwise.
+ :param state: the expected state.
+ """
+ if self._current_state != state:
+ raise InvalidStateException(
+ f"Excepted state = {state} but current state = {self._current_state}"
+ )
+
+ def confirm_state(self, *states: Union[T, Tuple[T]]) -> bool:
+ """
+ Check if current state is in one of the states.
+
+ :param states: Union[T, Tuple[T]], a series of states to be checked.
+ :return: bool
+ """
+ return any(self._current_state == state for state in states)
+
+ def transit_to(self, state: T) -> None:
+ """
+ Transit the current state into the target state.
+
+ :param state: T, the target state to transit to.
+ :return:
+ """
+
+ # do nothing if the current state is already the target state
+ if state == self._current_state:
+ return
+
+ if state not in self._state_transition_graph.get(self._current_state, set()):
+ raise InvalidTransitionException(
+ f"Cannot transit from {self._current_state} to {state}"
+ )
+
+ self._current_state = state
+
+ def get_current_state(self) -> T:
+ """
+ Return the current state.
+ :return:
+ """
+ return self._current_state
diff --git a/amber/src/main/python/core/architecture/managers/state_processing_manager.py b/amber/src/main/python/core/architecture/managers/state_processing_manager.py
new file mode 100644
index 00000000000..442b6e03564
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/state_processing_manager.py
@@ -0,0 +1,34 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Optional
+
+from core.models.state import State
+
+
+class StateProcessingManager:
+ def __init__(self):
+ self.current_input_state: Optional[State] = None
+ self.current_output_state: Optional[State] = None
+
+ def get_input_state(self) -> Optional[State]:
+ ret, self.current_input_state = self.current_input_state, None
+ return ret
+
+ def get_output_state(self) -> Optional[State]:
+ ret, self.current_output_state = self.current_output_state, None
+ return ret
diff --git a/amber/src/main/python/core/architecture/managers/statistics_manager.py b/amber/src/main/python/core/architecture/managers/statistics_manager.py
new file mode 100644
index 00000000000..6b36b78e577
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/statistics_manager.py
@@ -0,0 +1,92 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from collections import defaultdict
+from typing import DefaultDict
+
+from proto.org.apache.texera.amber.core import PortIdentity
+from proto.org.apache.texera.amber.engine.architecture.worker import (
+ WorkerStatistics,
+ PortTupleMetricsMapping,
+ TupleMetrics,
+)
+
+
+class StatisticsManager:
+ def __init__(self) -> None:
+ # Initialize metrics with default values
+ self._input_tuple_metrics: DefaultDict[PortIdentity, TupleMetrics] = (
+ defaultdict(lambda: (0, 0))
+ )
+ self._output_tuple_metrics: DefaultDict[PortIdentity, TupleMetrics] = (
+ defaultdict(lambda: (0, 0))
+ )
+ self._data_processing_time: int = 0
+ self._control_processing_time: int = 0
+ self._total_execution_time: int = 0
+ self._worker_start_time: int = 0
+
+ def get_statistics(self) -> WorkerStatistics:
+ # Compile and return worker statistics
+ return WorkerStatistics(
+ [
+ PortTupleMetricsMapping(port_id, TupleMetrics(*tuple_metrics))
+ for port_id, tuple_metrics in self._input_tuple_metrics.items()
+ ],
+ [
+ PortTupleMetricsMapping(port_id, TupleMetrics(*tuple_metrics))
+ for port_id, tuple_metrics in self._output_tuple_metrics.items()
+ ],
+ self._data_processing_time,
+ self._control_processing_time,
+ self._total_execution_time
+ - self._data_processing_time
+ - self._control_processing_time,
+ )
+
+ def increase_input_statistics(self, port_id: PortIdentity, size: int) -> None:
+ if size < 0:
+ raise ValueError("Tuple size must be non-negative")
+ count, total_size = self._input_tuple_metrics[port_id]
+ self._input_tuple_metrics[port_id] = (count + 1, total_size + size)
+
+ def increase_output_statistics(self, port_id: PortIdentity, size: int) -> None:
+ if size < 0:
+ raise ValueError("Tuple size must be non-negative")
+ count, total_size = self._output_tuple_metrics[port_id]
+ self._output_tuple_metrics[port_id] = (count + 1, total_size + size)
+
+ def increase_data_processing_time(self, time: int) -> None:
+ if time < 0:
+ raise ValueError("Time must be non-negative")
+ self._data_processing_time += time
+
+ def increase_control_processing_time(self, time: int) -> None:
+ if time < 0:
+ raise ValueError("Time must be non-negative")
+ self._control_processing_time += time
+
+ def update_total_execution_time(self, time: int) -> None:
+ if time < self._worker_start_time:
+ raise ValueError(
+ "Current time must be greater than or equal to worker start time"
+ )
+ self._total_execution_time = time - self._worker_start_time
+
+ def initialize_worker_start_time(self, time: int) -> None:
+ # Set the worker start time
+ self._worker_start_time = time
diff --git a/amber/src/main/python/core/architecture/managers/test_debug_manager.py b/amber/src/main/python/core/architecture/managers/test_debug_manager.py
new file mode 100644
index 00000000000..248a2e134ef
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/test_debug_manager.py
@@ -0,0 +1,116 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from threading import Condition
+
+import pytest
+
+from core.architecture.managers.debug_manager import DebugManager
+
+
+class TestDebugManager:
+ @pytest.fixture
+ def debug_manager(self):
+ return DebugManager(Condition())
+
+ def test_it_can_init(self, debug_manager):
+ assert debug_manager.debugger is not None
+ assert debug_manager.debugger.prompt == ""
+
+ def test_it_has_no_command_initially(self, debug_manager):
+ assert not debug_manager.has_debug_command()
+
+ def test_it_has_no_event_initially(self, debug_manager):
+ assert not debug_manager.has_debug_event()
+
+ def test_put_command_sets_has_debug_command(self, debug_manager):
+ debug_manager.put_debug_command("n")
+ assert debug_manager.has_debug_command()
+
+ def test_get_debug_event_returns_flushed_output(self, debug_manager):
+ # Pdb writes to its stdout via the SingleBlockingIO; simulate that path
+ # directly so we don't have to spin up a real debugging session.
+ debug_manager.debugger.stdout.write("hit breakpoint")
+ debug_manager.debugger.stdout.flush()
+ assert debug_manager.has_debug_event()
+ assert debug_manager.get_debug_event() == "hit breakpoint\n"
+ assert not debug_manager.has_debug_event()
+
+ def test_command_pipe_and_event_pipe_are_independent(self, debug_manager):
+ debug_manager.put_debug_command("step")
+ assert debug_manager.has_debug_command()
+ assert not debug_manager.has_debug_event()
+
+ debug_manager.debugger.stdout.write("event")
+ debug_manager.debugger.stdout.flush()
+ # Putting a command must not consume an event, and vice versa.
+ assert debug_manager.has_debug_command()
+ assert debug_manager.has_debug_event()
+
+ def test_pdb_is_wired_to_debug_pipes(self, debug_manager):
+ # The Pdb instance must read from the same IO that put_debug_command
+ # writes to, and write to the same IO that get_debug_event reads from.
+ debug_manager.put_debug_command("c")
+ # Reading via the debugger's stdin must see the queued command.
+ assert debug_manager.debugger.stdin.readline() == "c\n"
+
+ debug_manager.debugger.stdout.write("paused")
+ debug_manager.debugger.stdout.flush()
+ assert debug_manager.get_debug_event() == "paused\n"
+
+ def test_event_pipe_supports_multiple_round_trips(self, debug_manager):
+ for line in ("first", "second", "third"):
+ debug_manager.debugger.stdout.write(line)
+ debug_manager.debugger.stdout.flush()
+ assert debug_manager.get_debug_event() == f"{line}\n"
+ assert not debug_manager.has_debug_event()
+
+ def test_debugger_uses_nosigint_to_avoid_signal_install(self, debug_manager):
+ # We construct Pdb with nosigint=True to avoid touching signal handlers
+ # in the worker thread. Guard against accidental flips.
+ assert debug_manager.debugger.nosigint is True
+
+ # ----- edge cases / quirks -----
+
+ def test_put_empty_command_still_marks_command_present(self, debug_manager):
+ # SingleBlockingIO.flush always commits buf + "\n" to value, so even
+ # an empty command becomes a "\n" line and shows up as a pending
+ # command. Documents current behavior.
+ debug_manager.put_debug_command("")
+ assert debug_manager.has_debug_command()
+ assert debug_manager.debugger.stdin.readline() == "\n"
+
+ def test_put_overwrites_unconsumed_command(self, debug_manager):
+ # The command pipe holds at most one value. A second put without an
+ # intervening consume silently overwrites the first — known data-loss
+ # quirk of SingleBlockingIO. Pinning this so callers don't accidentally
+ # rely on queued semantics.
+ debug_manager.put_debug_command("first")
+ debug_manager.put_debug_command("second")
+ assert debug_manager.debugger.stdin.readline() == "second\n"
+
+ def test_put_command_with_embedded_newline_is_passed_verbatim(self, debug_manager):
+ # An embedded newline is not sanitized; pdb would see the raw bytes.
+ debug_manager.put_debug_command("step\nlist")
+ assert debug_manager.debugger.stdin.readline() == "step\nlist\n"
+
+ def test_event_pipe_overwrites_unconsumed_event(self, debug_manager):
+ debug_manager.debugger.stdout.write("first")
+ debug_manager.debugger.stdout.flush()
+ debug_manager.debugger.stdout.write("second")
+ debug_manager.debugger.stdout.flush()
+ assert debug_manager.get_debug_event() == "second\n"
diff --git a/amber/src/main/python/core/architecture/managers/test_executor_manager.py b/amber/src/main/python/core/architecture/managers/test_executor_manager.py
new file mode 100644
index 00000000000..901f768a216
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/test_executor_manager.py
@@ -0,0 +1,248 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import sys
+import pytest
+from unittest.mock import MagicMock
+
+from core.architecture.managers.executor_manager import ExecutorManager
+
+
+# Sample operator code for testing
+SAMPLE_OPERATOR_CODE = """
+from pytexera import *
+
+class TestOperator(UDFOperatorV2):
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ yield tuple_
+"""
+
+SAMPLE_SOURCE_OPERATOR_CODE = """
+from pytexera import *
+
+class TestSourceOperator(UDFSourceOperator):
+ def produce(self) -> Iterator[Union[TupleLike, TableLike, None]]:
+ yield Tuple({"test": "data"})
+"""
+
+
+class TestExecutorManager:
+ """Test suite for ExecutorManager, focusing on R UDF plugin support."""
+
+ @pytest.fixture
+ def executor_manager(self):
+ """Create a fresh ExecutorManager instance for each test."""
+ manager = ExecutorManager()
+ yield manager
+ # Cleanup: close the temp filesystem
+ if hasattr(manager, "_fs"):
+ manager.close()
+
+ def _mock_r_plugin(self, executor_class_name, is_source):
+ """
+ Helper to mock the texera_r plugin module.
+
+ :param executor_class_name: Name of the executor class (e.g., 'RTupleExecutor')
+ :param is_source: Whether the executor is a source operator
+ :return: Tuple of (mock_texera_r, mock_executor_instance)
+ """
+ from core.models import SourceOperator, Operator
+
+ mock_texera_r = MagicMock()
+ mock_executor_class = MagicMock()
+ setattr(mock_texera_r, executor_class_name, mock_executor_class)
+
+ # Use appropriate spec based on operator type
+ spec_class = SourceOperator if is_source else Operator
+ mock_executor_instance = MagicMock(spec=spec_class)
+ mock_executor_instance.is_source = is_source
+ mock_executor_class.return_value = mock_executor_instance
+
+ sys.modules["texera_r"] = mock_texera_r
+ return mock_texera_r, mock_executor_instance
+
+ def _cleanup_r_plugin(self):
+ """Remove the mocked texera_r module from sys.modules."""
+ if "texera_r" in sys.modules:
+ del sys.modules["texera_r"]
+
+ def test_initialization(self, executor_manager):
+ """Test that ExecutorManager initializes correctly."""
+ assert executor_manager.executor is None
+ assert executor_manager.operator_module_name is None
+ assert executor_manager.executor_version == 0
+
+ def test_reject_r_tuple_language(self, executor_manager):
+ """Test that 'r-tuple' language is rejected with ImportError when plugin is not available."""
+ with pytest.raises(ImportError) as exc_info:
+ executor_manager.initialize_executor(
+ code=SAMPLE_OPERATOR_CODE, is_source=False, language="r-tuple"
+ )
+
+ # Verify the error message mentions R operators require the texera-rudf package
+ assert "texera-rudf" in str(exc_info.value) or "R operators require" in str(
+ exc_info.value
+ )
+
+ def test_reject_r_table_language(self, executor_manager):
+ """Test that 'r-table' language is rejected with ImportError when plugin is not available."""
+ with pytest.raises(ImportError) as exc_info:
+ executor_manager.initialize_executor(
+ code=SAMPLE_OPERATOR_CODE, is_source=False, language="r-table"
+ )
+
+ # Verify the error message mentions R operators require the texera-rudf package
+ assert "texera-rudf" in str(exc_info.value) or "R operators require" in str(
+ exc_info.value
+ )
+
+ def test_accept_r_tuple_language_with_plugin(self, executor_manager):
+ """Test that 'r-tuple' language is accepted when plugin is available."""
+ _, mock_executor = self._mock_r_plugin("RTupleExecutor", is_source=False)
+ try:
+ executor_manager.initialize_executor(
+ code="# R code", is_source=False, language="r-tuple"
+ )
+ assert executor_manager.executor == mock_executor
+ finally:
+ self._cleanup_r_plugin()
+
+ def test_accept_r_table_language_with_plugin(self, executor_manager):
+ """Test that 'r-table' language is accepted when plugin is available."""
+ _, mock_executor = self._mock_r_plugin("RTableExecutor", is_source=False)
+ try:
+ executor_manager.initialize_executor(
+ code="# R code", is_source=False, language="r-table"
+ )
+ assert executor_manager.executor == mock_executor
+ finally:
+ self._cleanup_r_plugin()
+
+ def test_accept_r_tuple_source_with_plugin(self, executor_manager):
+ """Test that 'r-tuple' source operators work when plugin is available."""
+ _, mock_executor = self._mock_r_plugin("RTupleSourceExecutor", is_source=True)
+ try:
+ executor_manager.initialize_executor(
+ code="# R code", is_source=True, language="r-tuple"
+ )
+ assert executor_manager.executor == mock_executor
+ finally:
+ self._cleanup_r_plugin()
+
+ def test_accept_r_table_source_with_plugin(self, executor_manager):
+ """Test that 'r-table' source operators work when plugin is available."""
+ _, mock_executor = self._mock_r_plugin("RTableSourceExecutor", is_source=True)
+ try:
+ executor_manager.initialize_executor(
+ code="# R code", is_source=True, language="r-table"
+ )
+ assert executor_manager.executor == mock_executor
+ finally:
+ self._cleanup_r_plugin()
+
+ def test_accept_python_language_regular_operator(self, executor_manager):
+ """Test that 'python' language is accepted for regular operators."""
+ # This should not raise any assertion error
+ executor_manager.initialize_executor(
+ code=SAMPLE_OPERATOR_CODE, is_source=False, language="python"
+ )
+
+ # Verify executor was initialized
+ assert executor_manager.executor is not None
+ assert executor_manager.operator_module_name == "udf-v1"
+ assert executor_manager.executor_version == 1
+ assert executor_manager.executor.is_source is False
+
+ def test_accept_python_language_source_operator(self, executor_manager):
+ """Test that 'python' language is accepted for source operators."""
+ # This should not raise any assertion error
+ executor_manager.initialize_executor(
+ code=SAMPLE_SOURCE_OPERATOR_CODE, is_source=True, language="python"
+ )
+
+ # Verify executor was initialized
+ assert executor_manager.executor is not None
+ assert executor_manager.operator_module_name == "udf-v1"
+ assert executor_manager.executor_version == 1
+ assert executor_manager.executor.is_source is True
+
+ def test_reject_other_unsupported_languages(self, executor_manager):
+ """Test that other arbitrary languages still work (no R-specific check)."""
+ # Languages other than r-tuple and r-table should be allowed to pass
+ # the assertion, though they may fail at code execution
+ try:
+ executor_manager.initialize_executor(
+ code=SAMPLE_OPERATOR_CODE,
+ is_source=False,
+ language="javascript", # arbitrary language
+ )
+ # If we get here, the assertion passed (which is correct behavior)
+ # But the code execution might fail, which is fine
+ except AssertionError:
+ # Should NOT raise AssertionError for non-R languages
+ pytest.fail("Should not raise AssertionError for non-R languages")
+ except Exception:
+ # Other exceptions (like import errors) are expected and acceptable
+ pass
+
+ def test_gen_module_file_name_increments(self, executor_manager):
+ """Test that module file names increment correctly."""
+ module1, file1 = executor_manager.gen_module_file_name()
+ assert module1 == "udf-v1"
+ assert file1 == "udf-v1.py"
+
+ module2, file2 = executor_manager.gen_module_file_name()
+ assert module2 == "udf-v2"
+ assert file2 == "udf-v2.py"
+
+ module3, file3 = executor_manager.gen_module_file_name()
+ assert module3 == "udf-v3"
+ assert file3 == "udf-v3.py"
+
+ def test_is_concrete_operator_static_method(self):
+ """Test the is_concrete_operator static method."""
+ from core.models import TupleOperatorV2
+
+ # Should return True for concrete operator classes
+ # Note: We can't easily test with actual concrete classes here without imports
+ # This test just verifies the method exists and is callable
+ assert hasattr(ExecutorManager, "is_concrete_operator")
+ assert callable(ExecutorManager.is_concrete_operator)
+
+ # Test with non-class
+ assert ExecutorManager.is_concrete_operator("not a class") is False
+ assert ExecutorManager.is_concrete_operator(123) is False
+
+ # Test with abstract base classes (TupleOperatorV2 has abstract methods)
+ assert ExecutorManager.is_concrete_operator(TupleOperatorV2) is False
+
+ def test_regular_operator_is_not_source(self, executor_manager):
+ """Test that regular operator with is_source=False works correctly."""
+ executor_manager.initialize_executor(
+ code=SAMPLE_OPERATOR_CODE, is_source=False, language="python"
+ )
+ assert executor_manager.executor.is_source is False
+
+ def test_source_operator_mismatch_raises_error(self, executor_manager):
+ """Test that mismatched source operator flag raises AssertionError."""
+ with pytest.raises(AssertionError) as exc_info:
+ executor_manager.initialize_executor(
+ code=SAMPLE_OPERATOR_CODE,
+ is_source=True, # Wrong: regular operator but marked as source
+ language="python",
+ )
+ assert "SourceOperator API" in str(exc_info.value)
diff --git a/amber/src/main/python/core/architecture/managers/test_pause_manager.py b/amber/src/main/python/core/architecture/managers/test_pause_manager.py
new file mode 100644
index 00000000000..501c1335005
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/test_pause_manager.py
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from core.architecture.managers import StateManager
+from core.architecture.managers.pause_manager import PauseManager, PauseType
+from core.models import InternalQueue
+from proto.org.apache.texera.amber.engine.architecture.worker import WorkerState
+
+
+class TestPauseManager:
+ @pytest.fixture
+ def input_queue(self):
+ return InternalQueue()
+
+ @pytest.fixture
+ def state_manager(self):
+ return StateManager(
+ {
+ WorkerState.UNINITIALIZED: {WorkerState.READY},
+ WorkerState.READY: {WorkerState.PAUSED, WorkerState.RUNNING},
+ WorkerState.RUNNING: {WorkerState.PAUSED, WorkerState.COMPLETED},
+ WorkerState.PAUSED: {WorkerState.RUNNING},
+ WorkerState.COMPLETED: set(),
+ },
+ WorkerState.READY, # initial state set to READY for testing purpose
+ )
+
+ @pytest.fixture
+ def pause_manager(self, input_queue, state_manager):
+ return PauseManager(input_queue, state_manager)
+
+ def test_it_can_init(self, pause_manager):
+ pass
+
+ def test_it_is_not_paused_initially(self, pause_manager):
+ assert not pause_manager.is_paused()
+
+ def test_it_can_be_paused_and_resumed(self, pause_manager):
+ pause_manager.pause(PauseType.USER_PAUSE)
+ assert pause_manager.is_paused()
+ pause_manager.resume(PauseType.USER_PAUSE)
+ assert not pause_manager.is_paused()
+
+ def test_it_can_be_paused_when_paused(self, pause_manager):
+ pause_manager.pause(PauseType.USER_PAUSE)
+ assert pause_manager.is_paused()
+ pause_manager.pause(PauseType.USER_PAUSE)
+ assert pause_manager.is_paused()
+ pause_manager.resume(PauseType.USER_PAUSE)
+ assert not pause_manager.is_paused()
+
+ def test_it_can_be_resumed_when_resumed(self, pause_manager):
+ pause_manager.pause(PauseType.USER_PAUSE)
+ assert pause_manager.is_paused()
+ pause_manager.resume(PauseType.USER_PAUSE)
+ assert not pause_manager.is_paused()
+ pause_manager.resume(PauseType.USER_PAUSE)
+ assert not pause_manager.is_paused()
diff --git a/amber/src/main/python/core/architecture/managers/test_state_manager.py b/amber/src/main/python/core/architecture/managers/test_state_manager.py
new file mode 100644
index 00000000000..9ebe6e847c8
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/test_state_manager.py
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from core.architecture.managers.state_manager import (
+ InvalidStateException,
+ InvalidTransitionException,
+ StateManager,
+)
+from proto.org.apache.texera.amber.engine.architecture.worker import WorkerState
+
+
+class TestStateManager:
+ @pytest.fixture
+ def state_manager(self):
+ return StateManager(
+ {
+ WorkerState.UNINITIALIZED: {WorkerState.READY},
+ WorkerState.READY: {WorkerState.PAUSED, WorkerState.RUNNING},
+ WorkerState.RUNNING: {WorkerState.PAUSED, WorkerState.COMPLETED},
+ WorkerState.PAUSED: {WorkerState.RUNNING},
+ WorkerState.COMPLETED: set(),
+ },
+ WorkerState.UNINITIALIZED,
+ )
+
+ def test_it_can_init(self, state_manager):
+ pass
+
+ def test_it_can_transit_to_defined_state(self, state_manager):
+ state_manager.assert_state(WorkerState.UNINITIALIZED)
+ for state in [
+ WorkerState.READY,
+ WorkerState.PAUSED,
+ WorkerState.RUNNING,
+ WorkerState.COMPLETED,
+ ]:
+ state_manager.transit_to(state)
+ assert state_manager.confirm_state(state)
+ state_manager.assert_state(state)
+
+ def test_it_raises_exception_when_transit_to_undefined_state(self, state_manager):
+ state_manager.assert_state(WorkerState.UNINITIALIZED)
+ for state in [WorkerState.READY, WorkerState.PAUSED]:
+ state_manager.transit_to(state)
+ assert state_manager.confirm_state(state)
+ state_manager.assert_state(state)
+ with pytest.raises(InvalidTransitionException):
+ state_manager.transit_to(WorkerState.READY)
+
+ def test_it_raises_exception_when_asserting_a_different_state(self, state_manager):
+ state_manager.assert_state(WorkerState.UNINITIALIZED)
+ for state in [WorkerState.READY, WorkerState.PAUSED]:
+ state_manager.transit_to(state)
+ assert state_manager.confirm_state(state)
+ state_manager.assert_state(state)
+
+ with pytest.raises(InvalidStateException):
+ state_manager.assert_state(WorkerState.COMPLETED)
diff --git a/amber/src/main/python/core/architecture/managers/tuple_processing_manager.py b/amber/src/main/python/core/architecture/managers/tuple_processing_manager.py
new file mode 100644
index 00000000000..a67949e6717
--- /dev/null
+++ b/amber/src/main/python/core/architecture/managers/tuple_processing_manager.py
@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from threading import Event, Condition
+from typing import Optional, Tuple, Iterator
+
+from core.models import InternalMarker
+from proto.org.apache.texera.amber.core import PortIdentity
+
+
+class TupleProcessingManager:
+ def __init__(self):
+ self.current_input_tuple: Optional[Tuple] = None
+ self.current_input_port_id: Optional[PortIdentity] = None
+ self.current_input_tuple_iter: Optional[Iterator[Tuple]] = None
+ self.current_output_tuple: Optional[Tuple] = None
+ self.current_internal_marker: Optional[InternalMarker] = None
+ self.context_switch_condition: Condition = Condition()
+ self.finished_current: Event = Event()
+
+ def get_internal_marker(self) -> Optional[InternalMarker]:
+ ret, self.current_internal_marker = self.current_internal_marker, None
+ return ret
+
+ def get_input_tuple(self) -> Optional[Tuple]:
+ ret, self.current_input_tuple = self.current_input_tuple, None
+ return ret
+
+ def get_output_tuple(self) -> Optional[Tuple]:
+ ret, self.current_output_tuple = self.current_output_tuple, None
+ return ret
+
+ def get_input_port_id(self) -> int:
+ port_id = self.current_input_port_id
+ # no upstream, special case for source executor.
+ if port_id is None:
+ return 0
+ return port_id.id
diff --git a/amber/src/main/python/core/architecture/packaging/__init__.py b/amber/src/main/python/core/architecture/packaging/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/architecture/packaging/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/architecture/packaging/input_manager.py b/amber/src/main/python/core/architecture/packaging/input_manager.py
new file mode 100644
index 00000000000..6cb6bdc08c4
--- /dev/null
+++ b/amber/src/main/python/core/architecture/packaging/input_manager.py
@@ -0,0 +1,175 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import threading
+from pyarrow.lib import Table
+from typing import Iterator, Optional, Union, Dict, List, Set
+
+from core.models import Tuple, ArrowTableTupleProvider, Schema, InternalQueue
+from core.models.internal_marker import InternalMarker
+from core.models.payload import DataFrame, DataPayload, StateFrame
+from core.storage.runnables.input_port_materialization_reader_runnable import (
+ InputPortMaterializationReaderRunnable,
+)
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ PortIdentity,
+ ChannelIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import Partitioning
+
+
+class Channel:
+ def __init__(self):
+ self.port_id: Optional[PortIdentity] = None
+
+ def set_port_id(self, port_id: PortIdentity) -> None:
+ self.port_id = port_id
+
+
+class WorkerPort:
+ def __init__(self, schema: Schema):
+ self._schema = schema
+ self._channels: Set[ChannelIdentity] = set()
+ self.completed = False
+
+ def add_channel(self, channel: ChannelIdentity) -> None:
+ self._channels.add(channel)
+
+ def get_channels(self) -> Set[ChannelIdentity]:
+ return self._channels
+
+ def get_schema(self) -> Schema:
+ return self._schema
+
+
+class InputManager:
+ SOURCE_STARTER = ActorVirtualIdentity("SOURCE_STARTER")
+
+ def __init__(self, worker_id: str, input_queue: InternalQueue):
+ self.worker_id = worker_id
+ self._ports: Dict[PortIdentity, WorkerPort] = dict()
+ self._channels: Dict[ChannelIdentity, Channel] = dict()
+ self._current_channel_id: Optional[ChannelIdentity] = None
+ self._input_queue = input_queue
+ self._input_port_mat_reader_runnables: Dict[
+ PortIdentity, List[InputPortMaterializationReaderRunnable]
+ ] = dict() # TODO: Merge this into WorkerPort
+
+ def complete_current_port(self, channel_id: ChannelIdentity) -> None:
+ channel = self._channels[channel_id]
+ self._ports[channel.port_id].completed = True
+
+ def all_ports_completed(self) -> bool:
+ return all(port.completed for port in self._ports.values())
+
+ def set_up_input_port_mat_reader_threads(
+ self, port_id: PortIdentity, uris: List[str], partitionings: List[Partitioning]
+ ) -> None:
+ assert len(uris) == len(partitionings)
+ if uris is not None:
+ reader_runnables = [
+ InputPortMaterializationReaderRunnable(
+ uri=uri,
+ queue=self._input_queue,
+ worker_actor_id=ActorVirtualIdentity(self.worker_id),
+ partitioning=partitioning,
+ )
+ for uri, partitioning in zip(uris, partitionings)
+ ]
+ self._input_port_mat_reader_runnables[port_id] = reader_runnables
+
+ def get_input_port_mat_reader_threads(
+ self,
+ ) -> Dict[PortIdentity, List[InputPortMaterializationReaderRunnable]]:
+ return self._input_port_mat_reader_runnables
+
+ def start_input_port_mat_reader_threads(self):
+ for port_reader_runnables in self._input_port_mat_reader_runnables.values():
+ for reader_runnable in port_reader_runnables:
+ # A completed reader port should not be started again
+ if not reader_runnable.finished():
+ thread_for_reader_runnable = threading.Thread(
+ target=reader_runnable.run,
+ daemon=True,
+ name=f"port_mat_reader_runnable_thread_"
+ f"{reader_runnable.channel_id}",
+ )
+ thread_for_reader_runnable.start()
+
+ def get_all_channel_ids(self) -> Dict[ChannelIdentity, Channel].keys:
+ return self._channels.keys()
+
+ def get_all_data_channel_ids(self) -> Set[ChannelIdentity]:
+ return {key for key in self._channels if not key.is_control}
+
+ def add_input_port(
+ self,
+ port_id: PortIdentity,
+ schema: Schema,
+ storage_uris: List[str],
+ partitionings: List[Partitioning],
+ ) -> None:
+ if port_id.id is None:
+ port_id.id = 0
+ if port_id.internal is None:
+ port_id.internal = False
+
+ # each port can only be added and initialized once.
+ if port_id not in self._ports:
+ self._ports[port_id] = WorkerPort(schema)
+
+ self.set_up_input_port_mat_reader_threads(port_id, storage_uris, partitionings)
+
+ def get_port_id(self, channel_id: ChannelIdentity) -> PortIdentity:
+ return self._channels[channel_id].port_id
+
+ def get_port(self, port_id: PortIdentity) -> WorkerPort:
+ return self._ports[port_id]
+
+ def register_input(
+ self, channel_id: ChannelIdentity, port_id: PortIdentity
+ ) -> None:
+ if port_id.id is None:
+ port_id.id = 0
+ if port_id.internal is None:
+ port_id.internal = False
+ channel = Channel()
+ channel.set_port_id(port_id)
+ self._channels[channel_id] = channel
+ self._ports[port_id].add_channel(channel_id)
+
+ def process_data_payload(
+ self, from_: ChannelIdentity, payload: DataPayload
+ ) -> Iterator[Union[Tuple, InternalMarker]]:
+ self._current_channel_id = from_
+
+ if isinstance(payload, DataFrame):
+ yield from self._process_data(payload.frame)
+ elif isinstance(payload, StateFrame):
+ yield payload.frame
+ else:
+ raise NotImplementedError()
+
+ def _process_data(self, table: Table) -> Iterator[Tuple]:
+ schema = self._ports[
+ self._channels[self._current_channel_id].port_id
+ ].get_schema()
+ for field_accessor in ArrowTableTupleProvider(table):
+ yield Tuple(
+ {name: field_accessor for name in table.column_names}, schema=schema
+ )
diff --git a/amber/src/main/python/core/architecture/packaging/output_manager.py b/amber/src/main/python/core/architecture/packaging/output_manager.py
new file mode 100644
index 00000000000..bf4afbf396f
--- /dev/null
+++ b/amber/src/main/python/core/architecture/packaging/output_manager.py
@@ -0,0 +1,270 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import threading
+import typing
+from collections import OrderedDict
+from itertools import chain
+from loguru import logger
+from pyarrow import Table
+from queue import Queue
+from threading import Thread
+from typing import Iterable, Iterator
+from typing import Union
+
+from core.architecture.packaging.input_manager import WorkerPort, Channel
+from core.architecture.sendsemantics.broad_cast_partitioner import (
+ BroadcastPartitioner,
+)
+from core.architecture.sendsemantics.hash_based_shuffle_partitioner import (
+ HashBasedShufflePartitioner,
+)
+from core.architecture.sendsemantics.one_to_one_partitioner import OneToOnePartitioner
+from core.architecture.sendsemantics.partitioner import Partitioner
+from core.architecture.sendsemantics.range_based_shuffle_partitioner import (
+ RangeBasedShufflePartitioner,
+)
+from core.architecture.sendsemantics.round_robin_partitioner import (
+ RoundRobinPartitioner,
+)
+from core.models import Tuple, Schema, StateFrame
+from core.models.payload import DataPayload, DataFrame
+from core.models.state import State
+from core.storage.document_factory import DocumentFactory
+from core.storage.runnables.port_storage_writer import (
+ PortStorageWriter,
+ PortStorageWriterElement,
+)
+from core.util import get_one_of
+from core.util.virtual_identity import get_worker_index
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ PhysicalLink,
+ PortIdentity,
+ ChannelIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ HashBasedShufflePartitioning,
+ OneToOnePartitioning,
+ Partitioning,
+ RoundRobinPartitioning,
+ RangeBasedShufflePartitioning,
+ BroadcastPartitioning,
+)
+
+
+class OutputManager:
+ def __init__(self, worker_id: str):
+ self.worker_id = worker_id
+ self._partitioners: OrderedDict[PhysicalLink, Partitioning] = OrderedDict()
+ self._partitioning_to_partitioner: dict[
+ type(Partitioning), type(Partitioner)
+ ] = {
+ OneToOnePartitioning: OneToOnePartitioner,
+ RoundRobinPartitioning: RoundRobinPartitioner,
+ HashBasedShufflePartitioning: HashBasedShufflePartitioner,
+ RangeBasedShufflePartitioning: RangeBasedShufflePartitioner,
+ BroadcastPartitioning: BroadcastPartitioner,
+ }
+ self._ports: typing.Dict[PortIdentity, WorkerPort] = dict()
+ self._channels: typing.Dict[ChannelIdentity, Channel] = dict()
+ self._port_storage_writers: typing.Dict[
+ PortIdentity, typing.Tuple[Queue, PortStorageWriter, Thread]
+ ] = dict()
+
+ def is_missing_output_ports(self):
+ """
+ This method is only used for ensuring correct region execution.
+ Some operators may have input port dependency relationships, for
+ which we currently use a two-phase region execution scheme.
+ (See `RegionExecutionCoordinator.scala` for details.)
+ This logic will only be executed when the worker is part of an
+ `executingDependeePortPhase` region-execution phase.
+ We currently assume that in this phase the operator (worker) will
+ not output any data, hence no output ports.
+ However we still need to keep this worker open for the next
+ `executingNonDependeePortPhase` phase.
+ :return: Whether this worker currently does not have any output port.
+ """
+ return not self._ports
+
+ def add_output_port(
+ self,
+ port_id: PortIdentity,
+ schema: Schema,
+ storage_uri: typing.Optional[str] = None,
+ ) -> None:
+ if port_id.id is None:
+ port_id.id = 0
+ if port_id.internal is None:
+ port_id.internal = False
+
+ if storage_uri is not None:
+ self.set_up_port_storage_writer(port_id, storage_uri)
+
+ # each port can only be added and initialized once.
+ if port_id not in self._ports:
+ self._ports[port_id] = WorkerPort(schema)
+
+ def set_up_port_storage_writer(self, port_id: PortIdentity, storage_uri: str):
+ """
+ Create a separate thread for saving output tuples of a port
+ to storage in batch.
+ """
+ document, _ = DocumentFactory.open_document(storage_uri)
+ buffered_item_writer = document.writer(str(get_worker_index(self.worker_id)))
+ writer_queue = Queue()
+ port_storage_writer = PortStorageWriter(
+ buffered_item_writer=buffered_item_writer, queue=writer_queue
+ )
+ writer_thread = threading.Thread(
+ target=port_storage_writer.run,
+ daemon=True,
+ name=f"port_storage_writer_thread_{port_id}",
+ )
+ writer_thread.start()
+ self._port_storage_writers[port_id] = (
+ writer_queue,
+ port_storage_writer,
+ writer_thread,
+ )
+
+ def get_port(self, port_id=None) -> WorkerPort:
+ return list(self._ports.values())[0]
+
+ def get_port_ids(self) -> typing.List[PortIdentity]:
+ return list(self._ports.keys())
+
+ def get_output_channel_ids(self) -> typing.List[ChannelIdentity]:
+ return self._channels.keys()
+
+ def save_tuple_to_storage_if_needed(self, tuple_: Tuple, port_id=None) -> None:
+ """
+ Optionally write the tuple to storage if the specified output port
+ is determined by the scheduler to need storage. This method is not blocking
+ because a separate thread is used to flush the tuple to storage in batch.
+ :param tuple_: A tuple produced by the data processor.
+ :param port_id: If not specified, the tuple will be written to all
+ output ports that need storage.
+ :return:
+ """
+ if port_id is None:
+ for writer_queue, _, _ in self._port_storage_writers.values():
+ writer_queue.put(PortStorageWriterElement(data_tuple=tuple_))
+ elif port_id in self._port_storage_writers.keys():
+ self._port_storage_writers[port_id][0].put(
+ PortStorageWriterElement(data_tuple=tuple_)
+ )
+
+ def close_port_storage_writers(self) -> None:
+ """
+ Flush the buffers of port storage writers and wait for all the
+ writer threads to finish, which indicates the port storage writing
+ are finished.
+ """
+ for _, writer, _ in self._port_storage_writers.values():
+ # This non-blocking stop call will let the storage writers
+ # flush the remaining buffer
+ writer.stop()
+ for _, _, writer_thread in self._port_storage_writers.values():
+ # This blocking call will wait for all the writer to finish commit
+ writer_thread.join()
+
+ def add_partitioning(self, tag: PhysicalLink, partitioning: Partitioning) -> None:
+ """
+ Add down stream operator and its transfer policy
+ :param tag:
+ :param partitioning:
+ :return:
+ """
+ the_partitioning = get_one_of(partitioning)
+ logger.debug(f"adding {the_partitioning}")
+ for channel_id in the_partitioning.channels:
+ if channel_id.from_worker_id.name == self.worker_id:
+ # Explicitly set is_control to trigger lazy computation.
+ # If not set, it may be computed at different times,
+ # causing hash inconsistencies.
+ channel_id.is_control = False
+ self._channels[channel_id] = Channel()
+ partitioner = self._partitioning_to_partitioner[type(the_partitioning)]
+ self._partitioners[tag] = (
+ partitioner(the_partitioning)
+ if partitioner != OneToOnePartitioner
+ else partitioner(the_partitioning, self.worker_id)
+ )
+
+ def tuple_to_batch(
+ self, tuple_: Tuple
+ ) -> Iterator[typing.Tuple[ActorVirtualIdentity, DataFrame]]:
+ return chain(
+ *(
+ (
+ (receiver, self.tuple_to_frame(tuples))
+ for receiver, tuples in partitioner.add_tuple_to_batch(tuple_)
+ )
+ for partitioner in self._partitioners.values()
+ )
+ )
+
+ def emit_ecm(
+ self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage
+ ) -> Iterable[Union[DataPayload, EmbeddedControlMessage]]:
+ return chain(
+ *(
+ (
+ (
+ payload
+ if isinstance(payload, EmbeddedControlMessage)
+ else self.tuple_to_frame(payload)
+ )
+ for payload in partitioner.flush(to, ecm)
+ )
+ for partitioner in self._partitioners.values()
+ )
+ )
+
+ def emit_state(
+ self, state: State
+ ) -> Iterable[typing.Tuple[ActorVirtualIdentity, DataPayload]]:
+ return chain(
+ *(
+ (
+ (
+ receiver,
+ (
+ StateFrame(payload)
+ if isinstance(payload, State)
+ else self.tuple_to_frame(payload)
+ ),
+ )
+ for receiver, payload in partitioner.flush_state(state)
+ )
+ for partitioner in self._partitioners.values()
+ )
+ )
+
+ def tuple_to_frame(self, tuples: typing.List[Tuple]) -> DataFrame:
+ return DataFrame(
+ frame=Table.from_pydict(
+ {
+ name: [t.get_serialized_field(name) for t in tuples]
+ for name in self.get_port().get_schema().get_attr_names()
+ },
+ schema=self.get_port().get_schema().as_arrow_schema(),
+ )
+ )
diff --git a/amber/src/main/python/core/architecture/rpc/async_rpc_client.py b/amber/src/main/python/core/architecture/rpc/async_rpc_client.py
new file mode 100644
index 00000000000..f11c5afdea5
--- /dev/null
+++ b/amber/src/main/python/core/architecture/rpc/async_rpc_client.py
@@ -0,0 +1,253 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+import inspect
+from collections import defaultdict
+from concurrent.futures import Future
+from functools import wraps
+from loguru import logger
+from typing import Dict, TypeVar, Callable, Any, Coroutine
+
+from core.architecture.managers.context import Context
+from core.models.internal_queue import InternalQueue, DCMElement
+from core.util import set_one_of
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ AsyncRpcContext,
+ ReturnInvocation,
+ ControlReturn,
+ ControlInvocation,
+ ControllerServiceStub,
+ WorkerServiceStub,
+ ControlRequest,
+)
+from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2
+
+R = TypeVar("R")
+
+
+def async_run(func: Callable[..., Any]) -> Callable[..., Any]:
+ @wraps(func)
+ def wrapper(*args, **kwargs) -> Any:
+ try:
+ # Try to get the current running loop
+ if asyncio.get_running_loop():
+ return func(*args, **kwargs)
+ except RuntimeError:
+ # If there is no running loop, use asyncio.run to start one
+ return asyncio.run(func(*args, **kwargs))
+
+ return wrapper
+
+
+class AsyncRPCClient:
+ def __init__(self, output_queue: InternalQueue, context: Context):
+ self._context = context
+ self._output_queue = output_queue
+ self._send_sequences: Dict[ActorVirtualIdentity, int] = defaultdict(int)
+ self._unfulfilled_promises: Dict[(ActorVirtualIdentity, int), Future] = dict()
+ # TODO: is this correct?
+ self._controller_service_stub = ControllerServiceStub("")
+ rpc_context = AsyncRpcContext(
+ ActorVirtualIdentity(self._context.worker_id),
+ ActorVirtualIdentity(name="CONTROLLER"),
+ )
+ self._controller_service_stub._unary_unary = AsyncRPCClient._assign_context(
+ self, rpc_context
+ )
+ # Apply async_run to all async methods of the controller service stub
+ self._wrap_all_async_methods_with_async_run(self._controller_service_stub)
+
+ def _assign_context(
+ self, rpc_context: AsyncRpcContext
+ ) -> Callable[..., Coroutine[Any, Any, Future]]:
+ """Creates an async RPC wrapper function with a context"""
+
+ async def wrapper(
+ route: str, request, response_type, timeout, deadline, metadata
+ ):
+ to = rpc_context.receiver
+ control_command = ControlInvocation(
+ method_name=route.split("/")[-1], # Extract the method name for RPC
+ command=set_one_of(ControlRequest, request),
+ context=rpc_context,
+ command_id=self._send_sequences[to],
+ )
+ payload = set_one_of(DirectControlMessagePayloadV2, control_command)
+ self._output_queue.put(
+ DCMElement(
+ tag=ChannelIdentity(
+ ActorVirtualIdentity(self._context.worker_id), to, True
+ ),
+ payload=payload,
+ )
+ )
+ return self._create_future(to)
+
+ return wrapper
+
+ def _wrap_all_async_methods_with_async_run(self, instance: Any) -> None:
+ """Decorates all async methods of an instance with async_run."""
+ for attr_name in dir(instance):
+ attr = getattr(instance, attr_name)
+ if inspect.iscoroutinefunction(attr):
+ setattr(instance, attr_name, async_run(attr))
+
+ def controller_stub(self) -> ControllerServiceStub:
+ """
+ Returns a proxy for interacting with the controller interface.
+ """
+ return self._controller_service_stub
+
+ def get_worker_interface(self, target_worker) -> WorkerServiceStub:
+ """
+ Returns a proxy for interacting with a worker interface.
+
+ :param target_worker: The identifier for the target worker.
+ """
+ return self._create_proxy(
+ WorkerServiceStub, ActorVirtualIdentity(target_worker)
+ )
+
+ def _create_proxy(self, service_class, target_worker: ActorVirtualIdentity):
+ """
+ Creates a dynamic proxy for the given service class, allowing
+ asynchronous RPC communication with the specified target actor.
+
+ :param service_class: The service class to be proxied.
+ :param target: The target actor's identity.
+ :return: An instance of the proxy class.
+ """
+ rpc_client = self # to distinguish outer and inner self
+
+ class Proxy(service_class):
+ def __init__(self, target_actor: ActorVirtualIdentity):
+ self.target_actor = target_actor
+
+ async def _unary_unary(
+ self, route: str, request, response_type, *, timeout, deadline, metadata
+ ):
+ """
+ Handles unary-unary RPC calls by creating a ControlInvocation command
+ and sending it to the target actor.
+
+ :param route: The RPC route name.
+ :param request: The request message to be sent.
+ :param response_type: The expected response type (unused here).
+ :param timeout: The RPC call timeout (unused here).
+ :param deadline: The RPC call deadline (unused here).
+ :param metadata: Metadata for the RPC call (unused here).
+ :return: A future representing the RPC response.
+ """
+ rpc_context: AsyncRpcContext = AsyncRpcContext(
+ ActorVirtualIdentity(rpc_client._context.worker_id),
+ self.target_actor,
+ )
+ to = rpc_context.receiver
+ control_command = ControlInvocation(
+ # to align with java side, only use the method name
+ method_name=route.split("/")[-1],
+ command=set_one_of(ControlRequest, request),
+ context=rpc_context,
+ command_id=rpc_client._send_sequences[to],
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ control_command,
+ )
+ rpc_client._output_queue.put(
+ DCMElement(
+ tag=ChannelIdentity(
+ rpc_context.sender, rpc_context.receiver, True
+ ),
+ payload=payload,
+ )
+ )
+ return rpc_client._create_future(to)
+
+ def _stream_unary(self, *args, **kwargs):
+ """Block the _stream_unary method."""
+ raise NotImplementedError(
+ "Rpc call invokes _stream_unary, which is not supported."
+ )
+
+ def _unary_stream(self, *args, **kwargs):
+ """Block the _unary_stream method."""
+ raise NotImplementedError(
+ "Rpc call invokes _unary_stream, which is not supported."
+ )
+
+ def _stream_stream(self, *args, **kwargs):
+ """Block the _stream_stream method."""
+ raise NotImplementedError(
+ "Rpc call invokes _stream_stream, which is not supported."
+ )
+
+ return Proxy(target_worker)
+
+ def _create_future(self, to: ActorVirtualIdentity) -> Future:
+ """
+ Create a promise for the target actor, recording the CommandInvocations sent
+ with a sequence, so that the promise can be fulfilled once the
+ ReturnInvocation is received for the CommandInvocation.
+
+ :param to: ActorVirtualIdentity, the receiver.
+ """
+ future = Future()
+ self._unfulfilled_promises[(to, self._send_sequences[to])] = future
+ self._send_sequences[to] += 1
+ return future
+
+ def receive(
+ self, from_: ChannelIdentity, return_invocation: ReturnInvocation
+ ) -> None:
+ """
+ Receive the ReturnInvocation from the given actor.
+ :param from_: ChannelIdentity, the sender.
+ :param return_invocation: ReturnInvocationV2, the return to be processed.
+ """
+ command_id = return_invocation.command_id
+ self._fulfill_promise(from_, command_id, return_invocation.return_value)
+
+ def _fulfill_promise(
+ self,
+ from_: ChannelIdentity,
+ command_id: int,
+ control_return: ControlReturn,
+ ) -> None:
+ """
+ Fulfill the promise with the CommandInvocation, referenced by the sequence id
+ with this sender of ReturnInvocation.
+
+ :param from_: ChannelIdentity, the sender.
+ :param command_id: int, paired with from_ to uniquely identify an unfulfilled
+ future.
+ :param control_return: ControlReturnV2m, to be used to fulfill the promise.
+ """
+
+ future: Future = self._unfulfilled_promises.get(
+ (from_.from_worker_id, command_id)
+ )
+ if future is not None:
+ future.set_result(control_return)
+ del self._unfulfilled_promises[(from_.from_worker_id, command_id)]
+ else:
+ logger.warning(
+ f"received unknown ControlReturn {control_return}, no corresponding"
+ " ControlCommand found."
+ )
diff --git a/amber/src/main/python/core/architecture/rpc/async_rpc_handler_initializer.py b/amber/src/main/python/core/architecture/rpc/async_rpc_handler_initializer.py
new file mode 100644
index 00000000000..146cf91b0d3
--- /dev/null
+++ b/amber/src/main/python/core/architecture/rpc/async_rpc_handler_initializer.py
@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.architecture.handlers.control.add_input_channel_handler import (
+ AddInputChannelHandler,
+)
+from core.architecture.handlers.control.add_partitioning_handler import (
+ AddPartitioningHandler,
+)
+from core.architecture.handlers.control.assign_port_handler import AssignPortHandler
+from core.architecture.handlers.control.debug_command_handler import (
+ WorkerDebugCommandHandler,
+)
+from core.architecture.handlers.control.end_channel_handler import EndChannelHandler
+from core.architecture.handlers.control.end_worker_handler import EndWorkerHandler
+from core.architecture.handlers.control.evaluate_expression_handler import (
+ EvaluateExpressionHandler,
+)
+from core.architecture.handlers.control.initialize_executor_handler import (
+ InitializeExecutorHandler,
+)
+from core.architecture.handlers.control.no_operation_handler import NoOperationHandler
+from core.architecture.handlers.control.open_executor_handler import OpenExecutorHandler
+from core.architecture.handlers.control.pause_worker_handler import PauseWorkerHandler
+from core.architecture.handlers.control.query_statistics_handler import (
+ QueryStatisticsHandler,
+)
+from core.architecture.handlers.control.replay_current_tuple_handler import (
+ RetryCurrentTupleHandler,
+)
+from core.architecture.handlers.control.resume_worker_handler import ResumeWorkerHandler
+from core.architecture.handlers.control.start_channel_handler import StartChannelHandler
+from core.architecture.handlers.control.start_worker_handler import StartWorkerHandler
+from core.architecture.handlers.control.update_executor_handler import (
+ UpdateExecutorHandler,
+)
+
+
+class AsyncRPCHandlerInitializer(
+ AddInputChannelHandler,
+ AddPartitioningHandler,
+ AssignPortHandler,
+ WorkerDebugCommandHandler,
+ EvaluateExpressionHandler,
+ InitializeExecutorHandler,
+ OpenExecutorHandler,
+ PauseWorkerHandler,
+ QueryStatisticsHandler,
+ RetryCurrentTupleHandler,
+ ResumeWorkerHandler,
+ StartWorkerHandler,
+ EndWorkerHandler,
+ StartChannelHandler,
+ EndChannelHandler,
+ NoOperationHandler,
+ UpdateExecutorHandler,
+):
+ pass
diff --git a/amber/src/main/python/core/architecture/rpc/async_rpc_server.py b/amber/src/main/python/core/architecture/rpc/async_rpc_server.py
new file mode 100644
index 00000000000..49dc5f05472
--- /dev/null
+++ b/amber/src/main/python/core/architecture/rpc/async_rpc_server.py
@@ -0,0 +1,154 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+import grpclib.const
+from loguru import logger
+
+from core.architecture.managers.context import Context
+from core.architecture.rpc.async_rpc_handler_initializer import (
+ AsyncRPCHandlerInitializer,
+)
+from core.models.internal_queue import InternalQueue, DCMElement
+from core.util import get_one_of, set_one_of
+from proto.org.apache.texera.amber.core import ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ReturnInvocation,
+ ControlRequest,
+ ControlInvocation,
+ ControlReturn,
+ ControlError,
+ ErrorLanguage,
+)
+from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2
+
+
+class AsyncRPCServer:
+ def __init__(self, output_queue: InternalQueue, context: Context):
+ self._output_queue = output_queue
+ rpc_mapping = AsyncRPCHandlerInitializer(context).__mapping__()
+ self._handlers: dict[str, grpclib.const.Handler] = {
+ k.split("/")[-1].lower(): v for k, v in rpc_mapping.items()
+ }
+
+ def _wrap_as_stream(self, request: ControlRequest) -> grpclib.server.Stream:
+ """
+ Wraps a ControlRequest as a grpclib server Stream.
+
+ :param request: The ControlRequest to be wrapped.
+ :return: A Stream object that provides asynchronous send and receive methods.
+
+ This allows the incoming ControlRequest to be treated as a streaming request
+ for compatibility with grpclib's handler interface.
+ """
+
+ class ControlRequestStream(grpclib.server.Stream):
+ def __init__(self):
+ self.result = None
+
+ async def recv_message(self):
+ return request
+
+ async def send_message(self, msg):
+ self.result = msg
+
+ return ControlRequestStream()
+
+ def receive(self, from_: ChannelIdentity, control_invocation: ControlInvocation):
+ """
+ Handles incoming ControlInvocation messages by invoking the appropriate handler.
+
+ :param from_: The sender's ChannelIdentity.
+ :param control_invocation: The incoming ControlInvocation message.
+
+ This method performs the following steps:
+ 1. Extracts the command from the ControlInvocation.
+ 2. Looks up the corresponding handler for the method name.
+ 3. Wraps the command as a stream and runs the handler asynchronously.
+ 4. Constructs a ControlReturn or ControlError based on the handler's result.
+ 5. Sends the response back to the sender, unless no reply is needed.
+ """
+ command: ControlRequest = get_one_of(control_invocation.command)
+ method_name = control_invocation.method_name
+ logger.debug(f"PYTHON receives a ControlInvocation: {control_invocation}")
+ try:
+ # Look up the handler based on the lowercase method name.
+ handler: grpclib.const.Handler = self.look_up(method_name.lower())
+ # Wrap the command as a streaming request.
+ control_payload_stream = self._wrap_as_stream(command)
+ # Run the handler asynchronously.
+ asyncio.run(handler.func(control_payload_stream))
+ # Set up a ControlReturn from the handler's result.
+ control_return: ControlReturn = set_one_of(
+ ControlReturn, control_payload_stream.result
+ )
+
+ except Exception as exception:
+ # Handle exceptions and log the error.
+ logger.exception(exception)
+ # Construct a ControlError message in case of an exception.
+ control_return: ControlReturn = set_one_of(
+ ControlReturn,
+ ControlError(
+ error_message=str(exception), language=ErrorLanguage.PYTHON
+ ),
+ )
+
+ # Construct the payload as a ReturnInvocation.
+ payload: DirectControlMessagePayloadV2 = set_one_of(
+ DirectControlMessagePayloadV2,
+ ReturnInvocation(
+ command_id=control_invocation.command_id,
+ return_value=control_return,
+ ),
+ )
+
+ # Check if a reply is needed; if not, return early.
+ if self._no_reply_needed(control_invocation.command_id):
+ return
+
+ # Reply to the actor that originated this ControlInvocation, identified
+ # by control_invocation.context.sender. For a normal RPC over a
+ # control channel this matches `from_.from_worker_id`; for an
+ # invocation carried in-band by an ECM along a data channel, `from_`
+ # is the data channel between two workers and the original sender
+ # lives only in the invocation's context.
+ # When the context is unset (e.g. unit-test inputs that construct
+ # ControlInvocation directly), fall back to swapping `from_`.
+ ctx = control_invocation.context
+ if ctx.sender.name and ctx.receiver.name:
+ target_channel_id = ChannelIdentity(
+ ctx.receiver, ctx.sender, is_control=True
+ )
+ else:
+ target_channel_id = ChannelIdentity(
+ from_.to_worker_id, from_.from_worker_id, is_control=True
+ )
+ logger.debug(
+ f"PYTHON returns a ReturnInvocation {payload}, replying the command"
+ f" {command}"
+ )
+ # Put the control element in the output queue.
+ self._output_queue.put(DCMElement(tag=target_channel_id, payload=payload))
+
+ def look_up(self, method_name: str) -> grpclib.const.Handler:
+ logger.debug(method_name)
+ return self._handlers[method_name]
+
+ @staticmethod
+ def _no_reply_needed(command_id: int) -> bool:
+ return command_id < 0
diff --git a/amber/src/main/python/core/architecture/sendsemantics/__init__.py b/amber/src/main/python/core/architecture/sendsemantics/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/architecture/sendsemantics/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/architecture/sendsemantics/broad_cast_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/broad_cast_partitioner.py
new file mode 100644
index 00000000000..bd1b1d29ecf
--- /dev/null
+++ b/amber/src/main/python/core/architecture/sendsemantics/broad_cast_partitioner.py
@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import typing
+from overrides import overrides
+from typing import Iterator
+
+from core.architecture.sendsemantics.partitioner import Partitioner
+from core.models import Tuple
+from core.models.state import State
+from core.util import set_one_of
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ Partitioning,
+ BroadcastPartitioning,
+)
+
+
+class BroadcastPartitioner(Partitioner):
+ def __init__(self, partitioning: BroadcastPartitioning):
+ super().__init__(set_one_of(Partitioning, partitioning))
+ self.batch_size = partitioning.batch_size
+ self.batch: list[Tuple] = list()
+ self.receivers = list(
+ {channel.to_worker_id for channel in partitioning.channels}
+ )
+
+ @overrides
+ def add_tuple_to_batch(
+ self, tuple_: Tuple
+ ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]:
+ self.batch.append(tuple_)
+ if len(self.batch) == self.batch_size:
+ for receiver in self.receivers:
+ yield receiver, self.batch
+ self.reset()
+
+ @overrides
+ def flush(
+ self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage
+ ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]:
+ if len(self.batch) > 0:
+ for receiver in self.receivers:
+ if receiver == to:
+ yield self.batch
+ self.reset()
+ for receiver in self.receivers:
+ if receiver == to:
+ yield ecm
+
+ @overrides
+ def flush_state(
+ self, state: State
+ ) -> Iterator[
+ typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]]
+ ]:
+ if len(self.batch) > 0:
+ for receiver in self.receivers:
+ yield receiver, self.batch
+
+ self.reset()
+ for receiver in self.receivers:
+ yield receiver, state
+
+ @overrides
+ def reset(self) -> None:
+ self.batch = list()
diff --git a/amber/src/main/python/core/architecture/sendsemantics/hash_based_shuffle_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/hash_based_shuffle_partitioner.py
new file mode 100644
index 00000000000..de018e5a3ce
--- /dev/null
+++ b/amber/src/main/python/core/architecture/sendsemantics/hash_based_shuffle_partitioner.py
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import typing
+from loguru import logger
+from overrides import overrides
+from typing import Iterator
+
+from core.architecture.sendsemantics.partitioner import Partitioner
+from core.models import Tuple
+from core.models.state import State
+from core.util import set_one_of
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ HashBasedShufflePartitioning,
+ Partitioning,
+)
+
+
+class HashBasedShufflePartitioner(Partitioner):
+ def __init__(self, partitioning: HashBasedShufflePartitioning):
+ super().__init__(set_one_of(Partitioning, partitioning))
+ logger.debug(f"got {partitioning}")
+ self.batch_size = partitioning.batch_size
+ # Partitioning contains an ordered list of downstream worker ids.
+ # Currently we are using the index of such an order to choose
+ # a downstream worker to send tuples to.
+ # Must use dict.fromkeys to ensure the order of receiver workers
+ # from partitioning is preserved (using `{}` to create a set
+ # does not preserve order and will not work correctly.)
+ self.receivers = [
+ (rid, [])
+ for rid in dict.fromkeys(
+ channel.to_worker_id for channel in partitioning.channels
+ )
+ ]
+ self.hash_attribute_names = partitioning.hash_attribute_names
+
+ @overrides
+ def add_tuple_to_batch(
+ self, tuple_: Tuple
+ ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]:
+ partial_tuple = (
+ tuple_
+ if not self.hash_attribute_names
+ else tuple_.get_partial_tuple(self.hash_attribute_names)
+ )
+ hash_code = hash(partial_tuple) % len(self.receivers)
+ receiver, batch = self.receivers[hash_code]
+ batch.append(tuple_)
+ if len(batch) == self.batch_size:
+ yield receiver, batch
+ self.receivers[hash_code] = (receiver, list())
+
+ @overrides
+ def flush(
+ self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage
+ ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]:
+ for receiver, batch in self.receivers:
+ if receiver == to:
+ if len(batch) > 0:
+ yield batch
+ yield ecm
+
+ @overrides
+ def flush_state(
+ self, state: State
+ ) -> Iterator[
+ typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]]
+ ]:
+ for receiver, batch in self.receivers:
+ if len(batch) > 0:
+ yield receiver, batch
+ yield receiver, state
diff --git a/amber/src/main/python/core/architecture/sendsemantics/one_to_one_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/one_to_one_partitioner.py
new file mode 100644
index 00000000000..d6a568be6a3
--- /dev/null
+++ b/amber/src/main/python/core/architecture/sendsemantics/one_to_one_partitioner.py
@@ -0,0 +1,75 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import typing
+from overrides import overrides
+from typing import Iterator
+
+from core.architecture.sendsemantics.partitioner import Partitioner
+from core.models import Tuple
+from core.models.state import State
+from core.util import set_one_of
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ OneToOnePartitioning,
+ Partitioning,
+)
+
+
+class OneToOnePartitioner(Partitioner):
+ def __init__(self, partitioning: OneToOnePartitioning, worker_id: str):
+ super().__init__(set_one_of(Partitioning, partitioning))
+ self.batch_size = partitioning.batch_size
+ self.batch: list[Tuple] = list()
+ for channel in partitioning.channels:
+ if channel.from_worker_id.name == worker_id:
+ self.receiver = channel.to_worker_id
+ break # one to one will have only one receiver.
+
+ @overrides
+ def add_tuple_to_batch(
+ self, tuple_: Tuple
+ ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]:
+ self.batch.append(tuple_)
+ if len(self.batch) == self.batch_size:
+ yield self.receiver, self.batch
+ self.reset()
+
+ @overrides
+ def flush(
+ self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage
+ ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]:
+ if len(self.batch) > 0:
+ yield self.batch
+ self.reset()
+ yield ecm
+
+ @overrides
+ def flush_state(
+ self, state: State
+ ) -> Iterator[
+ typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]]
+ ]:
+ if len(self.batch) > 0:
+ yield self.receiver, self.batch
+ self.reset()
+ yield self.receiver, state
+
+ @overrides
+ def reset(self) -> None:
+ self.batch = list()
diff --git a/amber/src/main/python/core/architecture/sendsemantics/partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/partitioner.py
new file mode 100644
index 00000000000..c4aac57cbe4
--- /dev/null
+++ b/amber/src/main/python/core/architecture/sendsemantics/partitioner.py
@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import typing
+from abc import ABC
+from betterproto import Message
+from typing import Iterator
+
+from core.models import Tuple
+from core.models.state import State
+from core.util import get_one_of
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import Partitioning
+
+
+class Partitioner(ABC):
+ def __init__(self, partitioning: Message):
+ self.partitioning: Partitioning = get_one_of(partitioning)
+
+ def add_tuple_to_batch(
+ self, tuple_: Tuple
+ ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]:
+ pass
+
+ def flush(
+ self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage
+ ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]:
+ pass
+
+ def flush_state(
+ self, state: State
+ ) -> Iterator[
+ typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]]
+ ]:
+ pass
+
+ def reset(self) -> None:
+ pass
+
+ def __repr__(self):
+ return f"Partitioner[partitioning={self.partitioning}]"
diff --git a/amber/src/main/python/core/architecture/sendsemantics/range_based_shuffle_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/range_based_shuffle_partitioner.py
new file mode 100644
index 00000000000..28aff35935f
--- /dev/null
+++ b/amber/src/main/python/core/architecture/sendsemantics/range_based_shuffle_partitioner.py
@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import typing
+from loguru import logger
+from overrides import overrides
+from typing import Iterator
+
+from core.architecture.sendsemantics.partitioner import Partitioner
+from core.models import Tuple
+from core.models.state import State
+from core.util import set_one_of
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ RangeBasedShufflePartitioning,
+ Partitioning,
+)
+
+
+class RangeBasedShufflePartitioner(Partitioner):
+ def __init__(self, partitioning: RangeBasedShufflePartitioning):
+ super().__init__(set_one_of(Partitioning, partitioning))
+ logger.info(f"got {partitioning}")
+ self.batch_size = partitioning.batch_size
+ # Partitioning contains an ordered list of downstream worker ids.
+ # Currently we are using the index of such an order to choose
+ # a downstream worker to send tuples to.
+ # Must use dict.fromkeys to ensure the order of receiver workers
+ # from partitioning is preserved (using `{}` to create a set
+ # does not preserve order and will not work correctly.)
+ self.receivers = [
+ (rid, [])
+ for rid in dict.fromkeys(
+ channel.to_worker_id for channel in partitioning.channels
+ )
+ ]
+ self.range_attribute_names = partitioning.range_attribute_names
+ self.range_min = partitioning.range_min
+ self.range_max = partitioning.range_max
+ self.keys_per_receiver = int(
+ (
+ (partitioning.range_max - partitioning.range_min)
+ // len(partitioning.channels)
+ )
+ + 1
+ )
+
+ def get_receiver_index(self, column_val) -> int:
+ if column_val < self.range_min:
+ return 0
+ elif column_val > self.range_max:
+ return len(self.receivers) - 1
+ else:
+ return int((column_val - self.range_min) // self.keys_per_receiver)
+
+ @overrides
+ def add_tuple_to_batch(
+ self, tuple_: Tuple
+ ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]:
+ column_val = tuple_[self.range_attribute_names[0]]
+ receiver_index = self.get_receiver_index(column_val)
+ receiver, batch = self.receivers[receiver_index]
+ batch.append(tuple_)
+ if len(batch) == self.batch_size:
+ yield receiver, batch
+ self.receivers[receiver_index] = (receiver, list())
+
+ @overrides
+ def flush(
+ self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage
+ ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]:
+ for receiver, batch in self.receivers:
+ if receiver == to:
+ if len(batch) > 0:
+ yield batch
+ yield ecm
+
+ @overrides
+ def flush_state(
+ self, state: State
+ ) -> Iterator[
+ typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]]
+ ]:
+ for receiver, batch in self.receivers:
+ if len(batch) > 0:
+ yield receiver, batch
+ yield receiver, state
diff --git a/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py
new file mode 100644
index 00000000000..87c3fee87d8
--- /dev/null
+++ b/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py
@@ -0,0 +1,85 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import typing
+from overrides import overrides
+from typing import Iterator
+
+from core.architecture.sendsemantics.partitioner import Partitioner
+from core.models import Tuple
+from core.models.state import State
+from core.util import set_one_of
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ Partitioning,
+ RoundRobinPartitioning,
+)
+
+
+class RoundRobinPartitioner(Partitioner):
+ def __init__(self, partitioning: RoundRobinPartitioning):
+ super().__init__(set_one_of(Partitioning, partitioning))
+ self.batch_size = partitioning.batch_size
+ # Partitioning contains an ordered list of downstream worker ids.
+ # Currently we are using the index of such an order to choose
+ # a downstream worker to send tuples to.
+ # Must use dict.fromkeys to ensure the order of receiver workers
+ # from partitioning is preserved (using `{}` to create a set
+ # does not preserve order and will not work with input-port
+ # materialization reader threads.)
+ self.receivers = [
+ (rid, [])
+ for rid in dict.fromkeys(
+ channel.to_worker_id for channel in partitioning.channels
+ )
+ ]
+ self.round_robin_index = 0
+
+ @overrides
+ def add_tuple_to_batch(
+ self, tuple_: Tuple
+ ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]:
+ receiver, batch = self.receivers[self.round_robin_index]
+ batch.append(tuple_)
+ if len(batch) == self.batch_size:
+ yield receiver, batch
+ self.receivers[self.round_robin_index] = (receiver, list())
+ self.round_robin_index = (self.round_robin_index + 1) % len(self.receivers)
+
+ @overrides
+ def flush(
+ self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage
+ ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]:
+ for receiver, batch in self.receivers:
+ if receiver == to:
+ if len(batch) > 0:
+ yield batch
+ batch.clear()
+ yield ecm
+
+ @overrides
+ def flush_state(
+ self, state: State
+ ) -> Iterator[
+ typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]]
+ ]:
+ for receiver, batch in self.receivers:
+ if len(batch) > 0:
+ yield receiver, batch
+ batch.clear()
+ yield receiver, state
diff --git a/amber/src/main/python/core/models/__init__.py b/amber/src/main/python/core/models/__init__.py
new file mode 100644
index 00000000000..d24fe0a277d
--- /dev/null
+++ b/amber/src/main/python/core/models/__init__.py
@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import builtins
+from inspect import Traceback
+from typing import NamedTuple
+
+from .internal_queue import InternalQueue
+from .internal_marker import InternalMarker
+from .tuple import Tuple, TupleLike, ArrowTableTupleProvider
+from .table import Table, TableLike
+from .batch import Batch, BatchLike
+from .schema import AttributeType, Field, Schema
+from .state import State
+from .operator import (
+ Operator,
+ TableOperator,
+ TupleOperatorV2,
+ BatchOperator,
+ SourceOperator,
+)
+from .payload import DataFrame, DataPayload, StateFrame
+
+
+class ExceptionInfo(NamedTuple):
+ exc: builtins.type
+ value: Exception
+ tb: Traceback
+
+
+__all__ = [
+ "InternalQueue",
+ "InternalMarker",
+ "Tuple",
+ "TupleLike",
+ "ArrowTableTupleProvider",
+ "Table",
+ "TableLike",
+ "Batch",
+ "BatchLike",
+ "Operator",
+ "TupleOperatorV2",
+ "TableOperator",
+ "BatchOperator",
+ "SourceOperator",
+ "DataFrame",
+ "DataPayload",
+ "StateFrame",
+ "ExceptionInfo",
+ "AttributeType",
+ "Field",
+ "Schema",
+ "State",
+]
diff --git a/amber/src/main/python/core/models/batch.py b/amber/src/main/python/core/models/batch.py
new file mode 100644
index 00000000000..efa8a7b0ab5
--- /dev/null
+++ b/amber/src/main/python/core/models/batch.py
@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pandas
+from typing import TypeVar
+
+BatchLike = TypeVar("BatchLike", pandas.DataFrame, pandas.DataFrame)
+
+
+class Batch(pandas.DataFrame):
+ def __init__(self, batch_like: BatchLike):
+ super().__init__(batch_like)
diff --git a/amber/src/main/python/core/models/internal_marker.py b/amber/src/main/python/core/models/internal_marker.py
new file mode 100644
index 00000000000..6c9c80bafc4
--- /dev/null
+++ b/amber/src/main/python/core/models/internal_marker.py
@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+class InternalMarker:
+ """
+ A special Data Message, only being generated in un-packaging a batch into Tuples.
+ Markers retain the order information and served as a indicator of data state.
+ """
+
+ pass
+
+
+class StartChannel(InternalMarker):
+ pass
+
+
+class EndChannel(InternalMarker):
+ pass
diff --git a/amber/src/main/python/core/models/internal_queue.py b/amber/src/main/python/core/models/internal_queue.py
new file mode 100644
index 00000000000..abc1793ff6c
--- /dev/null
+++ b/amber/src/main/python/core/models/internal_queue.py
@@ -0,0 +1,161 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from threading import RLock
+from typing import TypeVar, Set
+
+from core.models.internal_marker import InternalMarker
+from core.models.payload import DataPayload
+from core.util.customized_queue.linked_blocking_multi_queue import (
+ LinkedBlockingMultiQueue,
+)
+from core.util.customized_queue.queue_base import IQueue, QueueElement
+from proto.org.apache.texera.amber.core import ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2
+
+
+@dataclass
+class InternalQueueElement(QueueElement):
+ tag: ChannelIdentity
+
+
+@dataclass
+class DataElement(InternalQueueElement):
+ payload: DataPayload
+
+
+@dataclass
+class DCMElement(InternalQueueElement):
+ payload: DirectControlMessagePayloadV2
+
+
+@dataclass
+class ECMElement(InternalQueueElement):
+ payload: EmbeddedControlMessage
+
+
+T = TypeVar("T", bound=InternalQueueElement)
+
+
+class InternalQueue(IQueue):
+ class DisableType(Enum):
+ DISABLE_BY_PAUSE = 1
+ DISABLE_BY_BACKPRESSURE = 2
+
+ def __init__(self):
+ self._queue = LinkedBlockingMultiQueue()
+ self._queue.add_sub_queue("SYSTEM", 0)
+ self._queue_ids: Set[ChannelIdentity] = set()
+ self._queue_state: Set[InternalQueue.DisableType] = set()
+ self._lock = RLock()
+
+ def is_empty(self, key=None) -> bool:
+ return self._queue.is_empty(key)
+
+ def get(self) -> T:
+ return self._queue.get()
+
+ def put(self, item: T) -> None:
+ if isinstance(item, InternalQueueElement):
+ if item.tag not in self._queue_ids:
+ self._queue.add_sub_queue(item.tag, 1 if item.tag.is_control else 2)
+ self._queue_ids.add(item.tag)
+ if isinstance(item, (DataElement, InternalMarker, ECMElement)):
+ self._queue.put(item.tag, item)
+ elif isinstance(item, DCMElement):
+ self._queue.put(item.tag, item)
+ else:
+ raise ValueError(f"item {item} is not recognized by internal queue")
+ else:
+ self._queue.put("SYSTEM", item)
+
+ def disable(self, channel_id: ChannelIdentity) -> None:
+ self._queue.disable(channel_id)
+
+ def enable(self, channel_id: ChannelIdentity) -> None:
+ self._queue.enable(channel_id)
+
+ def is_control_empty(self) -> bool:
+ return all(
+ self.is_empty(queue_id)
+ for queue_id in self._queue_ids
+ if queue_id.is_control
+ )
+
+ def is_data_empty(self) -> bool:
+ return all(
+ self.is_empty(queue_id)
+ for queue_id in self._queue_ids
+ if not queue_id.is_control
+ )
+
+ def __len__(self) -> int:
+ return self.size()
+
+ def size(self) -> int:
+ return self._queue.size()
+
+ def size_control(self) -> int:
+ return sum(
+ self._queue.size(queue_id)
+ for queue_id in self._queue_ids
+ if queue_id.is_control
+ )
+
+ def size_data(self) -> int:
+ return sum(
+ self._queue.size(queue_id)
+ for queue_id in self._queue_ids
+ if not queue_id.is_control
+ )
+
+ def enable_data(self, disable_type: DisableType) -> bool:
+ with self._lock:
+ if disable_type in self._queue_state:
+ self._queue_state.remove(disable_type)
+ if self._queue_state:
+ return False
+ for queue_id in self._queue_ids:
+ if not queue_id.is_control:
+ self._queue.enable(queue_id)
+ return True
+
+ def disable_data(self, disable_type: DisableType) -> None:
+ with self._lock:
+ self._queue_state.add(disable_type)
+ for queue_id in self._queue_ids:
+ if not queue_id.is_control:
+ self._queue.disable(queue_id)
+
+ def in_mem_size(self) -> int:
+ return sum(
+ self._queue.in_mem_size(queue_id)
+ for queue_id in self._queue_ids
+ if not queue_id.is_control
+ )
+
+ def is_data_enabled(self) -> bool:
+ return any(
+ self._queue.is_enabled(queue_id)
+ for queue_id in self._queue_ids
+ if not queue_id.is_control
+ )
diff --git a/amber/src/main/python/core/models/operator.py b/amber/src/main/python/core/models/operator.py
new file mode 100644
index 00000000000..952e2a12c81
--- /dev/null
+++ b/amber/src/main/python/core/models/operator.py
@@ -0,0 +1,293 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import overrides
+import pandas
+from functools import lru_cache
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import Iterator, List, Mapping, Optional, Union, MutableMapping, Protocol
+
+from . import Table, TableLike, Tuple, TupleLike, Batch, BatchLike
+from .state import State
+from .table import all_output_to_tuple
+
+import base64
+
+
+class Operator(ABC):
+ """
+ Abstract base class for all operators.
+ """
+
+ class PythonTemplateDecoder:
+ class Decoder(Protocol):
+ """Pluggable base64 decoder interface."""
+
+ def to_str(self, data: Union[str, bytes]) -> str: ...
+
+ class StdlibBase64Decoder:
+ """Default decoder using Python's stdlib base64."""
+
+ def to_str(self, data: Union[str, bytes]) -> str:
+ b64_bytes = data.encode("ascii") if isinstance(data, str) else data
+ raw = base64.b64decode(b64_bytes, validate=False)
+ return raw.decode("utf-8", errors="strict")
+
+ def __init__(
+ self,
+ decoder: Optional["Operator.PythonTemplateDecoder.Decoder"] = None,
+ cache_size: int = 256,
+ ) -> None:
+ self._decoder = decoder or self.StdlibBase64Decoder()
+ self._decode_cached = self._build_cached_decoder(cache_size)
+
+ def _build_cached_decoder(self, cache_size: int):
+ @lru_cache(maxsize=cache_size)
+ def _cached(data: Union[str, bytes]) -> str:
+ return self._decoder.to_str(data)
+
+ return _cached
+
+ def decode(self, data: Union[str, bytes]) -> str:
+ return self._decode_cached(data)
+
+ def _get_template_decoder(self) -> "Operator.PythonTemplateDecoder":
+ if not hasattr(self, "_python_template_decoder"):
+ self._python_template_decoder = self.PythonTemplateDecoder(cache_size=256)
+ return self._python_template_decoder
+
+ def decode_python_template(self, data: Union[str, bytes]) -> str:
+ return self._get_template_decoder().decode(data)
+
+ __internal_is_source: bool = False
+
+ @property
+ @overrides.final
+ def is_source(self) -> bool:
+ """
+ Whether the operator is a source operator. Source operators generate output
+ Tuples without having input Tuples.
+
+ :return:
+ """
+ return self.__internal_is_source
+
+ @is_source.setter
+ @overrides.final
+ def is_source(self, value: bool) -> None:
+ self.__internal_is_source = value
+
+ def open(self) -> None:
+ """
+ Open a context of the operator. Usually can be used for loading/initiating some
+ resources, such as a file, a model, or an API client.
+ """
+ pass
+
+ def close(self) -> None:
+ """
+ Close the context of the operator.
+ """
+ pass
+
+ def process_state(self, state: State, port: int) -> Optional[State]:
+ """
+ Process an input State from the given link.
+ The default implementation is to pass the State to all downstream operators.
+ :param state: State, a State from an input port to be processed.
+ :param port: int, input port index of the current exhausted port.
+ :return: State, producing one State object
+ """
+ return state
+
+ def produce_state_on_start(self, port: int) -> Optional[State]:
+ """
+ Produce a State when the given link started.
+
+ :param port: int, input port index of the current initialized port.
+ :return: State, producing one State object
+ """
+ pass
+
+ def produce_state_on_finish(self, port: int) -> Optional[State]:
+ """
+ Produce a State after the input port is exhausted.
+
+ :param port: int, input port index of the current exhausted port.
+ :return: State, producing one State object
+ """
+ pass
+
+
+class TupleOperatorV2(Operator):
+ """
+ Base class for tuple-oriented operators. A concrete implementation must
+ be provided upon using.
+ """
+
+ @abstractmethod
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ """
+ Process an input Tuple from the given link.
+
+ :param tuple_: Tuple, a Tuple from an input port to be processed.
+ :param port: int, input port index of the current Tuple.
+ :return: Iterator[Optional[TupleLike]], producing one TupleLike object at a
+ time, or None.
+ """
+ yield
+
+ def on_finish(self, port: int) -> Iterator[Optional[TupleLike]]:
+ """
+ Callback when one input port is exhausted.
+
+ :param port: int, input port index of the current exhausted port.
+ :return: Iterator[Optional[TupleLike]], producing one TupleLike object at a
+ time, or None.
+ """
+ yield
+
+
+class SourceOperator(TupleOperatorV2):
+ __internal_is_source = True
+
+ @abstractmethod
+ def produce(self) -> Iterator[Union[TupleLike, TableLike, None]]:
+ """
+ Produce Tuples or Tables. Used by the source operator only.
+
+ :return: Iterator[Union[TupleLike, TableLike, None]], producing
+ one TupleLike object, one TableLike object, or None, at a time.
+ """
+ yield
+
+ @overrides.final
+ def on_finish(self, port: int) -> Iterator[Optional[TupleLike]]:
+ # TODO: change on_finish to output Iterator[Union[TupleLike, TableLike, None]]
+ for i in self.produce():
+ yield from all_output_to_tuple(i)
+
+ @overrides.final
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ yield
+
+
+class BatchOperator(TupleOperatorV2):
+ """
+ Base class for batch-oriented operators. A concrete implementation must
+ be provided upon using.
+ """
+
+ BATCH_SIZE: int = 10 # must be a positive integer
+
+ def __init__(self):
+ super().__init__()
+ self.__batch_data: MutableMapping[int, List[Tuple]] = defaultdict(list)
+ self._validate_batch_size(self.BATCH_SIZE)
+
+ @staticmethod
+ @overrides.final
+ def _validate_batch_size(value):
+ if value is None:
+ raise ValueError("BATCH_SIZE cannot be None.")
+ if type(value) is not int:
+ raise ValueError("BATCH_SIZE cannot be {type(value))}.")
+ if value <= 0:
+ raise ValueError("BATCH_SIZE should be positive.")
+
+ @overrides.final
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ self.__batch_data[port].append(tuple_)
+ if (
+ self.BATCH_SIZE is not None
+ and len(self.__batch_data[port]) >= self.BATCH_SIZE
+ ):
+ yield from self._process_batch(port)
+
+ @overrides.final
+ def _process_batch(self, port: int) -> Iterator[Optional[BatchLike]]:
+ batch = Batch(
+ pandas.DataFrame(
+ [
+ self.__batch_data[port].pop(0).as_series()
+ for _ in range(min(len(self.__batch_data[port]), self.BATCH_SIZE))
+ ]
+ )
+ )
+ for output_batch in self.process_batch(batch, port):
+ if output_batch is not None:
+ if isinstance(output_batch, pandas.DataFrame):
+ # TODO: integrate into Batch as a helper function.
+ # convert from Batch to Tuple, only supports pandas.DataFrames for
+ # now.
+ for _, output_tuple in output_batch.iterrows():
+ yield output_tuple
+ else:
+ yield output_batch
+
+ @overrides.final
+ def on_finish(self, port: int) -> Iterator[Optional[BatchLike]]:
+ while len(self.__batch_data[port]) != 0:
+ yield from self._process_batch(port)
+
+ @abstractmethod
+ def process_batch(self, batch: Batch, port: int) -> Iterator[Optional[BatchLike]]:
+ """
+ Process an input Batch from the given link. The Batch is represented as a
+ pandas.DataFrame.
+
+ :param batch: Batch, a batch to be processed.
+ :param port: int, input port index of the current Batch.
+ :return: Iterator[Optional[BatchLike]], producing one BatchLike object at a
+ time, or None.
+ """
+ yield
+
+
+class TableOperator(TupleOperatorV2):
+ """
+ Base class for table-oriented operators. A concrete implementation must
+ be provided upon using.
+ """
+
+ def __init__(self):
+ super().__init__()
+ self.__internal_is_source: bool = False
+ self.__table_data: Mapping[int, List[Tuple]] = defaultdict(list)
+
+ @overrides.final
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ self.__table_data[port].append(tuple_)
+ yield
+
+ def on_finish(self, port: int) -> Iterator[Optional[TableLike]]:
+ table = Table(self.__table_data[port])
+ yield from self.process_table(table, port)
+
+ @abstractmethod
+ def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:
+ """
+ Process an input Table from the given link. The Table is represented as a
+ pandas.DataFrame.
+
+ :param table: Table, a table to be processed.
+ :param port: int, input port index of the current Tuple.
+ :return: Iterator[Optional[TableLike]], producing one TableLike object at a
+ time, or None.
+ """
+ yield
diff --git a/amber/src/main/python/core/models/payload.py b/amber/src/main/python/core/models/payload.py
new file mode 100644
index 00000000000..61a33294882
--- /dev/null
+++ b/amber/src/main/python/core/models/payload.py
@@ -0,0 +1,36 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from dataclasses import dataclass
+from pyarrow.lib import Table
+
+from core.models.state import State
+
+
+@dataclass
+class DataPayload:
+ pass
+
+
+@dataclass
+class DataFrame(DataPayload):
+ frame: Table
+
+
+@dataclass
+class StateFrame(DataPayload):
+ frame: State
diff --git a/amber/src/main/python/core/models/schema/__init__.py b/amber/src/main/python/core/models/schema/__init__.py
new file mode 100644
index 00000000000..306fe4e1d4e
--- /dev/null
+++ b/amber/src/main/python/core/models/schema/__init__.py
@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .attribute_type import AttributeType
+from core.models.type.large_binary import largebinary
+from .field import Field
+from .schema import Schema
+
+
+__all__ = [
+ "AttributeType",
+ "largebinary",
+ "Field",
+ "Schema",
+]
diff --git a/amber/src/main/python/core/models/schema/arrow_schema_utils.py b/amber/src/main/python/core/models/schema/arrow_schema_utils.py
new file mode 100644
index 00000000000..527e0095e65
--- /dev/null
+++ b/amber/src/main/python/core/models/schema/arrow_schema_utils.py
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Utilities for converting between Arrow schemas and Amber schemas,
+handling LARGE_BINARY metadata preservation.
+"""
+
+import pyarrow as pa
+from typing import Mapping
+
+from core.models.schema.attribute_type import AttributeType
+from core.models.schema.attribute_type_utils import (
+ detect_attribute_type_from_arrow_field,
+ create_arrow_field_with_metadata,
+)
+
+
+def arrow_schema_to_attr_types(arrow_schema: pa.Schema) -> dict[str, AttributeType]:
+ """
+ Converts an Arrow schema to a dictionary of attribute name to AttributeType.
+ Handles LARGE_BINARY metadata detection.
+
+ :param arrow_schema: PyArrow schema that may contain LARGE_BINARY metadata
+ :return: Dictionary mapping attribute names to AttributeTypes
+ """
+ attr_types = {}
+ for attr_name in arrow_schema.names:
+ field = arrow_schema.field(attr_name)
+ attr_types[attr_name] = detect_attribute_type_from_arrow_field(field)
+ return attr_types
+
+
+def attr_types_to_arrow_schema(
+ attr_types: Mapping[str, AttributeType],
+) -> pa.Schema:
+ """
+ Converts a mapping of attribute name to AttributeType into an Arrow schema.
+ Adds metadata for LARGE_BINARY types.
+ Preserves the order of attributes from the input mapping.
+
+ :param attr_types: Mapping of attribute names to AttributeTypes (e.g., OrderedDict)
+ :return: PyArrow schema with metadata for LARGE_BINARY types
+ """
+ fields = [
+ create_arrow_field_with_metadata(attr_name, attr_type)
+ for attr_name, attr_type in attr_types.items()
+ ]
+ return pa.schema(fields)
diff --git a/amber/src/main/python/core/models/schema/attribute_type.py b/amber/src/main/python/core/models/schema/attribute_type.py
new file mode 100644
index 00000000000..24d0745f41e
--- /dev/null
+++ b/amber/src/main/python/core/models/schema/attribute_type.py
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import pyarrow as pa
+from bidict import bidict
+from enum import Enum
+from pyarrow import lib
+from core.models.type.large_binary import largebinary
+
+
+class AttributeType(Enum):
+ """
+ Types supported by PyTexera & PyAmber.
+
+ The definitions are mapped and following the AttributeType.java
+ (src/main/scala/org/apache/texera/workflow/common/tuple/schema/AttributeType.java)
+ """
+
+ STRING = 1
+ INT = 2
+ LONG = 3
+ BOOL = 4
+ DOUBLE = 5
+ TIMESTAMP = 6
+ BINARY = 7
+ LARGE_BINARY = 8
+
+
+RAW_TYPE_MAPPING = bidict(
+ {
+ "STRING": AttributeType.STRING,
+ "INTEGER": AttributeType.INT,
+ "LONG": AttributeType.LONG,
+ "DOUBLE": AttributeType.DOUBLE,
+ "BOOLEAN": AttributeType.BOOL,
+ "TIMESTAMP": AttributeType.TIMESTAMP,
+ "BINARY": AttributeType.BINARY,
+ "LARGE_BINARY": AttributeType.LARGE_BINARY,
+ }
+)
+
+TO_ARROW_MAPPING = {
+ AttributeType.INT: pa.int32(),
+ AttributeType.LONG: pa.int64(),
+ AttributeType.STRING: pa.string(),
+ AttributeType.DOUBLE: pa.float64(),
+ AttributeType.BOOL: pa.bool_(),
+ AttributeType.BINARY: pa.binary(),
+ AttributeType.TIMESTAMP: pa.timestamp("us"),
+ AttributeType.LARGE_BINARY: pa.string(), # Serialized as URI string
+}
+
+FROM_ARROW_MAPPING = {
+ lib.Type_INT32: AttributeType.INT,
+ lib.Type_INT64: AttributeType.LONG,
+ lib.Type_STRING: AttributeType.STRING,
+ lib.Type_LARGE_STRING: AttributeType.STRING,
+ lib.Type_DOUBLE: AttributeType.DOUBLE,
+ lib.Type_BOOL: AttributeType.BOOL,
+ lib.Type_BINARY: AttributeType.BINARY,
+ lib.Type_LARGE_BINARY: AttributeType.BINARY,
+ lib.Type_TIMESTAMP: AttributeType.TIMESTAMP,
+}
+
+
+# Only single-directional mapping.
+TO_PYOBJECT_MAPPING = {
+ AttributeType.STRING: str,
+ AttributeType.INT: int,
+ AttributeType.LONG: int, # Python3 unifies long into int.
+ AttributeType.DOUBLE: float,
+ AttributeType.BOOL: bool,
+ AttributeType.BINARY: bytes,
+ AttributeType.TIMESTAMP: datetime.datetime,
+ AttributeType.LARGE_BINARY: largebinary,
+}
+
+FROM_PYOBJECT_MAPPING = {
+ str: AttributeType.STRING,
+ int: AttributeType.INT,
+ float: AttributeType.DOUBLE,
+ bool: AttributeType.BOOL,
+ bytes: AttributeType.BINARY,
+ datetime.datetime: AttributeType.TIMESTAMP,
+ largebinary: AttributeType.LARGE_BINARY,
+}
diff --git a/amber/src/main/python/core/models/schema/attribute_type_utils.py b/amber/src/main/python/core/models/schema/attribute_type_utils.py
new file mode 100644
index 00000000000..3918fdfc342
--- /dev/null
+++ b/amber/src/main/python/core/models/schema/attribute_type_utils.py
@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Utilities for converting between AttributeTypes and Arrow field types,
+handling LARGE_BINARY metadata preservation.
+"""
+
+import pyarrow as pa
+
+from core.models.schema.attribute_type import (
+ AttributeType,
+ FROM_ARROW_MAPPING,
+ TO_ARROW_MAPPING,
+)
+
+# Metadata key used to mark LARGE_BINARY fields in Arrow schemas
+TEXERA_TYPE_METADATA_KEY = b"texera_type"
+LARGE_BINARY_METADATA_VALUE = b"LARGE_BINARY"
+
+
+def detect_attribute_type_from_arrow_field(field: pa.Field) -> AttributeType:
+ """
+ Detects the AttributeType from an Arrow field, checking metadata for LARGE_BINARY.
+
+ :param field: PyArrow field that may contain metadata
+ :return: The detected AttributeType
+ """
+ # Check metadata for LARGE_BINARY type
+ # (can be stored by either Scala ArrowUtils or Python)
+ is_large_binary = (
+ field.metadata
+ and field.metadata.get(TEXERA_TYPE_METADATA_KEY) == LARGE_BINARY_METADATA_VALUE
+ )
+
+ if is_large_binary:
+ return AttributeType.LARGE_BINARY
+ else:
+ return FROM_ARROW_MAPPING[field.type.id]
+
+
+def create_arrow_field_with_metadata(
+ attr_name: str, attr_type: AttributeType
+) -> pa.Field:
+ """
+ Creates a PyArrow field with appropriate metadata for the given AttributeType.
+
+ :param attr_name: Name of the attribute
+ :param attr_type: The AttributeType
+ :return: PyArrow field with metadata if needed
+ """
+ metadata = (
+ {TEXERA_TYPE_METADATA_KEY: LARGE_BINARY_METADATA_VALUE}
+ if attr_type == AttributeType.LARGE_BINARY
+ else None
+ )
+
+ return pa.field(attr_name, TO_ARROW_MAPPING[attr_type], metadata=metadata)
diff --git a/amber/src/main/python/core/models/schema/field.py b/amber/src/main/python/core/models/schema/field.py
new file mode 100644
index 00000000000..59fef5e4d67
--- /dev/null
+++ b/amber/src/main/python/core/models/schema/field.py
@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import TypeVar, Optional
+
+Field = TypeVar(
+ "Field",
+ Optional[str],
+ Optional[int], # for both INT and LONG
+ Optional[float],
+ Optional[bool],
+ Optional[datetime.datetime],
+ Optional[bytes],
+)
diff --git a/amber/src/main/python/core/models/schema/schema.py b/amber/src/main/python/core/models/schema/schema.py
new file mode 100644
index 00000000000..d349807ab79
--- /dev/null
+++ b/amber/src/main/python/core/models/schema/schema.py
@@ -0,0 +1,150 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+from collections import OrderedDict
+from typing import MutableMapping, Optional, Mapping, List, Tuple
+
+from core.models.schema.attribute_type import (
+ AttributeType,
+ RAW_TYPE_MAPPING,
+)
+from core.models.schema.arrow_schema_utils import (
+ arrow_schema_to_attr_types,
+ attr_types_to_arrow_schema,
+)
+
+
+class Schema:
+ """
+ Schema describes a sequence of attributes, maintaining a name-to-type mapping.
+
+ Schema is mapped to PyArrow's Schema (pyarrow.Schema), with each
+ AttributeType mapped to a pyarrow.DataType.
+
+ Schema is mapped to a Tuple Field (which is a collection TypeVar of
+ python objects).
+
+ See AttributeType for detailed mappings.
+
+ Note: Schema is to be used by the engine only, and should be invisible to the
+ users of Tuple. It only gets assigned to a Tuple to finalize the Tuple for
+ serialization purpose.
+ """
+
+ def __init__(
+ self,
+ arrow_schema: Optional[pa.Schema] = None,
+ raw_schema: Optional[Mapping[str, str]] = None,
+ ):
+ self._name_type_mapping: MutableMapping[str, AttributeType] = OrderedDict()
+
+ if arrow_schema is not None:
+ self._from_arrow_schema(arrow_schema)
+ if raw_schema is not None:
+ self._from_raw_schema(raw_schema)
+
+ def add(self, attr_name: str, attr_type: AttributeType) -> None:
+ """
+ Append a new attribute with its name and type to the Schema.
+ :param attr_name: new attribute's name, must not be in the Schema already.
+ :param attr_type: the type of the attribute.
+ :return:
+ """
+ if attr_name in self._name_type_mapping:
+ raise KeyError(f"Adding a duplicated attribute {repr(attr_name)}.")
+ self._name_type_mapping[attr_name] = attr_type
+
+ def _from_raw_schema(self, raw_schema: Mapping[str, str]) -> None:
+ """
+ Resets the Schema by converting a raw schema.
+ :param raw_schema: a map of attr_name -> type_str.
+ :return:
+ """
+ self._name_type_mapping = OrderedDict()
+ for attr_name, raw_type in raw_schema.items():
+ attr_type = RAW_TYPE_MAPPING[raw_type]
+ self.add(attr_name, attr_type)
+
+ def _from_arrow_schema(self, arrow_schema: pa.Schema) -> None:
+ """
+ Resets the Schema by converting a pyarrow.Schema.
+ :param arrow_schema: a pyarrow.Schema.
+ :return:
+ """
+ self._name_type_mapping = OrderedDict()
+ attr_types = arrow_schema_to_attr_types(arrow_schema)
+ # Preserve field order from arrow_schema
+ for attr_name in arrow_schema.names:
+ self.add(attr_name, attr_types[attr_name])
+
+ def as_arrow_schema(self) -> pa.Schema:
+ """
+ Creates a new pyarrow.Schema according to the current Schema.
+ :return: pyarrow.Schema
+ """
+ return attr_types_to_arrow_schema(self._name_type_mapping)
+
+ def get_attr_names(self) -> List[str]:
+ """
+ Get all the attributes' names.
+ :return: a list of attribute names.
+ """
+ return list(self._name_type_mapping.keys())
+
+ def get_attr_type(self, attr_name: str) -> AttributeType:
+ """
+ Get an attribute's type specified by an attribute name.
+ :param attr_name: the name of the target attribute.
+ :return: the AttributeType of the target attribute.
+ """
+ return self._name_type_mapping[attr_name]
+
+ def as_key_value_pairs(self) -> List[Tuple[str, AttributeType]]:
+ """
+ Creates all attributes information according to the current Schema.
+ :return: A list of (name, type) tuples.
+ """
+ return [(k, v) for k, v in self._name_type_mapping.items()]
+
+ def get_partial_schema(self, attribute_names: List[str]) -> "Schema":
+ """
+ Creates a partial Schema with fields specified by the attribute names.
+
+ :param attribute_names: A list of attribute names for which to create the
+ partial schema.
+ :return: A new Schema instance containing only the specified fields, preserving
+ the order specified by the attribute names.
+ """
+ raw_schema = OrderedDict()
+ for name in attribute_names:
+ raw_schema[name] = RAW_TYPE_MAPPING.inverse[self.get_attr_type(name)]
+ return Schema(raw_schema=raw_schema)
+
+ def __eq__(self, other: object) -> bool:
+ if not isinstance(other, Schema):
+ return False
+ left_pairs = self.as_key_value_pairs()
+ right_pairs = other.as_key_value_pairs()
+ return left_pairs == right_pairs
+
+ def __str__(self) -> str:
+ content = ",\n".join(
+ f"({index}){repr(attr_name)} -> {attr_type}"
+ for index, (attr_name, attr_type) in enumerate(self.as_key_value_pairs(), 0)
+ )
+ return f"Schema[\n{content}\n]"
diff --git a/amber/src/main/python/core/models/schema/test_schema.py b/amber/src/main/python/core/models/schema/test_schema.py
new file mode 100644
index 00000000000..60e4c848a54
--- /dev/null
+++ b/amber/src/main/python/core/models/schema/test_schema.py
@@ -0,0 +1,156 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+import pytest
+
+from core.models.schema.attribute_type import AttributeType
+from core.models.schema.schema import Schema
+
+
+class TestSchema:
+ @pytest.fixture
+ def raw_schema(self):
+ return {
+ "field-1": "STRING",
+ "field-2": "INTEGER",
+ "field-3": "LONG",
+ "field-4": "DOUBLE",
+ "field-5": "BOOLEAN",
+ "field-6": "TIMESTAMP",
+ "field-7": "BINARY",
+ }
+
+ @pytest.fixture
+ def arrow_schema(self):
+ return pa.schema(
+ [
+ pa.field("field-1", pa.string()),
+ pa.field("field-2", pa.int32()),
+ pa.field("field-3", pa.int64()),
+ pa.field("field-4", pa.float64()),
+ pa.field("field-5", pa.bool_()),
+ pa.field("field-6", pa.timestamp("us")),
+ pa.field("field-7", pa.binary()),
+ ]
+ )
+
+ @pytest.fixture
+ def schema(self):
+ s = Schema()
+ s.add("field-1", AttributeType.STRING)
+ s.add("field-2", AttributeType.INT)
+ s.add("field-3", AttributeType.LONG)
+ s.add("field-4", AttributeType.DOUBLE)
+ s.add("field-5", AttributeType.BOOL)
+ s.add("field-6", AttributeType.TIMESTAMP)
+ s.add("field-7", AttributeType.BINARY)
+ return s
+
+ def test_accessors_and_mutators(self, schema):
+ assert schema.get_attr_names() == [f"field-{i}" for i in range(1, 8)]
+ assert schema.get_attr_type("field-2") == AttributeType.INT
+ assert schema.get_attr_type("field-6") == AttributeType.TIMESTAMP
+ assert schema.as_key_value_pairs() == [
+ ("field-1", AttributeType.STRING),
+ ("field-2", AttributeType.INT),
+ ("field-3", AttributeType.LONG),
+ ("field-4", AttributeType.DOUBLE),
+ ("field-5", AttributeType.BOOL),
+ ("field-6", AttributeType.TIMESTAMP),
+ ("field-7", AttributeType.BINARY),
+ ]
+ with pytest.raises(KeyError):
+ schema.get_attr_type("does not exist")
+ with pytest.raises(TypeError):
+ schema["illegal_assign"] = "value"
+ with pytest.raises(TypeError):
+ _ = schema["illegal_access"]
+ with pytest.raises(KeyError):
+ schema.add("field-2", AttributeType.LONG)
+
+ def test_convert_from_raw_schema(self, raw_schema, schema):
+ assert schema == Schema(raw_schema=raw_schema)
+
+ def test_convert_from_arrow_schema(self, arrow_schema, schema):
+ assert schema == Schema(arrow_schema=arrow_schema)
+ assert schema.as_arrow_schema() == arrow_schema
+
+ def test_large_binary_in_raw_schema(self):
+ """Test creating schema with LARGE_BINARY from raw schema."""
+ raw_schema = {
+ "regular_field": "STRING",
+ "large_binary_field": "LARGE_BINARY",
+ }
+ schema = Schema(raw_schema=raw_schema)
+ assert schema.get_attr_type("regular_field") == AttributeType.STRING
+ assert schema.get_attr_type("large_binary_field") == AttributeType.LARGE_BINARY
+
+ def test_large_binary_in_arrow_schema_with_metadata(self):
+ """Test creating schema with LARGE_BINARY from Arrow schema with metadata."""
+ arrow_schema = pa.schema(
+ [
+ pa.field("regular_field", pa.string()),
+ pa.field(
+ "large_binary_field",
+ pa.string(),
+ metadata={b"texera_type": b"LARGE_BINARY"},
+ ),
+ ]
+ )
+ schema = Schema(arrow_schema=arrow_schema)
+ assert schema.get_attr_type("regular_field") == AttributeType.STRING
+ assert schema.get_attr_type("large_binary_field") == AttributeType.LARGE_BINARY
+
+ def test_large_binary_as_arrow_schema_includes_metadata(self):
+ """Test that LARGE_BINARY fields include metadata in Arrow schema."""
+ schema = Schema()
+ schema.add("regular_field", AttributeType.STRING)
+ schema.add("large_binary_field", AttributeType.LARGE_BINARY)
+
+ arrow_schema = schema.as_arrow_schema()
+
+ # Regular field should have no metadata
+ regular_field = arrow_schema.field("regular_field")
+ assert (
+ regular_field.metadata is None
+ or b"texera_type" not in regular_field.metadata
+ )
+
+ # LARGE_BINARY field should have metadata
+ large_binary_field = arrow_schema.field("large_binary_field")
+ assert large_binary_field.metadata is not None
+ assert large_binary_field.metadata.get(b"texera_type") == b"LARGE_BINARY"
+ assert (
+ large_binary_field.type == pa.string()
+ ) # LARGE_BINARY is stored as string
+
+ def test_round_trip_large_binary_schema(self):
+ """Test round-trip conversion of schema with LARGE_BINARY."""
+ original_schema = Schema()
+ original_schema.add("field1", AttributeType.STRING)
+ original_schema.add("field2", AttributeType.LARGE_BINARY)
+ original_schema.add("field3", AttributeType.INT)
+
+ # Convert to Arrow and back
+ arrow_schema = original_schema.as_arrow_schema()
+ round_trip_schema = Schema(arrow_schema=arrow_schema)
+
+ assert round_trip_schema == original_schema
+ assert round_trip_schema.get_attr_type("field1") == AttributeType.STRING
+ assert round_trip_schema.get_attr_type("field2") == AttributeType.LARGE_BINARY
+ assert round_trip_schema.get_attr_type("field3") == AttributeType.INT
diff --git a/amber/src/main/python/core/models/single_blocking_io.py b/amber/src/main/python/core/models/single_blocking_io.py
new file mode 100644
index 00000000000..3dd1c761ff0
--- /dev/null
+++ b/amber/src/main/python/core/models/single_blocking_io.py
@@ -0,0 +1,133 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from threading import Condition
+from types import TracebackType
+from typing import IO, Type, AnyStr, Iterator, Iterable, Optional
+
+
+class SingleBlockingIO(IO):
+ """
+ An implementation of single-element IO that can be blocked on reading.
+
+ Some highlights:
+ - The IO only has one value.
+ - Each write() will append to the value.
+ - Each flush() will make the value readable resets the value to be written next.
+ - Each readline() will fetch the value and clear the IO.
+ - When there is no value to read, it blocks in readline() until there is a value.
+ """
+
+ def __init__(self, condition: Condition):
+ self.value: Optional[str] = None
+ self.buf: str = ""
+ self.condition: Condition = condition
+
+ def write(self, s: str) -> None:
+ """
+ Writes a partial string, append to the buffer.
+ :param s: a string.
+ :return:
+ """
+ self.buf += s
+
+ def flush(self) -> None:
+ """
+ Denotes the end of buffer, adds a "\n" to complete the string.
+ Flushes the completed string in the buffer to value.
+ Resets the buffer to accept the next complete string.
+ :return:
+ """
+ self.write("\n")
+ self.value, self.buf = self.buf, ""
+
+ def readline(self, limit=None) -> str:
+ """
+ Fetches a string value by removing it from the IO. It blocks the current
+ thread until there is a valid string to fetch.
+ :param limit: parent's API, not implemented here. It is always None.
+ :return str: A completed string value.
+ """
+ try:
+ with self.condition:
+ # keeps waiting until a value is available
+ while self.value is None:
+ self.condition.notify()
+ self.condition.wait()
+
+ # noinspection PyTypeChecker
+ return self.value
+ finally:
+ self.value = None
+
+ ####################################################################################
+ # The following IO methods are not implemented as they are not used in pdb.
+ ####################################################################################
+ def close(self) -> None:
+ pass
+
+ def fileno(self) -> int:
+ pass
+
+ def isatty(self) -> bool:
+ pass
+
+ def read(self, __n: int = ...) -> AnyStr:
+ pass
+
+ def readable(self) -> bool:
+ pass
+
+ def readlines(self, __hint: int = ...) -> list[AnyStr]:
+ pass
+
+ def seek(self, __offset: int, __whence: int = ...) -> int:
+ pass
+
+ def seekable(self) -> bool:
+ pass
+
+ def tell(self) -> int:
+ pass
+
+ def truncate(self, __size: int | None = ...) -> int:
+ pass
+
+ def writable(self) -> bool:
+ pass
+
+ def writelines(self, __lines: Iterable[AnyStr]) -> None:
+ pass
+
+ def __next__(self) -> AnyStr:
+ pass
+
+ def __iter__(self) -> Iterator[AnyStr]:
+ pass
+
+ def __enter__(self) -> IO[AnyStr]:
+ pass
+
+ def __exit__(
+ self,
+ __t: Type[BaseException] | None,
+ __value: BaseException | None,
+ __traceback: TracebackType | None,
+ ) -> bool | None:
+ pass
diff --git a/amber/src/main/python/core/models/state.py b/amber/src/main/python/core/models/state.py
new file mode 100644
index 00000000000..003aaa212ac
--- /dev/null
+++ b/amber/src/main/python/core/models/state.py
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import base64
+import json
+from typing import Any
+
+from .schema import Schema
+from .tuple import Tuple
+
+
+class State(dict):
+ CONTENT = "content"
+ SCHEMA = Schema(raw_schema={CONTENT: "STRING"})
+
+ def to_json(self) -> str:
+ return json.dumps(_to_json_value(self), separators=(",", ":"))
+
+ def to_tuple(self) -> Tuple:
+ return Tuple({State.CONTENT: self.to_json()}, schema=State.SCHEMA)
+
+ @classmethod
+ def from_json(cls, payload: str) -> "State":
+ return cls(_from_json_value(json.loads(payload)))
+
+ @classmethod
+ def from_tuple(cls, row: Tuple) -> "State":
+ return cls.from_json(row[cls.CONTENT])
+
+
+_TYPE_MARKER = "__texera_type__"
+_PAYLOAD_MARKER = "payload"
+_BYTES_TYPE = "bytes"
+
+
+def _to_json_value(value: Any) -> Any:
+ if value is None or isinstance(value, (bool, int, float, str)):
+ return value
+ if isinstance(value, bytes):
+ return {
+ _TYPE_MARKER: _BYTES_TYPE,
+ _PAYLOAD_MARKER: base64.b64encode(value).decode("ascii"),
+ }
+ if isinstance(value, dict):
+ return {str(key): _to_json_value(inner) for key, inner in value.items()}
+ if isinstance(value, (list, tuple)):
+ return [_to_json_value(inner) for inner in value]
+ raise TypeError(
+ f"State value of type {type(value).__name__} is not JSON serializable"
+ )
+
+
+def _from_json_value(value: Any) -> Any:
+ if isinstance(value, list):
+ return [_from_json_value(inner) for inner in value]
+ if isinstance(value, dict):
+ if value.get(_TYPE_MARKER) == _BYTES_TYPE:
+ return base64.b64decode(value[_PAYLOAD_MARKER])
+ return {key: _from_json_value(inner) for key, inner in value.items()}
+ return value
diff --git a/amber/src/main/python/core/models/table.py b/amber/src/main/python/core/models/table.py
new file mode 100644
index 00000000000..4716e0eba06
--- /dev/null
+++ b/amber/src/main/python/core/models/table.py
@@ -0,0 +1,105 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pandas
+from pampy import match
+from typing import Iterator, TypeVar, List
+
+from core.models import Tuple, TupleLike
+
+TableLike = TypeVar("TableLike", pandas.DataFrame, List[TupleLike])
+
+
+class Table(pandas.DataFrame):
+ @staticmethod
+ def from_table(table):
+ return table
+
+ @staticmethod
+ def from_data_frame(df):
+ return df
+
+ @staticmethod
+ def from_tuple_likes(tuple_likes: Iterator[TupleLike]):
+ # TODO: currently only validate all Tuples have the same fields.
+ # should validate types as well
+ column_names = None
+ records = []
+ for tuple_like in tuple_likes:
+ tuple_ = Tuple(tuple_like)
+ field_names = tuple_.get_field_names()
+
+ if column_names is not None:
+ assert field_names == column_names
+ else:
+ column_names = field_names
+
+ records.append(tuple_.get_fields())
+
+ return pandas.DataFrame.from_records(records, columns=column_names)
+
+ def __init__(self, table_like):
+ df: pandas.DataFrame
+
+ if isinstance(table_like, Table):
+ df = self.from_table(table_like)
+ elif isinstance(table_like, pandas.DataFrame):
+ df = self.from_data_frame(table_like)
+ elif isinstance(table_like, list):
+ # only supports List[TupleLike]
+ df = self.from_tuple_likes(table_like)
+ else:
+ raise TypeError(f"unsupported tablelike type {type(table_like)}")
+ super().__init__(df)
+
+ def as_tuples(self) -> Iterator[Tuple]:
+ """
+ Convert rows of the table into Tuples, and returning an iterator of Tuples
+ following their row index order.
+ :return:
+ """
+ for raw_tuple in self.itertuples(index=False, name=None):
+ yield Tuple(dict(zip(self.columns, raw_tuple)))
+
+ def __eq__(self, other: "Table") -> bool:
+ if isinstance(other, Table):
+ return all(a == b for a, b in zip(self.as_tuples(), other.as_tuples()))
+ else:
+ return super().__eq__(other).all()
+
+
+def all_output_to_tuple(output) -> Iterator[Tuple]:
+ """
+ Convert all kinds of types into Tuples.
+ :param output:
+ :return:
+ """
+ yield from match(
+ output,
+ None,
+ iter([None]),
+ Table,
+ lambda x: x.as_tuples(),
+ pandas.DataFrame,
+ lambda x: Table(x).as_tuples(),
+ List[TupleLike],
+ lambda x: (Tuple(t) for t in x),
+ TupleLike,
+ lambda x: iter([Tuple(x)]),
+ Tuple,
+ lambda x: iter([x]),
+ )
diff --git a/amber/src/main/python/core/models/test_state.py b/amber/src/main/python/core/models/test_state.py
new file mode 100644
index 00000000000..aef2297130b
--- /dev/null
+++ b/amber/src/main/python/core/models/test_state.py
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from core.models.state import State
+
+
+class TestState:
+ def test_state_subclasses_dict(self):
+ state = State({"a": 1})
+ assert isinstance(state, dict)
+ assert state["a"] == 1
+ assert State() == {}
+
+ def test_class_attributes(self):
+ assert State.CONTENT == "content"
+ assert State.SCHEMA.get_attr_names() == ["content"]
+
+ def test_json_round_trip_primitives(self):
+ original = State(
+ {
+ "string": "hello",
+ "int": 42,
+ "float": 3.14,
+ "bool_true": True,
+ "bool_false": False,
+ "none_value": None,
+ }
+ )
+ decoded = State.from_json(original.to_json())
+ assert decoded == original
+
+ def test_json_round_trip_empty(self):
+ assert State.from_json(State().to_json()) == State()
+
+ def test_json_round_trip_bytes(self):
+ original = State({"payload": b"\x00\x01\x02\xff"})
+ decoded = State.from_json(original.to_json())
+ assert decoded["payload"] == b"\x00\x01\x02\xff"
+ assert isinstance(decoded["payload"], bytes)
+
+ def test_json_round_trip_nested_dict(self):
+ original = State({"outer": {"inner": {"value": 1}}})
+ decoded = State.from_json(original.to_json())
+ assert decoded == original
+
+ def test_json_round_trip_list_of_mixed_values(self):
+ original = State({"items": [1, "two", 3.0, True, None]})
+ decoded = State.from_json(original.to_json())
+ assert decoded == original
+
+ def test_json_round_trip_bytes_inside_list_and_nested_dict(self):
+ original = State(
+ {
+ "blobs": [b"first", b"second"],
+ "nested": {"sub_blob": b"inside"},
+ }
+ )
+ decoded = State.from_json(original.to_json())
+ assert decoded["blobs"] == [b"first", b"second"]
+ assert decoded["nested"]["sub_blob"] == b"inside"
+
+ def test_to_json_rejects_non_serializable_value(self):
+ class Custom:
+ pass
+
+ with pytest.raises(TypeError):
+ State({"bad": Custom()}).to_json()
+
+ def test_tuple_round_trip(self):
+ original = State({"loop_counter": 3, "label": "outer", "blob": b"\x01\x02"})
+ decoded = State.from_tuple(original.to_tuple())
+ assert decoded == original
+
+ def test_to_tuple_uses_state_schema(self):
+ tuple_ = State({"x": 1}).to_tuple()
+ # Single STRING column whose value is the JSON serialization.
+ assert tuple_[State.CONTENT] == '{"x":1}'
+
+ def test_nested_dict_decodes_to_plain_dict(self):
+ # Top-level returns a State; nested dicts come back as plain dict.
+ # This is intentional -- only the outermost mapping is wrapped.
+ decoded = State.from_json('{"outer":{"inner":1}}')
+ assert isinstance(decoded, State)
+ assert isinstance(decoded["outer"], dict)
+ assert not isinstance(decoded["outer"], State)
diff --git a/amber/src/main/python/core/models/test_table.py b/amber/src/main/python/core/models/test_table.py
new file mode 100644
index 00000000000..368220779d5
--- /dev/null
+++ b/amber/src/main/python/core/models/test_table.py
@@ -0,0 +1,144 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import pandas
+import pickle
+import pytest
+from pandas import RangeIndex
+
+from core.models import Table, Tuple
+
+
+class TestTable:
+ @pytest.fixture
+ def a_timestamp(self):
+ return datetime.datetime.now()
+
+ @pytest.fixture
+ def target_raw_tuples(self, a_timestamp):
+ return [
+ {
+ "field1": 1,
+ "field2": "hello",
+ "field3": 2.3,
+ "field4": True,
+ "field5": a_timestamp,
+ "field6": b"some binary",
+ "7_special-name": None,
+ "none": None,
+ },
+ {
+ "field1": 2,
+ "field2": "world",
+ "field3": 0.0,
+ "field4": False,
+ "field5": datetime.datetime.fromtimestamp(1000000000),
+ "field6": pickle.dumps([1, 2, 3]),
+ "7_special-name": "a strange value",
+ "none": None,
+ },
+ ]
+
+ @pytest.fixture
+ def target_tuples(self, target_raw_tuples):
+ return [Tuple(raw_tuple) for raw_tuple in target_raw_tuples]
+
+ @pytest.fixture
+ def target_table(self, target_raw_tuples):
+ return Table(target_raw_tuples)
+
+ @pytest.fixture
+ def target_data_frame(self, a_timestamp):
+ return pandas.DataFrame(
+ {
+ "field1": [1, 2],
+ "field2": ["hello", "world"],
+ "field3": [2.3, 0.0],
+ "field4": [True, False],
+ "field5": [
+ a_timestamp,
+ datetime.datetime.fromtimestamp(1000000000),
+ ],
+ "field6": [b"some binary", pickle.dumps([1, 2, 3])],
+ "7_special-name": [None, "a strange value"],
+ "none": [None, None],
+ },
+ columns=[
+ "field1",
+ "field2",
+ "field3",
+ "field4",
+ "field5",
+ "field6",
+ "7_special-name",
+ "none",
+ ],
+ )
+
+ def test_table_creation(self, target_table, a_timestamp):
+ assert target_table["field1"][0] == 1
+ assert target_table["field1"][1] == 2
+ assert target_table["field2"][0] == "hello"
+ assert target_table["field2"][1] == "world"
+ assert target_table["field3"][0] == 2.3
+ assert target_table["field3"][1] == 0.0
+ assert target_table["field4"][0]
+ assert not target_table["field4"][1]
+ assert target_table["field5"][0] == a_timestamp
+ assert target_table["field5"][1] == datetime.datetime.fromtimestamp(1000000000)
+ assert target_table["field6"][0] == b"some binary"
+ assert target_table["field6"][1] == pickle.dumps([1, 2, 3])
+ assert target_table["7_special-name"][0] is None
+ assert target_table["7_special-name"][1] == "a strange value"
+ assert target_table["none"][0] is None
+ assert target_table["none"][1] is None
+
+ def test_as_tuples_preserve_types(self, target_table, target_tuples):
+ assert list(target_table.as_tuples()) == target_tuples
+
+ def test_table_from_data_frame(self, target_table, target_data_frame):
+ assert Table(target_data_frame) == target_table
+
+ def test_table_from_list_of_tuples(self, target_table, target_tuples):
+ table = Table(target_tuples)
+ assert table == target_table
+ assert list(table.as_tuples()) == target_tuples
+
+ def test_table_from_list_of_series(
+ self, target_table, a_timestamp, target_raw_tuples, target_tuples
+ ):
+ table = Table([pandas.Series(raw_tuple) for raw_tuple in target_raw_tuples])
+
+ assert table == target_table
+ assert list(table.as_tuples()) == target_tuples
+
+ def test_table_from_table(self, target_table, target_tuples):
+ table = Table(target_table)
+ assert table == target_table
+ assert list(table.as_tuples()) == target_tuples
+
+ def test_use_table_as_data_frame(self, target_table, target_data_frame):
+ df = target_table
+ assert (df.index == RangeIndex(start=0, stop=2, step=1)).all()
+ concat_df = pandas.concat([df, df])
+ assert len(concat_df) == 4
+ assert target_table.equals(target_data_frame)
+
+ def test_validation_of_schema(self):
+ with pytest.raises(AssertionError):
+ Table([{"text": "hello"}, {"book": "harry"}])
diff --git a/amber/src/main/python/core/models/test_tuple.py b/amber/src/main/python/core/models/test_tuple.py
new file mode 100644
index 00000000000..efb4fdf5c71
--- /dev/null
+++ b/amber/src/main/python/core/models/test_tuple.py
@@ -0,0 +1,321 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import pandas
+import pyarrow
+import pytest
+import numpy as np
+from copy import deepcopy
+
+from core.models import Tuple, ArrowTableTupleProvider
+from core.models.schema.schema import Schema
+
+
+class TestTuple:
+ @pytest.fixture
+ def target_tuple(self):
+ return Tuple({"x": 1, "y": "a"})
+
+ def test_tuple_from_list(self, target_tuple):
+ assert Tuple([("x", 1), ("y", "a")]) == target_tuple
+
+ def test_tuple_from_dict(self, target_tuple):
+ assert Tuple({"x": 1, "y": "a"}) == target_tuple
+
+ def test_tuple_from_series(self, target_tuple):
+ assert Tuple(pandas.Series({"x": 1, "y": "a"})) == target_tuple
+
+ def test_tuple_as_key_value_pairs(self, target_tuple):
+ assert target_tuple.as_key_value_pairs() == [("x", 1), ("y", "a")]
+
+ def test_tuple_as_dict(self, target_tuple):
+ assert target_tuple.as_dict() == {"x": 1, "y": "a"}
+
+ def test_tuple_as_series(self, target_tuple):
+ assert (target_tuple.as_series() == pandas.Series({"x": 1, "y": "a"})).all()
+
+ def test_tuple_get_fields(self, target_tuple):
+ assert target_tuple.get_fields() == (1, "a")
+
+ def test_tuple_get_field_names(self, target_tuple):
+ assert target_tuple.get_field_names() == ("x", "y")
+
+ def test_tuple_get_item(self, target_tuple):
+ assert target_tuple["x"] == 1
+ assert target_tuple["y"] == "a"
+ assert target_tuple[0] == 1
+ assert target_tuple[1] == "a"
+
+ def test_tuple_set_item(self, target_tuple):
+ target_tuple["x"] = 3
+ assert target_tuple["x"] == 3
+ assert target_tuple["y"] == "a"
+ assert target_tuple[0] == 3
+ assert target_tuple[1] == "a"
+ target_tuple["z"] = 1.1
+ assert target_tuple[2] == 1.1
+ assert target_tuple["z"] == 1.1
+
+ def test_tuple_str(self, target_tuple):
+ assert str(target_tuple) == "Tuple['x': 1, 'y': 'a']"
+
+ def test_tuple_repr(self, target_tuple):
+ assert repr(target_tuple) == "Tuple['x': 1, 'y': 'a']"
+
+ def test_tuple_eq(self, target_tuple):
+ assert target_tuple == target_tuple
+ assert not Tuple({"x": 2, "y": "a"}) == target_tuple
+
+ def test_tuple_ne(self, target_tuple):
+ assert not target_tuple != target_tuple
+ assert Tuple({"x": 1, "y": "b"}) != target_tuple
+
+ def test_reject_empty_tuplelike(self):
+ with pytest.raises(AssertionError):
+ Tuple([])
+ with pytest.raises(AssertionError):
+ Tuple({})
+ with pytest.raises(AssertionError):
+ Tuple(pandas.Series(dtype=pandas.StringDtype()))
+
+ def test_reject_invalid_tuplelike(self):
+ with pytest.raises(TypeError):
+ Tuple(1)
+ with pytest.raises(TypeError):
+ Tuple([1])
+ with pytest.raises(TypeError):
+ Tuple([None])
+
+ def test_tuple_lazy_get_from_arrow(self):
+ def field_accessor(field_name):
+ return chr(96 + int(field_name))
+
+ chr_tuple = Tuple({"1": "a", "3": "c"})
+ tuple_ = Tuple({"1": field_accessor, "3": field_accessor})
+ assert tuple_ == Tuple({"1": "a", "3": "c"})
+ tuple_ = Tuple({"1": field_accessor, "3": field_accessor})
+ assert deepcopy(tuple_) == chr_tuple
+
+ def test_retrieve_tuple_from_empty_arrow_table(self):
+ arrow_schema = pyarrow.schema([])
+ arrow_table = arrow_schema.empty_table()
+ tuple_provider = ArrowTableTupleProvider(arrow_table)
+ tuples = [
+ Tuple({name: field_accessor for name in arrow_table.column_names})
+ for field_accessor in tuple_provider
+ ]
+ assert tuples == []
+
+ def test_finalize_tuple(self):
+ tuple_ = Tuple(
+ {"name": "texera", "age": 21, "scores": [85, 94, 100], "height": np.nan}
+ )
+ schema = Schema(
+ raw_schema={
+ "name": "STRING",
+ "age": "INTEGER",
+ "scores": "BINARY",
+ "height": "DOUBLE",
+ }
+ )
+ tuple_.finalize(schema)
+ assert isinstance(tuple_["scores"], bytes)
+ assert tuple_["height"] is None
+
+ def test_hash(self):
+ schema = Schema(
+ raw_schema={
+ "col-int": "INTEGER",
+ "col-string": "STRING",
+ "col-bool": "BOOLEAN",
+ "col-long": "LONG",
+ "col-double": "DOUBLE",
+ "col-timestamp": "TIMESTAMP",
+ "col-binary": "BINARY",
+ }
+ )
+
+ tuple_ = Tuple(
+ {
+ "col-int": 922323,
+ "col-string": "string-attr",
+ "col-bool": True,
+ "col-long": 1123213213213,
+ "col-double": 214214.9969346,
+ "col-timestamp": datetime.datetime.fromtimestamp(100000000),
+ "col-binary": b"hello",
+ },
+ schema,
+ )
+ assert hash(tuple_) == -1335416166 # calculated with Java
+
+ tuple2 = Tuple(
+ {
+ "col-int": 0,
+ "col-string": "",
+ "col-bool": False,
+ "col-long": 0,
+ "col-double": 0.0,
+ "col-timestamp": datetime.datetime.fromtimestamp(0),
+ "col-binary": b"",
+ },
+ schema,
+ )
+
+ assert hash(tuple2) == -1409761483 # calculated with Java
+
+ tuple3 = Tuple(
+ {
+ "col-int": None,
+ "col-string": None,
+ "col-bool": None,
+ "col-long": None,
+ "col-double": None,
+ "col-timestamp": None,
+ "col-binary": None,
+ },
+ schema,
+ )
+
+ assert hash(tuple3) == 1742810335 # calculated with Java
+
+ tuple4 = Tuple(
+ {
+ "col-int": -3245763,
+ "col-string": "\n\r\napple",
+ "col-bool": True,
+ "col-long": -8965536434247,
+ "col-double": 1 / 3,
+ "col-timestamp": datetime.datetime.fromtimestamp(-1990),
+ "col-binary": None,
+ },
+ schema,
+ )
+ assert hash(tuple4) == -592643630 # calculated with Java
+
+ tuple5 = Tuple(
+ {
+ "col-int": 0x7FFFFFFF,
+ "col-string": "",
+ "col-bool": True,
+ "col-long": 0x7FFFFFFFFFFFFFFF,
+ "col-double": 7 / 17,
+ "col-timestamp": datetime.datetime.fromtimestamp(1234567890),
+ "col-binary": b"o" * 4097,
+ },
+ schema,
+ )
+ assert hash(tuple5) == -2099556631 # calculated with Java
+
+ def test_tuple_with_large_binary(self):
+ """Test tuple with largebinary field."""
+ from core.models.type.large_binary import largebinary
+
+ schema = Schema(
+ raw_schema={
+ "regular_field": "STRING",
+ "large_binary_field": "LARGE_BINARY",
+ }
+ )
+
+ large_binary = largebinary("s3://test-bucket/path/to/object")
+ tuple_ = Tuple(
+ {
+ "regular_field": "test string",
+ "large_binary_field": large_binary,
+ },
+ schema=schema,
+ )
+
+ assert tuple_["regular_field"] == "test string"
+ assert tuple_["large_binary_field"] == large_binary
+ assert isinstance(tuple_["large_binary_field"], largebinary)
+ assert tuple_["large_binary_field"].uri == "s3://test-bucket/path/to/object"
+
+ def test_tuple_from_arrow_with_large_binary(self):
+ """Test creating tuple from Arrow table with LARGE_BINARY metadata."""
+ import pyarrow as pa
+ from core.models.type.large_binary import largebinary
+
+ # Create Arrow schema with LARGE_BINARY metadata
+ arrow_schema = pa.schema(
+ [
+ pa.field("regular_field", pa.string()),
+ pa.field(
+ "large_binary_field",
+ pa.string(),
+ metadata={b"texera_type": b"LARGE_BINARY"},
+ ),
+ ]
+ )
+
+ # Create Arrow table with URI string for large_binary_field
+ arrow_table = pa.Table.from_pydict(
+ {
+ "regular_field": ["test"],
+ "large_binary_field": ["s3://test-bucket/path/to/object"],
+ },
+ schema=arrow_schema,
+ )
+
+ # Create tuple from Arrow table
+ tuple_provider = ArrowTableTupleProvider(arrow_table)
+ tuples = [
+ Tuple({name: field_accessor for name in arrow_table.column_names})
+ for field_accessor in tuple_provider
+ ]
+
+ assert len(tuples) == 1
+ tuple_ = tuples[0]
+ assert tuple_["regular_field"] == "test"
+ assert isinstance(tuple_["large_binary_field"], largebinary)
+ assert tuple_["large_binary_field"].uri == "s3://test-bucket/path/to/object"
+
+ def test_tuple_with_null_large_binary(self):
+ """Test tuple with null largebinary field."""
+ import pyarrow as pa
+
+ # Create Arrow schema with LARGE_BINARY metadata
+ arrow_schema = pa.schema(
+ [
+ pa.field(
+ "large_binary_field",
+ pa.string(),
+ metadata={b"texera_type": b"LARGE_BINARY"},
+ ),
+ ]
+ )
+
+ # Create Arrow table with null value
+ arrow_table = pa.Table.from_pydict(
+ {
+ "large_binary_field": [None],
+ },
+ schema=arrow_schema,
+ )
+
+ # Create tuple from Arrow table
+ tuple_provider = ArrowTableTupleProvider(arrow_table)
+ tuples = [
+ Tuple({name: field_accessor for name in arrow_table.column_names})
+ for field_accessor in tuple_provider
+ ]
+
+ assert len(tuples) == 1
+ tuple_ = tuples[0]
+ assert tuple_["large_binary_field"] is None
diff --git a/amber/src/main/python/core/models/tuple.py b/amber/src/main/python/core/models/tuple.py
new file mode 100644
index 00000000000..916301406f5
--- /dev/null
+++ b/amber/src/main/python/core/models/tuple.py
@@ -0,0 +1,446 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import ctypes
+import pandas
+import pickle
+import pyarrow
+import struct
+import typing
+from collections import OrderedDict
+from copy import deepcopy
+from loguru import logger
+from pandas._libs.missing import checknull
+from pympler import asizeof
+from typing import Any, List, Iterator, Callable
+from typing_extensions import Protocol, runtime_checkable
+
+from core.models.type.large_binary import largebinary
+from .schema.attribute_type import TO_PYOBJECT_MAPPING, AttributeType
+from .schema.field import Field
+from .schema.schema import Schema
+
+
+@runtime_checkable
+class TupleLike(Protocol):
+ def __getitem__(self, item: typing.Union[str, int]) -> Field: ...
+
+ def __setitem__(self, key: typing.Union[str, int], value: Field) -> None: ...
+
+
+class ArrowTableTupleProvider:
+ """
+ This class provides "view"s for tuple from a pyarrow.Table.
+ """
+
+ def __init__(self, table: pyarrow.Table):
+ """
+ Construct a provider from a pyarrow.Table.
+ Keep the current chunk and tuple idx as its state.
+ """
+ self._table = table
+ self._current_idx = 0
+ self._current_chunk = 0
+
+ def __iter__(self) -> Iterator[Callable]:
+ """
+ Return itself as it is iterable.
+ """
+ return self
+
+ def __next__(self) -> Callable:
+ """
+ Provide the field accessor of the next tuple.
+ If current chunk is exhausted, move to the first tuple of the next chunk.
+ """
+ if self._table.num_columns == 0:
+ # empty table
+ raise StopIteration
+ if self._current_idx >= len(self._table.column(0).chunks[self._current_chunk]):
+ self._current_idx = 0
+ self._current_chunk += 1
+ if self._current_chunk >= self._table.column(0).num_chunks:
+ raise StopIteration
+
+ chunk_idx = self._current_chunk
+ tuple_idx = self._current_idx
+
+ def field_accessor(field_name: str) -> Field:
+ """
+ Retrieve the field value by a given field name.
+ This abstracts and hides the underlying implementation of the tuple data
+ storage from the user.
+ """
+ value = self._table.column(field_name).chunks[chunk_idx][tuple_idx].as_py()
+ field_type = self._table.schema.field(field_name).type
+ field_metadata = self._table.schema.field(field_name).metadata
+
+ # for binary types, convert pickled objects back.
+ if (
+ field_type == pyarrow.binary()
+ and value is not None
+ and value[:6] == b"pickle"
+ ):
+ value = pickle.loads(value[10:])
+
+ # Convert URI string to largebinary for LARGE_BINARY types
+ # Metadata is set by Scala ArrowUtils or Python iceberg_utils
+ elif (
+ value is not None
+ and field_metadata
+ and field_metadata.get(b"texera_type") == b"LARGE_BINARY"
+ ):
+ value = largebinary(value)
+
+ return value
+
+ self._current_idx += 1
+ return field_accessor
+
+
+def double_to_long(value: float) -> int:
+ """
+ Convert a double value into a long value.
+ :param value: A double (Python float) value.
+ :return: The converted long (Python int) value.
+ """
+ # Pack the double value into a binary string of 8 bytes
+ packed_value = struct.pack("d", value)
+ # Unpack the binary string to a 64-bit integer (int in Python 3)
+ long_value = struct.unpack("Q", packed_value)[0]
+ return long_value
+
+
+def int_32(value: int) -> int:
+ """
+ Convert a Python int (unbounded) to a 32-bit int with overflow.
+ :param value: A Python int value.
+ :return: The converted 32-bit integer, with overflow.
+ """
+ return ctypes.c_int32(value).value
+
+
+def java_hash_bool(value: bool) -> int:
+ """
+ Java's hash function for a boolean value.
+ :param value: A boolean value.
+ :return: Java's hash value in a 32-bit integer.
+ """
+ return 1231 if value else 1237
+
+
+def java_hash_long(value: int) -> int:
+ """
+ Java's hash function for a long value.
+ :param value: A long (Python int) value.
+ :return: Java's hash value in a 32-bit integer.
+ """
+ return int_32(value ^ (value >> 32))
+
+
+def java_hash_bytes(bytes: Iterator[int], init: int, salt: int):
+ """
+ Java's hash function for an array of bytes.
+ :param bytes: An iterator of int (byte) values.
+ :param init: An init hash value.
+ :param salt: A hash salt value.
+ :return: Java's hash value in a 32-bit integer.
+ """
+ h = init
+ for b in bytes:
+ h = int_32(salt * h + b)
+ return h
+
+
+class Tuple:
+ """
+ Lazy-Tuple implementation.
+ """
+
+ def __init__(
+ self,
+ tuple_like: typing.Optional["TupleLike"] = None,
+ schema: typing.Optional[Schema] = None,
+ ):
+ """
+ Construct a lazy-tuple with given TupleLike object. If the field value is a
+ accessor callable, the actual value is fetched upon first reference.
+
+ :param tuple_like: in which the field value could be the actual value in
+ memory, or a callable accessor.
+ """
+ assert len(tuple_like) != 0
+ self._field_data: "OrderedDict[str, Field]"
+ if isinstance(tuple_like, Tuple):
+ self._field_data = tuple_like._field_data
+ elif isinstance(tuple_like, pandas.Series):
+ self._field_data = OrderedDict(tuple_like.to_dict())
+ else:
+ self._field_data = OrderedDict(tuple_like) if tuple_like else OrderedDict()
+ self._schema: typing.Optional[Schema] = schema
+
+ def __getitem__(self, item: typing.Union[int, str]) -> Field:
+ """
+ Get a field value with given item. If the value is an accessor, fetch it from
+ the accessor.
+
+ :param item: field name or field index
+ :return: field value
+ """
+ assert isinstance(item, (int, str)), (
+ "field can only be retrieved by index or name"
+ )
+
+ if isinstance(item, int):
+ item: str = self.get_field_names()[item]
+
+ if (
+ callable(self._field_data[item])
+ and getattr(self._field_data[item], "__name__", "Unknown")
+ == "field_accessor"
+ ):
+ # evaluate the field now
+ field_accessor = self._field_data[item]
+ self._field_data[item] = field_accessor(field_name=item)
+ return self._field_data[item]
+
+ def __setitem__(self, field_name: str, field_value: Field) -> None:
+ """
+ Set a field with the given value.
+ :param field_name
+ :param field_value
+ """
+ assert isinstance(field_name, str), "field can only be set by name"
+ assert not callable(field_value), "field cannot be of type callable"
+ self._field_data[field_name] = field_value
+
+ def as_series(self) -> pandas.Series:
+ """Convert the tuple to Pandas series format"""
+ return pandas.Series(self.as_dict())
+
+ def as_dict(self) -> "OrderedDict[str, Field]":
+ """
+ Return a dictionary copy of this tuple.
+ Fields will be fetched from accessor if absent.
+ :return: dict with all the fields
+ """
+ # evaluate all the fields now
+ for i in self.get_field_names():
+ self.__getitem__(i)
+ return deepcopy(self._field_data)
+
+ def as_key_value_pairs(self) -> List[typing.Tuple[str, Field]]:
+ return [(k, v) for k, v in self.as_dict().items()]
+
+ def get_serialized_field(self, field_name: str) -> Field:
+ """
+ Get a field value serialized for Arrow table conversion.
+ For LARGE_BINARY fields, converts largebinary instances to URI strings.
+ For other fields, returns the value as-is.
+
+ :param field_name: field name
+ :return: field value (URI string for LARGE_BINARY fields with largebinary values)
+ """
+ value = self[field_name]
+
+ # Convert largebinary to URI string for LARGE_BINARY fields when schema available
+ if (
+ self._schema is not None
+ and self._schema.get_attr_type(field_name) == AttributeType.LARGE_BINARY
+ and isinstance(value, largebinary)
+ ):
+ return value.uri
+
+ return value
+
+ def get_field_names(self) -> typing.Tuple[str]:
+ return tuple(map(str, self._field_data.keys()))
+
+ def get_fields(self, output_field_names=None) -> typing.Tuple[Field, ...]:
+ """
+ Get values from tuple for selected fields.
+ """
+ if output_field_names is None:
+ output_field_names = self.get_field_names()
+ return tuple(self[i] for i in output_field_names)
+
+ def finalize(self, schema: Schema) -> None:
+ """
+ Finalizes a Tuple by adding a schema to it. This convert all Fields into the
+ AttributeType defined in the Schema and make the Tuple immutable.
+
+ A Tuple can have no Schema initially. The types of Fields are not restricted.
+ This is to provide the maximum flexibility for users to construct Tuples as
+ they wish. When a Schema is added, the Tuple is finalized to match the Schema.
+
+ :param schema: target Schema to finalize the Tuple.
+ :return:
+ """
+ assert self._schema is None
+ self.cast_to_schema(schema)
+ self.validate_schema(schema)
+ self._schema = schema
+
+ def cast_to_schema(self, schema: Schema) -> None:
+ """
+ Safely cast each field value to match the target schema.
+ If failed, the value will stay not changed.
+ This current conducts two kinds of casts:
+ 1. cast NaN to None;
+ 2. cast any object to bytes (using pickle).
+ :param schema: The target Schema that describes the target AttributeType to
+ cast.
+ :return:
+ """
+ for field_name in self.get_field_names():
+ try:
+ field_value: Field = self[field_name]
+
+ # convert NaN to None to support null value conversion
+ if checknull(field_value):
+ self[field_name] = None
+
+ if field_value is not None:
+ field_type = schema.get_attr_type(field_name)
+ if field_type == AttributeType.BINARY and not isinstance(
+ field_value, bytes
+ ):
+ self[field_name] = b"pickle " + pickle.dumps(field_value)
+ except Exception as err:
+ # Surpass exceptions during cast.
+ # Keep the value as it is if the cast fails, and continue to attempt
+ # on the next one.
+ logger.warning(err)
+ continue
+
+ def validate_schema(self, schema: Schema) -> None:
+ """
+ Checks if the field values in the Tuple matches the expected Schema.
+ :param schema: Schema
+ :return:
+ """
+
+ schema_fields = schema.get_attr_names()
+ tuple_fields = self.get_field_names()
+ expected_but_missing = set(schema_fields) - set(tuple_fields)
+ unexpected = set(tuple_fields) - set(schema_fields)
+ if expected_but_missing:
+ raise KeyError(
+ f"field{'' if len(expected_but_missing) == 1 else 's'} "
+ f"{', '.join(map(repr, expected_but_missing))} "
+ f"{'is' if len(expected_but_missing) == 1 else 'are'} "
+ f"expected but missing in the {self}."
+ )
+
+ if unexpected:
+ raise KeyError(
+ f"{self} contains {'an' if len(unexpected) == 1 else ''} unexpected "
+ f"field{'' if len(unexpected) == 1 else 's'}: "
+ f"{', '.join(map(repr, unexpected))}."
+ )
+
+ for field_name, field_value in self.as_key_value_pairs():
+ expected = schema.get_attr_type(field_name)
+ if not isinstance(
+ field_value, (TO_PYOBJECT_MAPPING.get(expected), type(None))
+ ):
+ raise TypeError(
+ f"Unmatched type for field '{field_name}', expected {expected}, "
+ f"got {field_value} ({type(field_value)}) instead."
+ )
+
+ def get_partial_tuple(self, attribute_names: List[str]) -> "Tuple":
+ """
+ Creates a partial Tuple with fields specified by the attribute names.
+
+ :param attribute_names: A list of attribute names for which to create the
+ partial tuple.
+ :return: A new Tuple instance containing only the specified fields,
+ preserving the order specified by the attribute names.
+ """
+ assert self._schema is not None
+ schema = self._schema.get_partial_schema(attribute_names)
+ new_raw_tuple = OrderedDict()
+ for name in attribute_names:
+ new_raw_tuple[name] = self[name]
+ return Tuple(new_raw_tuple, schema=schema)
+
+ def __iter__(self) -> Iterator[Field]:
+ return iter(self.get_fields())
+
+ def __str__(self) -> str:
+ content = ", ".join(
+ [repr(key) + ": " + repr(value) for key, value in self.as_key_value_pairs()]
+ )
+ return f"Tuple[{content}]"
+
+ __repr__ = __str__
+
+ def __eq__(self, other: Any) -> bool:
+ return (
+ isinstance(other, Tuple)
+ and self.get_field_names() == other.get_field_names()
+ and all(self[i] == other[i] for i in self.get_field_names())
+ )
+
+ def __ne__(self, other) -> bool:
+ return not self.__eq__(other)
+
+ def __len__(self) -> int:
+ return len(self._field_data)
+
+ def __contains__(self, __x: object) -> bool:
+ return __x in self._field_data
+
+ def __hash__(self) -> int:
+ """
+ Aligned with Java's built-in hash algorithm implementation described in
+ _Josh Bloch's Effective Java_.
+ This algorithm is taken by
+ - Built-in Java (java.util.Objects.hash)
+ - Guava (com.google.common.base.Objects.hashCode)
+ :return: A 32-bit integer value.
+ """
+ result = 1
+ salt = 31 # for ease of optimization
+
+ mapping = {
+ AttributeType.BOOL: lambda f: java_hash_bool(f),
+ AttributeType.INT: lambda f: int_32(f),
+ AttributeType.LONG: lambda f: java_hash_long(f),
+ AttributeType.DOUBLE: lambda f: java_hash_long(double_to_long(f)),
+ AttributeType.STRING: lambda f: java_hash_bytes(map(ord, f), 0, salt),
+ AttributeType.TIMESTAMP: lambda f: java_hash_long(int(f.timestamp())),
+ AttributeType.BINARY: lambda f: java_hash_bytes(f, 1, salt),
+ }
+
+ for name, field in self.as_key_value_pairs():
+ attr_type = self._schema.get_attr_type(name)
+ if field is None:
+ hash_value = 0
+ else:
+ hash_value = mapping[attr_type](field)
+ result = result * salt + hash_value
+
+ return int_32(result)
+
+ def in_mem_size(self) -> int:
+ """
+ Calculate the in-memory size of the Tuple instance.
+ :return: The size in bytes.
+ """
+ return asizeof.asizeof(self)
diff --git a/amber/src/main/python/core/models/type/__init__.py b/amber/src/main/python/core/models/type/__init__.py
new file mode 100644
index 00000000000..41344433aab
--- /dev/null
+++ b/amber/src/main/python/core/models/type/__init__.py
@@ -0,0 +1,20 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .large_binary import largebinary
+
+__all__ = ["largebinary"]
diff --git a/amber/src/main/python/core/models/type/large_binary.py b/amber/src/main/python/core/models/type/large_binary.py
new file mode 100644
index 00000000000..581a688912b
--- /dev/null
+++ b/amber/src/main/python/core/models/type/large_binary.py
@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+largebinary represents a reference to a large object stored externally (e.g., S3).
+This is a schema type class used throughout the system for handling
+LARGE_BINARY attribute types.
+"""
+
+from typing import Optional
+from urllib.parse import urlparse
+
+
+class largebinary:
+ """
+ largebinary represents a reference to a large object stored in S3.
+
+ Each largebinary is identified by an S3 URI (s3://bucket/path/to/object).
+ largebinary objects are automatically tracked and cleaned up when the workflow
+ execution completes.
+
+ Usage:
+ from pytexera import largebinary, LargeBinaryInputStream, LargeBinaryOutputStream
+
+ # Create a new largebinary for writing
+ large_binary = largebinary()
+ with LargeBinaryOutputStream(large_binary) as out:
+ out.write(b"data")
+ # large_binary is now ready to be added to tuples
+
+ # Read from an existing largebinary
+ with LargeBinaryInputStream(large_binary) as stream:
+ content = stream.read()
+
+ # Create from existing URI (e.g., from deserialization)
+ large_binary = largebinary("s3://bucket/path/to/object")
+ """
+
+ def __init__(self, uri: Optional[str] = None):
+ """
+ Create a largebinary.
+
+ Args:
+ uri: Optional S3 URI in the format s3://bucket/path/to/object.
+ If None, creates a new largebinary with a unique S3 URI.
+
+ Raises:
+ ValueError: If URI is provided but doesn't start with "s3://"
+ """
+ if uri is None:
+ # Lazy import to avoid circular dependencies
+ from pytexera.storage import large_binary_manager
+
+ uri = large_binary_manager.create()
+
+ if not uri.startswith("s3://"):
+ raise ValueError(f"largebinary URI must start with 's3://', got: {uri}")
+
+ self._uri = uri
+
+ @property
+ def uri(self) -> str:
+ """Get the S3 URI of this largebinary."""
+ return self._uri
+
+ def get_bucket_name(self) -> str:
+ """Get the S3 bucket name from the URI."""
+ return urlparse(self._uri).netloc
+
+ def get_object_key(self) -> str:
+ """Get the S3 object key (path) from the URI, without leading slash."""
+ return urlparse(self._uri).path.lstrip("/")
+
+ def __str__(self) -> str:
+ return self._uri
+
+ def __repr__(self) -> str:
+ return f"largebinary('{self._uri}')"
+
+ def __eq__(self, other) -> bool:
+ return isinstance(other, largebinary) and self._uri == other._uri
+
+ def __hash__(self) -> int:
+ return hash(self._uri)
diff --git a/amber/src/main/python/core/models/type/test_large_binary.py b/amber/src/main/python/core/models/type/test_large_binary.py
new file mode 100644
index 00000000000..36310e1dd53
--- /dev/null
+++ b/amber/src/main/python/core/models/type/test_large_binary.py
@@ -0,0 +1,88 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+from unittest.mock import patch
+from core.models.type.large_binary import largebinary
+
+
+class TestLargeBinary:
+ def test_create_with_uri(self):
+ """Test creating largebinary with a valid S3 URI."""
+ uri = "s3://test-bucket/path/to/object"
+ large_binary = largebinary(uri)
+ assert large_binary.uri == uri
+ assert str(large_binary) == uri
+ assert repr(large_binary) == f"largebinary('{uri}')"
+
+ def test_create_without_uri(self):
+ """Test creating largebinary without URI (calls large_binary_manager.create)."""
+ with patch("pytexera.storage.large_binary_manager.create") as mock_create:
+ mock_create.return_value = "s3://bucket/objects/123/uuid"
+ large_binary = largebinary()
+ assert large_binary.uri == "s3://bucket/objects/123/uuid"
+ mock_create.assert_called_once()
+
+ def test_invalid_uri_raises_value_error(self):
+ """Test that invalid URI (not starting with s3://) raises ValueError."""
+ with pytest.raises(ValueError, match="largebinary URI must start with 's3://'"):
+ largebinary("http://invalid-uri")
+
+ with pytest.raises(ValueError, match="largebinary URI must start with 's3://'"):
+ largebinary("invalid-uri")
+
+ def test_get_bucket_name(self):
+ """Test extracting bucket name from URI."""
+ large_binary = largebinary("s3://my-bucket/path/to/object")
+ assert large_binary.get_bucket_name() == "my-bucket"
+
+ def test_get_object_key(self):
+ """Test extracting object key from URI."""
+ large_binary = largebinary("s3://my-bucket/path/to/object")
+ assert large_binary.get_object_key() == "path/to/object"
+
+ def test_get_object_key_with_leading_slash(self):
+ """Test extracting object key when URI has leading slash."""
+ large_binary = largebinary("s3://my-bucket/path/to/object")
+ # urlparse includes leading slash, but get_object_key removes it
+ assert large_binary.get_object_key() == "path/to/object"
+
+ def test_equality(self):
+ """Test largebinary equality comparison."""
+ uri = "s3://bucket/path"
+ obj1 = largebinary(uri)
+ obj2 = largebinary(uri)
+ obj3 = largebinary("s3://bucket/different")
+
+ assert obj1 == obj2
+ assert obj1 != obj3
+ assert obj1 != "not a largebinary"
+
+ def test_hash(self):
+ """Test largebinary hashing."""
+ uri = "s3://bucket/path"
+ obj1 = largebinary(uri)
+ obj2 = largebinary(uri)
+
+ assert hash(obj1) == hash(obj2)
+ assert hash(obj1) == hash(uri)
+
+ def test_uri_property(self):
+ """Test URI property access."""
+ uri = "s3://test-bucket/test/path"
+ large_binary = largebinary(uri)
+ assert large_binary.uri == uri
diff --git a/amber/src/main/python/core/proxy/__init__.py b/amber/src/main/python/core/proxy/__init__.py
new file mode 100644
index 00000000000..d2324dbdca5
--- /dev/null
+++ b/amber/src/main/python/core/proxy/__init__.py
@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .proxy_client import ProxyClient
+from .proxy_server import ProxyServer
+
+__all__ = ["ProxyClient", "ProxyServer"]
diff --git a/amber/src/main/python/core/proxy/proxy_client.py b/amber/src/main/python/core/proxy/proxy_client.py
new file mode 100644
index 00000000000..4f0055e5391
--- /dev/null
+++ b/amber/src/main/python/core/proxy/proxy_client.py
@@ -0,0 +1,111 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from loguru import logger
+from pyarrow import Table, Buffer
+from pyarrow.flight import (
+ Action,
+ FlightCallOptions,
+ FlightClient,
+ FlightDescriptor,
+ FlightStreamWriter,
+ FlightMetadataReader,
+)
+from typing import Optional
+
+
+class ProxyClient(FlightClient):
+ def __init__(
+ self,
+ scheme: str = "grpc+tcp",
+ host: str = "localhost",
+ port: int = 5005,
+ handshake_port: Optional[int] = None,
+ timeout=1000,
+ *args,
+ **kwargs,
+ ):
+ location = f"{scheme}://{host}:{port}"
+ super().__init__(location, *args, **kwargs)
+ logger.debug(f"Connected to server at {location}")
+ self._timeout = timeout
+ if handshake_port is not None:
+ self._handshake(handshake_port=handshake_port)
+
+ @logger.catch(reraise=True)
+ def call_action(
+ self,
+ action_name: str,
+ payload: bytes = bytes(),
+ options: Optional[FlightCallOptions] = None,
+ ) -> bytes:
+ """
+ Call a specific remote action specified by the name, pass along a payload.
+ :param action_name: the registered action name to be invoked.
+ :param payload: the action payload in bytes, user should take the
+ responsibility to deserialize it.
+ :param options: FlightCallOption to config the call.
+ :return: exactly one result in bytes.
+ """
+
+ action = Action(action_name, payload)
+ if options is None:
+ options = FlightCallOptions(timeout=self._timeout)
+
+ # Arrow allows multiple results from the Action call return as a stream (
+ # interator). In Arrow 11, it alerts if the results are not consumed fully.
+ # As we do our own Async RPC management, we are currently not using results
+ # from Action call. In the future, this results can include credits for flow
+ # control purpose.
+ results = list(self.do_action(action, options))
+
+ # However, we will only expect exactly one result for now.
+ assert len(results) == 1
+
+ return results[0].body.to_pybytes()
+
+ @logger.catch(reraise=True)
+ def send_data(self, command: bytes, table: Optional[Table]) -> int:
+ """
+ Send a data batch to the server.
+ :param command: a command to in descriptor to pass along, user should take
+ the responsibility to deserialize it.
+ :param table: a PyArrow.Table of column-stored records.
+ :return: an integer representing credit values received from ack
+ """
+ descriptor = FlightDescriptor.for_command(command)
+ table = Table.from_arrays([]) if table is None else table
+ writer, reader = self.do_put(descriptor, table.schema)
+ writer: FlightStreamWriter
+ reader: FlightMetadataReader
+ with writer:
+ writer.write_table(table)
+ credit_buf: Buffer = reader.read()
+ credit_count: int = int.from_bytes(
+ credit_buf.to_pybytes(), byteorder="little"
+ )
+ return credit_count
+
+ def _handshake(self, handshake_port: int) -> None:
+ """
+ Send the handshake port to Java Proxy Server, which will be forwarded to
+ the Java Proxy Client to use.
+ :param handshake_port: int, the port number for Java Proxy Client to connect
+ to Python Proxy Server.
+ :return:
+ """
+ self.call_action("handshake", bytes(str(handshake_port), "utf-8"))
diff --git a/amber/src/main/python/core/proxy/proxy_server.py b/amber/src/main/python/core/proxy/proxy_server.py
new file mode 100644
index 00000000000..b3bc3afea06
--- /dev/null
+++ b/amber/src/main/python/core/proxy/proxy_server.py
@@ -0,0 +1,319 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import socket
+import threading
+from functools import wraps
+from inspect import signature
+from loguru import logger
+from overrides import overrides
+from pyarrow import Table, py_buffer, Buffer
+from pyarrow.flight import (
+ Action,
+ FlightDescriptor,
+ FlightServerBase,
+ MetadataRecordBatchReader,
+ FlightMetadataWriter,
+ Result,
+ ServerCallContext,
+)
+from typing import Callable, Dict, Iterator, Optional, Tuple
+
+
+def get_free_local_port():
+ # results a free random port
+ with socket.socket() as s:
+ s.bind(("", 0))
+ return s.getsockname()[1]
+
+
+class ProxyServer(FlightServerBase):
+ """
+ There are three kinds of messages supported by the ProxyServer:
+ 1. Data Messages.
+ Data Messages are passed through the endpoint do_put. It will contain a
+ command and a data batch.
+ The user should provide deserializer for the command and a handler for the
+ data batch.
+ 2. ProxyInternal Messages.
+ ProxyInternal Messages are passed through the endpoint do_action. It will be
+ used to control the life cycle of the ProxyServer. Some example messages:
+ - heartbeat: checks if the ProxyServer is alive.
+ - shutdown: shutdown the ProxyServer.
+ - control: passing a Control Message.
+ 3. Control Messages.
+ Control Messages are passed through the endpoint do_action. It will contain a
+ command and a payload.
+ The user should provide deserializer for both the command and the payload.
+ """
+
+ @staticmethod
+ def ack(original_func: Optional[Callable] = None, msg="ack"):
+ """
+ Decorator for returning an ack message after the action. It is a Proxy level
+ ack, only to be used by ProxyServer actions.
+
+ Example usage:
+ ```
+ @ack
+ def hello():
+ return None
+ server.register("hello", hello)
+ msg = client.call("hello") # msg will be "ack"
+ ```
+
+ or
+ ```
+ @ack(msg="other msg")
+ def hello():
+ return None
+ server.register("hello", hello)
+ msg = client.call("hello") # msg will be "other msg"
+ ```
+
+ :param original_func: decorated function, usually is a callable to be
+ registered.
+ :param msg: the return message from the decorator, "ack" by default.
+ :return:
+ """
+
+ def ack_decorator(func: Callable):
+ @wraps(func)
+ def wrapper(*args, **kwargs):
+ func(*args, **kwargs)
+ return msg
+
+ return wrapper
+
+ if original_func:
+ return ack_decorator(original_func)
+ return ack_decorator
+
+ def __init__(
+ self,
+ scheme: str = "grpc+tcp",
+ host: str = "localhost",
+ port: Optional[int] = None,
+ ):
+ if port is None:
+ port = get_free_local_port()
+ location = f"{scheme}://{host}:{port}"
+ super(ProxyServer, self).__init__(location)
+ logger.debug(f"Serving on {location}")
+
+ self._port_number = port
+
+ # action name to callable map, will contain registered actions,
+ # identified by action name.
+ self._procedures: Dict[str, Tuple[Callable, str]] = dict()
+
+ # register heartbeat, this is the default action for the client to
+ # check the aliveness of the server.
+ self.register(name="heartbeat", action=ProxyServer.ack()(lambda: None))
+
+ # register shutdown, this is the default action for the client to
+ # terminate the server.
+ self.register(
+ name="shutdown",
+ action=ProxyServer.ack(msg="Bye bye!")(
+ lambda: threading.Thread(target=self.graceful_shutdown).start()
+ ),
+ description="Shut down this server.",
+ )
+
+ # register control, set default action for the client to invoke
+ # after receiving control. it should invoke the control_handler defined
+ # in network_receiver and return number of batches in internal_queue to be
+ # used for credit calculation
+ self.register(
+ name="control",
+ action=lambda control_message: self.process_control(control_message),
+ description="Process the control message",
+ )
+
+ self.register(
+ name="actor",
+ action=lambda message: self.process_actor(message),
+ description="Process the actor message",
+ )
+
+ # the data message handler for each data message, needs to be
+ # implemented during runtime.
+ self.process_data = lambda *args, **kwargs: (_ for _ in ()).throw(
+ NotImplementedError
+ )
+
+ # the control message handler for each control message, needs to be
+ # implemented during runtime.
+ self.process_control = lambda *args, **kwargs: (_ for _ in ()).throw(
+ NotImplementedError
+ )
+
+ # the actor command message handler for each actor message, needs to be
+ # implemented during runtime.
+ self.process_actor = lambda *args, **kwargs: (_ for _ in ()).throw(
+ NotImplementedError
+ )
+
+ ###########################
+ # Flights related methods #
+ ###########################
+ @overrides(check_signature=False)
+ def do_put(
+ self,
+ context: ServerCallContext,
+ descriptor: FlightDescriptor,
+ reader: MetadataRecordBatchReader,
+ writer: FlightMetadataWriter,
+ ):
+ """
+ Put a data table into the server, the data will be handled by the
+ `self.process_data()` handler. Also send back number of sender batches
+ currently in internal queue for credit calculations
+
+ :param context: server context, containing information of middlewares.
+ :param descriptor: the descriptor of this batch of data.
+ :param reader: the input stream of batches of records.
+ :param writer: the output stream.
+ :return:
+ """
+
+ data: Table = reader.read_all()
+ command: bytes = descriptor.command
+ logger.debug(f"getting a data batch {data}")
+
+ sender_credits = self.process_data(command, data)
+ if isinstance(sender_credits, int):
+ sender_credits_buf: Buffer = py_buffer(
+ sender_credits.to_bytes(length=8, byteorder="little")
+ )
+ writer.write(sender_credits_buf)
+
+ ###############################
+ # Actions related methods #
+ ###############################
+ @overrides(check_signature=False)
+ def list_actions(self, context: ServerCallContext) -> Iterator[Tuple[str, str]]:
+ """
+ List all actions that are being registered with the server, it will
+ return the action name and description for each registered action.
+
+ :param context: server context, containing information of middlewares.
+ :return: iterator of (action_name, action_description) pairs.
+ """
+ return map(lambda x: (x[0], x[1][1]), self._procedures.items())
+
+ @overrides(check_signature=False)
+ def do_action(self, context: ServerCallContext, action: Action) -> Iterator[Result]:
+ """
+ Perform an action that previously registered with a action,
+ return a result in bytes.
+
+ :param context: server context, containing information of middlewares.
+ :param action: the action to perform, including
+ action.type: the action name to invoke
+ action.body: the action arguments in bytes
+ :return: yield the encoded result back to client.
+ """
+
+ action_name = action.type
+ logger.debug(f"python getting a call on {action_name}")
+ # get action by name
+ if action_name in self._procedures:
+ procedure, _ = self._procedures.get(action_name)
+ if not action:
+ raise KeyError("Unknown action {!r}".format(action_name))
+
+ payload = action.body.to_pybytes()
+ # invoke the action
+ if payload:
+ result = procedure(payload)
+ else:
+ result = procedure()
+
+ # serialize the result
+ if isinstance(result, bytes):
+ encoded = result
+ else:
+ encoded = str(result).encode("utf-8")
+ yield Result(py_buffer(encoded))
+ else:
+ raise KeyError("Unknown action {!r}".format(action_name))
+
+ @logger.catch(reraise=True)
+ def register(self, name: str, action: Callable, description: str = "") -> None:
+ """
+ Register an action with the action name.
+
+ :param name: the name of the action, it should be matching Action's type.
+ :param action: a callable, could be class, function, or lambda.
+ :param description: describes the action.
+ :return:
+ """
+
+ # wrap the given action so that its error can be logged.
+ @logger.catch(level="WARNING", reraise=True)
+ def wrapper(*args, **kwargs):
+ return action(*args, **kwargs)
+
+ # update the actions, which overwrites the previous registration.
+ self._procedures[name] = (wrapper, description)
+ logger.debug(f"registered action {name}")
+
+ @logger.catch(reraise=True)
+ def register_data_handler(self, handler: Callable) -> None:
+ """
+ Register the data handler function, which will be invoked after each `do_put`.
+
+ :param handler: a callable with at least two arguments, for
+ 1) the command and 2) the data batch.
+ :return:
+ """
+
+ # the handler should have at least 2 arguments
+ assert len(signature(handler).parameters) >= 2
+ self.process_data = handler
+
+ @logger.catch(reraise=True)
+ def register_control_handler(self, handler: Callable) -> None:
+ """
+ Register a control handler function, which will be invoked after each
+ `do_action` with `control` as the command.
+
+ :param handler: a callable with at least two arguments, for 1) the command
+ and 2) the control payload.
+ :return:
+ """
+ # the handler should have at least 1 argument
+ assert len(signature(handler).parameters) >= 1
+ self.process_control = handler
+
+ @logger.catch(reraise=True)
+ def register_actor_message_handler(self, handler: Callable) -> None:
+ self.process_actor = handler
+
+ ##################
+ # helper methods #
+ ##################
+ def graceful_shutdown(self):
+ """Shut down after a delay."""
+ logger.debug("Server is shutting down...")
+ super().shutdown()
+ logger.debug("Server is shutdown.")
+
+ def get_port_number(self):
+ return self._port_number
diff --git a/amber/src/main/python/core/proxy/test_proxy_client.py b/amber/src/main/python/core/proxy/test_proxy_client.py
new file mode 100644
index 00000000000..b28b2cfe999
--- /dev/null
+++ b/amber/src/main/python/core/proxy/test_proxy_client.py
@@ -0,0 +1,155 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+from pandas import DataFrame
+from pyarrow import ArrowNotImplementedError, Table
+from queue import Queue
+
+from .proxy_client import ProxyClient
+from .proxy_server import ProxyServer
+
+
+class TestProxyClient:
+ @pytest.fixture
+ def data_queue(self):
+ return Queue()
+
+ @pytest.fixture
+ def server(self):
+ server = ProxyServer(port=5005)
+ yield server
+ server.graceful_shutdown()
+
+ @pytest.fixture
+ def server_with_dp(self, data_queue):
+ server = ProxyServer(port=5005)
+ server.register_data_handler(
+ lambda _, table: list(
+ map(data_queue.put, map(lambda t: t[1], table.to_pandas().iterrows()))
+ )
+ )
+ yield server
+ server.graceful_shutdown()
+
+ class MockFlightMetadataReader:
+ """
+ MockFlightMetadataReader is a mocked FlightMetadataReader class to ultimately
+ mock a credit value to be returned from Scala server to Python client
+ """
+
+ class MockBuffer:
+ def to_pybytes(self):
+ dummy_credit = 31
+ return dummy_credit.to_bytes(8, "little")
+
+ def read(self):
+ return self.MockBuffer()
+
+ @pytest.fixture
+ def client(self):
+ mock_client = ProxyClient()
+
+ def mock_do_put(
+ self,
+ FlightDescriptor_descriptor,
+ Schema_schema,
+ FlightCallOptions_options=None,
+ ):
+ """
+ Mocking FlightClient.do_put that is called in ProxyClient to return
+ a MockFlightMetadataReader instead of a FlightMetadataReader
+
+ :param self: an instance of FlightClient (would be ProxyClient in this case)
+ :param FlightDescriptor_descriptor: descriptor
+ :param Schema_schema: schema
+ :param FlightCallOptions_options: options, None by default
+ :return: writer : FlightStreamWriter, reader : MockFlightMetadataReader
+ """
+ writer, _ = super(ProxyClient, self).do_put(
+ FlightDescriptor_descriptor, Schema_schema, FlightCallOptions_options
+ )
+ reader = TestProxyClient.MockFlightMetadataReader()
+ return writer, reader
+
+ mock_client.do_put = mock_do_put.__get__(
+ mock_client, ProxyClient
+ ) # override do_put with mock_do_put
+
+ yield mock_client
+
+ @pytest.fixture
+ def data_table(self):
+ df_to_sent = DataFrame(
+ {
+ "Brand": ["Honda Civic", "Toyota Corolla", "Ford Focus", "Audi A4"],
+ "Price": [22000, 25000, 27000, 35000],
+ },
+ columns=["Brand", "Price"],
+ )
+ return Table.from_pandas(df_to_sent)
+
+ def test_client_can_connect_to_server(self, server, client):
+ assert client.call_action("heartbeat") == b"ack"
+
+ def test_client_can_shutdown_server(self, server, client):
+ assert client.call_action("shutdown") == b"Bye bye!"
+
+ def test_client_can_call_registered_lambdas(self, server, client):
+ action_count = len(client.list_actions())
+ server.register("hello", lambda: "hello")
+ server.register("this is another call", lambda: "ack!!!")
+ assert len(client.list_actions()) == action_count + 2
+ assert client.call_action("hello") == b"hello"
+ assert client.call_action("this is another call") == b"ack!!!"
+ assert client.call_action("shutdown") == b"Bye bye!"
+
+ def test_client_can_call_registered_function(self, server, client):
+ def hello():
+ return "hello-function"
+
+ action_count = len(client.list_actions())
+ server.register("hello-function", hello)
+ assert len(client.list_actions()) == action_count + 1
+ assert client.call_action("hello-function") == b"hello-function"
+ assert client.call_action("shutdown") == b"Bye bye!"
+
+ def test_client_can_call_registered_callable_class(self, server, client):
+ class HelloClass:
+ def __call__(self):
+ return "hello-class"
+
+ action_count = len(client.list_actions())
+ server.register("hello-class", HelloClass())
+ assert len(client.list_actions()) == action_count + 1
+ assert client.call_action("hello-class") == b"hello-class"
+ assert client.call_action("shutdown") == b"Bye bye!"
+
+ def test_client_cannot_send_data_without_handler(self, server, client, data_table):
+ # send the pyarrow table to server as a flight
+ with pytest.raises(ArrowNotImplementedError):
+ client.send_data(command=bytes(), table=data_table)
+
+ def test_client_can_send_data_with_handler(
+ self, data_queue: Queue, server_with_dp, client, data_table
+ ):
+ # send the pyarrow table to server as a flight
+ client.send_data(bytes(), data_table)
+
+ assert data_queue.qsize() == 4
+ for i, row in data_table.to_pandas().iterrows():
+ assert data_queue.get().equals(row)
diff --git a/amber/src/main/python/core/proxy/test_proxy_server.py b/amber/src/main/python/core/proxy/test_proxy_server.py
new file mode 100644
index 00000000000..a4a422e16f8
--- /dev/null
+++ b/amber/src/main/python/core/proxy/test_proxy_server.py
@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+from pyarrow.flight import Action
+
+from .proxy_server import ProxyServer
+
+
+class TestProxyServer:
+ @pytest.fixture()
+ def server(self):
+ server = ProxyServer()
+ yield server
+ server.graceful_shutdown()
+
+ def test_server_can_register_control_actions_with_lambda(self, server):
+ assert "hello" not in server._procedures
+ server.register("hello", lambda: None)
+ assert "hello" in server._procedures
+
+ def test_server_can_register_control_actions_with_function(self, server):
+ def hello():
+ return None
+
+ assert "hello" not in server._procedures
+ server.register("hello", hello)
+ assert "hello" in server._procedures
+
+ def test_server_can_register_control_actions_with_callable_class(self, server):
+ class Hello:
+ def __call__(self):
+ return None
+
+ assert "hello" not in server._procedures
+ server.register("hello", Hello())
+ assert "hello" in server._procedures
+
+ def test_server_can_invoke_registered_control_actions(self, server):
+ procedure_contents = {
+ "hello": "hello world",
+ "get an int": 12,
+ "get a float": 1.23,
+ "get a tuple": (5, None, 123.4),
+ "get a list": [5, (None, 123.4)],
+ "get a dict": {"entry": [5, (None, 123.4)]},
+ }
+
+ for name, result in procedure_contents.items():
+ server.register(name, lambda: result)
+ assert name in server._procedures
+ assert next(
+ server.do_action(None, Action(name, b""))
+ ).body.to_pybytes() == str(result).encode("utf-8")
diff --git a/amber/src/main/python/core/python_worker.py b/amber/src/main/python/core/python_worker.py
new file mode 100644
index 00000000000..bcd0652d596
--- /dev/null
+++ b/amber/src/main/python/core/python_worker.py
@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from overrides import overrides
+from threading import Thread, Event
+
+from core.models.internal_queue import InternalQueue
+from core.runnables import MainLoop, NetworkReceiver, NetworkSender, Heartbeat
+from core.util.runnable.runnable import Runnable
+from core.util.stoppable.stoppable import Stoppable
+
+
+class PythonWorker(Runnable, Stoppable):
+ def __init__(self, worker_id: str, host: str, output_port: int):
+ self._input_queue = InternalQueue()
+ self._output_queue = InternalQueue()
+ # start the server
+ self._network_receiver = NetworkReceiver(self._input_queue, host=host)
+ # let Java knows where Python starts (do handshake)
+ self._network_sender = NetworkSender(
+ self._output_queue,
+ host=host,
+ port=output_port,
+ handshake_port=self._network_receiver.proxy_server.get_port_number(),
+ )
+ self._stop_event = Event()
+ self._heartbeat = Heartbeat(host, output_port, 5, self._stop_event)
+
+ self._main_loop = MainLoop(worker_id, self._input_queue, self._output_queue)
+ self._network_receiver.register_shutdown(self.stop)
+
+ @overrides
+ def run(self) -> None:
+ network_sender_thread = Thread(
+ target=self._network_sender.run, name="network_sender"
+ )
+ main_loop_thread = Thread(target=self._main_loop.run, name="main_loop_thread")
+
+ heartbeat_thread = Thread(
+ target=self._heartbeat.run,
+ name="heartbeat_thread",
+ )
+
+ network_sender_thread.start()
+ main_loop_thread.start()
+ heartbeat_thread.start()
+ main_loop_thread.join()
+ network_sender_thread.join()
+
+ # if everything finishes, the heartbeat should stop
+ self._stop_event.set()
+
+ heartbeat_thread.join()
+
+ @overrides
+ def stop(self):
+ self._main_loop.stop()
+ self._network_sender.stop()
+ self._heartbeat.stop()
diff --git a/amber/src/main/python/core/runnables/__init__.py b/amber/src/main/python/core/runnables/__init__.py
new file mode 100644
index 00000000000..d132bd3cdeb
--- /dev/null
+++ b/amber/src/main/python/core/runnables/__init__.py
@@ -0,0 +1,23 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .network_receiver import NetworkReceiver
+from .network_sender import NetworkSender
+from .main_loop import MainLoop
+from .heartbeat import Heartbeat
+
+__all__ = ["NetworkReceiver", "NetworkSender", "MainLoop", "Heartbeat"]
diff --git a/amber/src/main/python/core/runnables/data_processor.py b/amber/src/main/python/core/runnables/data_processor.py
new file mode 100644
index 00000000000..276a1669f55
--- /dev/null
+++ b/amber/src/main/python/core/runnables/data_processor.py
@@ -0,0 +1,226 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+import sys
+import traceback
+from loguru import logger
+from threading import Event
+from typing import Iterator, Optional
+
+from core.architecture.managers import Context
+from core.models import ExceptionInfo, State, TupleLike, InternalMarker
+from core.models.internal_marker import StartChannel, EndChannel
+from core.models.table import all_output_to_tuple
+from core.util import Stoppable
+from core.util.console_message.replace_print import replace_print
+from core.util.console_message.timestamp import current_time_in_local_timezone
+from core.util.runnable.runnable import Runnable
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ConsoleMessage,
+ ConsoleMessageType,
+)
+
+
+class DataProcessor(Runnable, Stoppable):
+ def __init__(self, context: Context):
+ self._running = Event()
+ self._context = context
+
+ def run(self) -> None:
+ """
+ Start the data processing loop. Wait for context switch conditions to be met,
+ then continuously process markers or tuples until stopped.
+ """
+ with self._context.tuple_processing_manager.context_switch_condition:
+ self._context.tuple_processing_manager.context_switch_condition.wait()
+ self._running.set()
+ self._pre_loop_checks()
+ while self._running.is_set():
+ tpm = self._context.tuple_processing_manager
+ spm = self._context.state_processing_manager
+ has_marker = tpm.current_internal_marker is not None
+ has_state = spm.current_input_state is not None
+ has_tuple = tpm.current_input_tuple is not None
+ queued = has_marker + has_state + has_tuple
+ # MainLoop is single-threaded and sets at most one of
+ # current_internal_marker / current_input_state /
+ # current_input_tuple per cycle before switching to here, so
+ # exactly one slot must be populated on every iteration.
+ if queued != 1:
+ raise RuntimeError(
+ "DataProcessor expected exactly one queued input per "
+ f"iteration, got marker={has_marker}, state={has_state}, "
+ f"tuple={has_tuple}"
+ )
+ if has_marker:
+ self.process_internal_marker(tpm.get_internal_marker())
+ elif has_state:
+ self.process_state(spm.get_input_state())
+ else:
+ self.process_tuple()
+
+ def process_internal_marker(self, internal_marker: InternalMarker) -> None:
+ try:
+ executor = self._context.executor_manager.executor
+ port_id = self._context.tuple_processing_manager.get_input_port_id()
+ with replace_print(
+ self._context.worker_id,
+ self._context.console_message_manager.print_buf,
+ ):
+ if isinstance(internal_marker, StartChannel):
+ self._set_output_state(executor.produce_state_on_start(port_id))
+ elif isinstance(internal_marker, EndChannel):
+ self._set_output_state(executor.produce_state_on_finish(port_id))
+ self._switch_context()
+ self._set_output_tuple(executor.on_finish(port_id))
+
+ except Exception as err:
+ logger.exception(err)
+ exc_info = sys.exc_info()
+ self._context.exception_manager.set_exception_info(exc_info)
+ self._report_exception(exc_info)
+
+ finally:
+ self._switch_context()
+
+ def process_state(self, state: State) -> None:
+ """
+ Process an input marker by invoking appropriate state
+ or tuple generation based on the marker type.
+ """
+ try:
+ executor = self._context.executor_manager.executor
+ port_id = self._context.tuple_processing_manager.get_input_port_id()
+ with replace_print(
+ self._context.worker_id,
+ self._context.console_message_manager.print_buf,
+ ):
+ self._set_output_state(executor.process_state(state, port_id))
+
+ except Exception as err:
+ logger.exception(err)
+ exc_info = sys.exc_info()
+ self._context.exception_manager.set_exception_info(exc_info)
+ self._report_exception(exc_info)
+
+ finally:
+ self._switch_context()
+
+ def process_tuple(self) -> None:
+ """
+ Process an input tuple by invoking the executor's tuple processing method.
+ """
+ finished_current = self._context.tuple_processing_manager.finished_current
+ while not finished_current.is_set():
+ try:
+ executor = self._context.executor_manager.executor
+ port_id = self._context.tuple_processing_manager.get_input_port_id()
+ tuple_ = self._context.tuple_processing_manager.get_input_tuple()
+ with replace_print(
+ self._context.worker_id,
+ self._context.console_message_manager.print_buf,
+ ):
+ self._set_output_tuple(executor.process_tuple(tuple_, port_id))
+
+ except Exception as err:
+ logger.exception(err)
+ exc_info = sys.exc_info()
+ self._context.exception_manager.set_exception_info(exc_info)
+ self._report_exception(exc_info)
+
+ finally:
+ self._switch_context()
+
+ def _set_output_tuple(self, output_iterator: Iterator[Optional[TupleLike]]) -> None:
+ """
+ Set the output tuple after processing by the executor.
+ """
+ for output in output_iterator:
+ # output could be a None, a TupleLike, or a TableLike.
+ for output_tuple in all_output_to_tuple(output):
+ if output_tuple is not None:
+ output_tuple.finalize(
+ self._context.output_manager.get_port().get_schema()
+ )
+ self._switch_context()
+ self._context.tuple_processing_manager.current_output_tuple = (
+ output_tuple
+ )
+ self._switch_context()
+ self._context.tuple_processing_manager.finished_current.set()
+
+ def _set_output_state(self, output_state: State) -> None:
+ """
+ Set the output state after processing by the executor.
+ """
+ if output_state is not None and not isinstance(output_state, State):
+ output_state = State(output_state)
+ self._context.state_processing_manager.current_output_state = output_state
+
+ def _switch_context(self) -> None:
+ """
+ Notify the MainLoop thread and wait here until being switched back.
+ """
+ with self._context.tuple_processing_manager.context_switch_condition:
+ self._context.tuple_processing_manager.context_switch_condition.notify()
+ self._context.tuple_processing_manager.context_switch_condition.wait()
+ self._post_switch_context_checks()
+
+ def _check_and_process_debug_command(self) -> None:
+ """
+ If a debug command is available, invokes the debugger from this frame.
+ """
+ if self._context.debug_manager.has_debug_command():
+ # Let debugger trace from the current frame.
+ # This line will also trigger cmdloop in the debugger.
+ # This line has no side effects on the current debugger state.
+ self._context.debug_manager.debugger.set_trace()
+
+ def _post_switch_context_checks(self):
+ self._check_and_process_debug_command()
+
+ def _pre_loop_checks(self) -> None:
+ # Runs once after init and before the first task so that a debug
+ # command queued during worker setup fires before any
+ # tuple / state / marker is processed. Only the debug-command
+ # check is needed here -- no task has run yet, so there is no
+ # exception to surface.
+ self._check_and_process_debug_command()
+
+ def _report_exception(self, exc_info: ExceptionInfo):
+ tb = traceback.extract_tb(exc_info[2])
+ filename, line_number, func_name, text = tb[-1]
+ base_name = os.path.basename(filename)
+ module_name, _ = os.path.splitext(base_name)
+ formatted_exception = traceback.format_exception(*exc_info)
+ title: str = formatted_exception[-1].strip()
+ message: str = "\n".join(formatted_exception)
+
+ self._context.console_message_manager.put_message(
+ ConsoleMessage(
+ worker_id=self._context.worker_id,
+ timestamp=current_time_in_local_timezone(),
+ msg_type=ConsoleMessageType.ERROR,
+ source=f"{module_name}:{func_name}:{line_number}",
+ title=title,
+ message=message,
+ )
+ )
+
+ def stop(self):
+ self._running.clear()
diff --git a/amber/src/main/python/core/runnables/heartbeat.py b/amber/src/main/python/core/runnables/heartbeat.py
new file mode 100644
index 00000000000..9199518f3f4
--- /dev/null
+++ b/amber/src/main/python/core/runnables/heartbeat.py
@@ -0,0 +1,112 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+import psutil
+import signal
+import socket
+import urllib.parse
+from loguru import logger
+from overrides import overrides
+from threading import Event
+
+from core.util.runnable.runnable import Runnable
+from core.util.stoppable.stoppable import Stoppable
+
+
+class Heartbeat(Runnable, Stoppable):
+ def __init__(
+ self,
+ host: str,
+ output_port: int,
+ interval: float,
+ event: Event,
+ ):
+ self._original_parent_pid = os.getppid()
+ server_url = urllib.parse.urlparse(f"grpc+tcp://{host}:{output_port}")
+ self._parsed_server_host = server_url.hostname
+ self._parsed_server_port = server_url.port
+ self._interval = interval
+ self._stop_event = event
+
+ @overrides
+ def run(self) -> None:
+ while not self._stop_event.wait(timeout=self._interval):
+ alive = self._check_heartbeat()
+ if not alive:
+ # double check
+ still_alive = self._check_heartbeat()
+
+ if not still_alive:
+ parent_pid = os.getppid()
+ try:
+ parent_status = psutil.Process(
+ self._original_parent_pid
+ ).status()
+ except Exception:
+ parent_status = "NOT FOUND"
+
+ logger.warning(
+ f"Parent process PID {self._original_parent_pid} "
+ "runs unusually."
+ + (
+ f" Parent PID changed to {parent_pid}."
+ if parent_pid != self._original_parent_pid
+ else " Parent PID hasn't changed."
+ )
+ + f" Original parent process Status: {parent_status}"
+ )
+ self.stop()
+ return
+
+ # If JVM crashed and main loop and network sender threads stop, we need
+ # to add this line:
+ # self.stop()
+
+ def _check_heartbeat(self) -> bool:
+ """
+ Attempt to connect to JVM on the specific port. If succeeds, it means the
+ socket is still available and the JVM is still alive. Otherwise, the JVM
+ might have been gone.
+
+ :return: bool, indicating if the socket is available.
+ """
+ try:
+ temp_socket = socket.create_connection(
+ (self._parsed_server_host, self._parsed_server_port), timeout=1
+ )
+ temp_socket.close()
+ return True
+ except Exception as e:
+ logger.warning(f"Server is down with exception: {e}")
+ return False
+
+ @overrides
+ def stop(self):
+ # clean up every process under the python worker
+ current_process = psutil.Process()
+ children = current_process.children(recursive=True)
+ for child in children:
+ if child.is_running():
+ try:
+ os.kill(child.pid, signal.SIGKILL)
+ except Exception as e:
+ logger.warning(
+ "Exception during process termination "
+ f"PID {str(child.pid)}: {e} "
+ )
+ os.kill(os.getpid(), signal.SIGTERM)
diff --git a/amber/src/main/python/core/runnables/main_loop.py b/amber/src/main/python/core/runnables/main_loop.py
new file mode 100644
index 00000000000..ab35cda81b9
--- /dev/null
+++ b/amber/src/main/python/core/runnables/main_loop.py
@@ -0,0 +1,481 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import threading
+import time
+import typing
+from loguru import logger
+from overrides import overrides
+from pampy import match
+from typing import Iterator, Optional
+
+from core.architecture.managers.context import Context
+from core.architecture.managers.pause_manager import PauseType
+from core.architecture.rpc.async_rpc_client import AsyncRPCClient
+from core.architecture.rpc.async_rpc_server import AsyncRPCServer
+from core.models import (
+ InternalQueue,
+ Tuple,
+)
+from core.models.internal_marker import StartChannel, EndChannel
+from core.models.internal_queue import (
+ DataElement,
+ DCMElement,
+ ECMElement,
+ InternalQueueElement,
+)
+from core.models.state import State
+from core.runnables.data_processor import DataProcessor
+from core.util import StoppableQueueBlockingRunnable, get_one_of
+from core.util.console_message.timestamp import current_time_in_local_timezone
+from core.util.customized_queue.queue_base import QueueElement
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ PortIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ConsoleMessage,
+ ControlInvocation,
+ ConsoleMessageType,
+ ReturnInvocation,
+ PortCompletedRequest,
+ EmptyRequest,
+ ConsoleMessageTriggeredRequest,
+ EmbeddedControlMessageType,
+ EmbeddedControlMessage,
+ AsyncRpcContext,
+ ControlRequest,
+)
+from proto.org.apache.texera.amber.engine.architecture.worker import (
+ WorkerState,
+)
+
+
+class MainLoop(StoppableQueueBlockingRunnable):
+ def __init__(
+ self,
+ worker_id: str,
+ input_queue: InternalQueue,
+ output_queue: InternalQueue,
+ ):
+ super().__init__(self.__class__.__name__, queue=input_queue)
+ self._input_queue: InternalQueue = input_queue
+ self._output_queue: InternalQueue = output_queue
+
+ self.context = Context(worker_id, input_queue)
+ self._async_rpc_server = AsyncRPCServer(output_queue, context=self.context)
+ self._async_rpc_client = AsyncRPCClient(output_queue, context=self.context)
+
+ self.data_processor = DataProcessor(self.context)
+ threading.Thread(
+ target=self.data_processor.run, daemon=True, name="data_processor_thread"
+ ).start()
+
+ def complete(self) -> None:
+ """
+ Complete the DataProcessor, marking state to COMPLETED, and notify the
+ controller.
+ """
+ # flush the buffered console prints
+ self._check_and_report_console_messages(force_flush=True)
+ self.context.executor_manager.executor.close()
+ # stop the data processing thread
+ self.data_processor.stop()
+ self.context.state_manager.transit_to(WorkerState.COMPLETED)
+ self.context.statistics_manager.update_total_execution_time(time.time_ns())
+ controller_interface = self._async_rpc_client.controller_stub()
+ controller_interface.worker_execution_completed(EmptyRequest())
+ self.context.close()
+
+ def _check_and_process_control(self) -> None:
+ """
+ Check if there exists any ControlElement(s) in the input_queue, if so, take and
+ process them one by one.
+
+ This is used very frequently as we want to prioritize the process of
+ ControlElement, and will be invoked many times during a DataElement's
+ processing lifecycle. Thus, this method's invocation could appear in any
+ stage while processing a DataElement.
+ """
+ while (
+ not self._input_queue.is_control_empty()
+ or not self._input_queue.is_data_enabled()
+ ):
+ next_entry = self.interruptible_get()
+ match(
+ next_entry,
+ DCMElement,
+ self._process_dcm,
+ ECMElement,
+ self._process_ecm,
+ )
+
+ @overrides
+ def pre_start(self) -> None:
+ self.context.state_manager.assert_state(WorkerState.UNINITIALIZED)
+ self.context.state_manager.transit_to(WorkerState.READY)
+ self.context.statistics_manager.initialize_worker_start_time(time.time_ns())
+
+ @overrides
+ def receive(self, next_entry: QueueElement) -> None:
+ """
+ Main entry point of the DataProcessor. Upon receipt of an next_entry,
+ process it respectfully.
+
+ :param next_entry: An entry from input_queue, could be one of the followings:
+ 1. a ControlElement;
+ 2. a DataElement.
+ """
+ if isinstance(next_entry, InternalQueueElement):
+ self.context.current_input_channel_id = next_entry.tag
+
+ match(
+ next_entry,
+ DataElement,
+ self._process_data_element,
+ DCMElement,
+ self._process_dcm,
+ ECMElement,
+ self._process_ecm,
+ )
+
+ def process_input_tuple(self) -> None:
+ """
+ Process the current input tuple with the current input link.
+ Send all result Tuples or State to downstream workers.
+
+ This is being invoked for each Tuple that are unpacked from the DataElement.
+ """
+ if isinstance(self.context.tuple_processing_manager.current_input_tuple, Tuple):
+ self.context.statistics_manager.increase_input_statistics(
+ self.context.tuple_processing_manager.current_input_port_id,
+ self.context.tuple_processing_manager.current_input_tuple.in_mem_size(),
+ )
+
+ for output_tuple in self.process_tuple_with_udf():
+ self._check_and_process_control()
+ if output_tuple is not None:
+ self.context.statistics_manager.increase_output_statistics(
+ PortIdentity(0), output_tuple.in_mem_size()
+ )
+ for to, batch in self.context.output_manager.tuple_to_batch(
+ output_tuple
+ ):
+ self._output_queue.put(
+ DataElement(
+ tag=ChannelIdentity(
+ ActorVirtualIdentity(self.context.worker_id), to, False
+ ),
+ payload=batch,
+ )
+ )
+ self.context.output_manager.save_tuple_to_storage_if_needed(
+ output_tuple
+ )
+
+ def process_input_state(self) -> None:
+ self._switch_context()
+ output_state = self.context.state_processing_manager.get_output_state()
+ if output_state is not None:
+ for to, batch in self.context.output_manager.emit_state(output_state):
+ self._output_queue.put(
+ DataElement(
+ tag=ChannelIdentity(
+ ActorVirtualIdentity(self.context.worker_id), to, False
+ ),
+ payload=batch,
+ )
+ )
+
+ def process_tuple_with_udf(self) -> Iterator[Optional[Tuple]]:
+ """
+ Process the Tuple/InputExhausted with the current link.
+
+ This is a wrapper to invoke processing of the executor.
+
+ :return: Iterator[Tuple], iterator of result Tuple(s).
+ """
+ finished_current = self.context.tuple_processing_manager.finished_current
+ finished_current.clear()
+
+ while not finished_current.is_set():
+ self._check_and_process_control()
+ self._switch_context()
+ yield self.context.tuple_processing_manager.get_output_tuple()
+
+ def _process_dcm(self, dcm_element: DCMElement) -> None:
+ """
+ Upon receipt of a ControlElement, unpack it into tag and payload to be handled.
+
+ :param dcm_element: DirectControlMessageElement to be handled.
+ """
+ start_time = time.time_ns()
+ match(
+ (dcm_element.tag, get_one_of(dcm_element.payload, sealed=False)),
+ typing.Tuple[ChannelIdentity, ControlInvocation],
+ self._async_rpc_server.receive,
+ typing.Tuple[ChannelIdentity, ReturnInvocation],
+ self._async_rpc_client.receive,
+ )
+ end_time = time.time_ns()
+ self.context.statistics_manager.increase_control_processing_time(
+ end_time - start_time
+ )
+ self.context.statistics_manager.update_total_execution_time(end_time)
+
+ def _process_tuple(self, tuple_: Tuple) -> None:
+ self.context.tuple_processing_manager.current_input_tuple = tuple_
+ self.process_input_tuple()
+ self._check_and_process_control()
+
+ def _process_state(self, state_: State) -> None:
+ self.context.state_processing_manager.current_input_state = state_
+ self.process_input_state()
+ self._check_and_process_control()
+
+ def _process_start_channel(self) -> None:
+ self._send_ecm_to_data_channels(
+ "StartChannel", EmbeddedControlMessageType.NO_ALIGNMENT
+ )
+ self.process_input_state()
+
+ def _process_end_channel(self) -> None:
+ self.process_input_state()
+ self.process_input_tuple()
+
+ input_port_id = self.context.input_manager.get_port_id(
+ self.context.current_input_channel_id
+ )
+
+ if input_port_id is not None:
+ self._async_rpc_client.controller_stub().port_completed(
+ PortCompletedRequest(
+ port_id=input_port_id,
+ input=True,
+ )
+ )
+
+ if self.context.input_manager.all_ports_completed():
+ # Special case for the hack of input port dependency.
+ # See documentation of is_missing_output_ports
+ if self.context.output_manager.is_missing_output_ports():
+ return
+ self.context.output_manager.close_port_storage_writers()
+
+ self._send_ecm_to_data_channels(
+ "EndChannel", EmbeddedControlMessageType.PORT_ALIGNMENT
+ )
+
+ # Need to send port completed even if there is no downstream link
+ for port_id in self.context.output_manager.get_port_ids():
+ self._async_rpc_client.controller_stub().port_completed(
+ PortCompletedRequest(port_id=port_id, input=False)
+ )
+ self.complete()
+
+ def _process_ecm(self, ecm_element: ECMElement):
+ """
+ Processes a received ECM and handles synchronization,
+ command execution, and forwarding to downstream channels if applicable.
+
+ Args:
+ ecm_element (ECMElement): The received ECM element.
+ """
+ ecm = ecm_element.payload
+ command = ecm.command_mapping.get(self.context.worker_id)
+ channel_id = self.context.current_input_channel_id
+ logger.info(
+ f"receive channel ECM from {channel_id}, id = {ecm.id}, cmd = {command}"
+ )
+ if ecm.ecm_type != EmbeddedControlMessageType.NO_ALIGNMENT:
+ self.context.pause_manager.pause_input_channel(
+ PauseType.ECM_PAUSE, channel_id
+ )
+
+ if self.context.ecm_manager.is_ecm_aligned(channel_id, ecm):
+ logger.info(
+ f"process channel ECM from {channel_id}, id = {ecm.id}, cmd = {command}"
+ )
+
+ if command is not None:
+ self._async_rpc_server.receive(channel_id, command)
+
+ downstream_channels_in_scope = {
+ scope
+ for scope in ecm.scope
+ if scope.from_worker_id == ActorVirtualIdentity(self.context.worker_id)
+ }
+ if downstream_channels_in_scope:
+ for (
+ active_channel_id
+ ) in self.context.output_manager.get_output_channel_ids():
+ if active_channel_id in downstream_channels_in_scope:
+ logger.info(
+ f"send ECM to {active_channel_id},"
+ f" id = {ecm.id}, cmd = {command}"
+ )
+ self._send_ecm_to_channel(active_channel_id, ecm)
+
+ if ecm.ecm_type != EmbeddedControlMessageType.NO_ALIGNMENT:
+ self.context.pause_manager.resume(PauseType.ECM_PAUSE)
+
+ if self.context.tuple_processing_manager.current_internal_marker:
+ {
+ StartChannel: self._process_start_channel,
+ EndChannel: self._process_end_channel,
+ }[type(self.context.tuple_processing_manager.current_internal_marker)]()
+
+ def _send_ecm_to_data_channels(
+ self, method_name: str, alignment: EmbeddedControlMessageType
+ ) -> None:
+ for active_channel_id in self.context.output_manager.get_output_channel_ids():
+ if not active_channel_id.is_control:
+ ecm = EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity(method_name),
+ alignment,
+ [],
+ {
+ active_channel_id.to_worker_id.name: ControlInvocation(
+ method_name,
+ ControlRequest(empty_request=EmptyRequest()),
+ AsyncRpcContext(
+ ActorVirtualIdentity(), ActorVirtualIdentity()
+ ),
+ -1,
+ )
+ },
+ )
+ self._send_ecm_to_channel(active_channel_id, ecm)
+
+ def _send_ecm_to_channel(
+ self, channel_id: ChannelIdentity, ecm: EmbeddedControlMessage
+ ) -> None:
+ for batch in self.context.output_manager.emit_ecm(channel_id.to_worker_id, ecm):
+ tag = channel_id
+ element = (
+ ECMElement(tag=tag, payload=batch)
+ if isinstance(batch, EmbeddedControlMessage)
+ else DataElement(tag=tag, payload=batch)
+ )
+ self._output_queue.put(element)
+
+ def _process_data_element(self, data_element: DataElement) -> None:
+ """
+ Upon receipt of a DataElement, unpack it into Tuples and States,
+ and process them one by one.
+
+ :param data_element: DataElement, a batch of data.
+ """
+
+ self.context.tuple_processing_manager.current_input_port_id = (
+ self.context.input_manager.get_port_id(
+ self.context.current_input_channel_id
+ )
+ )
+
+ # Update state to RUNNING
+ if self.context.state_manager.confirm_state(WorkerState.READY):
+ self.context.state_manager.transit_to(WorkerState.RUNNING)
+
+ self.context.tuple_processing_manager.current_input_tuple_iter = (
+ self.context.input_manager.process_data_payload(
+ data_element.tag, data_element.payload
+ )
+ )
+
+ if self.context.tuple_processing_manager.current_input_tuple_iter is None:
+ return
+ # here the self.context.processing_manager.current_input_iter
+ # could be modified during iteration, thus we are using the while :=
+ # way to iterate through the iterator, instead of the for-each-loop
+ # syntax sugar.
+ while (
+ element := next(
+ self.context.tuple_processing_manager.current_input_tuple_iter, None
+ )
+ ) is not None:
+ try:
+ match(
+ element,
+ Tuple,
+ self._process_tuple,
+ State,
+ self._process_state,
+ )
+ except Exception as err:
+ logger.exception(err)
+
+ def _send_console_message(self, console_message: ConsoleMessage):
+ self._async_rpc_client.controller_stub().console_message_triggered(
+ ConsoleMessageTriggeredRequest(console_message=console_message)
+ )
+
+ def _switch_context(self) -> None:
+ """
+ Notify the DataProcessor thread and wait here until being switched back.
+ """
+ start_time = time.time_ns()
+ with self.context.tuple_processing_manager.context_switch_condition:
+ self.context.tuple_processing_manager.context_switch_condition.notify()
+ self.context.tuple_processing_manager.context_switch_condition.wait()
+ self._post_switch_context_checks()
+ end_time = time.time_ns()
+ self.context.statistics_manager.increase_data_processing_time(
+ end_time - start_time
+ )
+ self.context.statistics_manager.update_total_execution_time(end_time)
+
+ def _check_and_report_debug_event(self) -> None:
+ if self.context.debug_manager.has_debug_event():
+ debug_event = self.context.debug_manager.get_debug_event()
+ self._send_console_message(
+ ConsoleMessage(
+ worker_id=self.context.worker_id,
+ timestamp=current_time_in_local_timezone(),
+ msg_type=ConsoleMessageType.DEBUGGER,
+ source="(Pdb)",
+ title=debug_event,
+ message="",
+ )
+ )
+ self._check_and_report_console_messages(force_flush=True)
+ self.context.pause_manager.pause(PauseType.DEBUG_PAUSE)
+
+ def _check_exception(self) -> None:
+ if self.context.exception_manager.has_exception():
+ self._check_and_report_console_messages(force_flush=True)
+ self.context.pause_manager.pause(PauseType.EXCEPTION_PAUSE)
+
+ def _check_and_report_console_messages(self, force_flush=False) -> None:
+ for msg in self.context.console_message_manager.get_messages(force_flush):
+ self._send_console_message(msg)
+
+ def _post_switch_context_checks(self) -> None:
+ """
+ Post callback for switch context.
+
+ One step in DataProcessor could produce some results, which includes
+ - print messages
+ - Debug Event
+ - Exception
+ We check and report them each time coming back from DataProcessor.
+ """
+ self._check_and_report_console_messages(force_flush=True)
+ self._check_and_report_debug_event()
+ self._check_exception()
diff --git a/amber/src/main/python/core/runnables/network_receiver.py b/amber/src/main/python/core/runnables/network_receiver.py
new file mode 100644
index 00000000000..659cd65c78d
--- /dev/null
+++ b/amber/src/main/python/core/runnables/network_receiver.py
@@ -0,0 +1,173 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from loguru import logger
+from overrides import overrides
+from pampy import match
+from pyarrow.lib import Table
+from typing import Optional
+
+from core.architecture.handlers.actorcommand.actor_handler_base import (
+ ActorCommandHandler,
+)
+from core.architecture.handlers.actorcommand.backpressure_handler import (
+ BackpressureHandler,
+)
+from core.architecture.handlers.actorcommand.credit_update_handler import (
+ CreditUpdateHandler,
+)
+from core.models import (
+ DataFrame,
+ State,
+ StateFrame,
+)
+from core.models.internal_queue import (
+ DataElement,
+ DCMElement,
+ InternalQueue,
+ ECMElement,
+)
+from core.proxy import ProxyServer
+from core.util import Stoppable, get_one_of
+from core.util.runnable.runnable import Runnable
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.common import (
+ PythonControlMessage,
+ PythonDataHeader,
+ PythonActorMessage,
+ ActorCommand,
+)
+
+
+class NetworkReceiver(Runnable, Stoppable):
+ """
+ Receive and deserialize messages.
+ """
+
+ @logger.catch(reraise=True)
+ def __init__(
+ self, shared_queue: InternalQueue, host: str, port: Optional[int] = None
+ ):
+ server_start = False
+ # try to start the server until it succeeds
+ while not server_start:
+ try:
+ self._proxy_server = ProxyServer(host=host, port=port)
+ server_start = True
+ except Exception as e:
+ logger.debug("Error occurred while starting the server:", repr(e))
+
+ self._handlers: dict[type(ActorCommand), ActorCommandHandler] = dict()
+
+ self.register_actor_command_handler(BackpressureHandler())
+ self.register_actor_command_handler(CreditUpdateHandler())
+
+ # register the data handler to deserialize data messages.
+ @logger.catch(reraise=True)
+ def data_handler(command: bytes, table: Table) -> int:
+ """
+ Data handler for deserializing data messages
+
+ :param command:
+ :param table:
+ :return: sender credits
+ """
+ data_header = PythonDataHeader().parse(command)
+ # Explicitly set is_control to trigger lazy computation.
+ # If not set, it may be computed at different times,
+ # causing hash inconsistencies.
+ data_header.tag.is_control = bool(data_header.tag.is_control)
+ payload = match(
+ data_header.payload_type,
+ "Data",
+ lambda _: DataFrame(table),
+ "State",
+ lambda _: StateFrame(State.from_json(table[State.CONTENT][0].as_py())),
+ "ECM",
+ lambda _: EmbeddedControlMessage().parse(table["payload"][0].as_py()),
+ )
+ if isinstance(payload, EmbeddedControlMessage):
+ for channel_id in payload.scope:
+ channel_id.is_control = bool(channel_id.is_control)
+ shared_queue.put(ECMElement(tag=data_header.tag, payload=payload))
+ else:
+ shared_queue.put(DataElement(tag=data_header.tag, payload=payload))
+ return shared_queue.in_mem_size()
+
+ self._proxy_server.register_data_handler(data_handler)
+
+ @logger.catch(reraise=True)
+ def control_handler(message: bytes) -> int:
+ """
+ Control handler for deserializing control messages
+
+ :param message:
+ :return: sender credits
+ """
+ python_control_message = PythonControlMessage().parse(message)
+ shared_queue.put(
+ DCMElement(
+ tag=python_control_message.tag,
+ payload=python_control_message.payload,
+ )
+ )
+ return shared_queue.in_mem_size()
+
+ self._proxy_server.register_control_handler(control_handler)
+
+ @logger.catch(reraise=True)
+ def actor_message_handler(message: bytes) -> int:
+ """
+ Control handler for deserializing actor messages
+
+ :param message:
+ :return: sender credits
+ """
+ python_actor_message = PythonActorMessage().parse(message)
+ command = get_one_of(python_actor_message.payload)
+ self.look_up(command)(command, shared_queue)
+ return shared_queue.in_mem_size()
+
+ self._proxy_server.register_actor_message_handler(actor_message_handler)
+
+ def register_shutdown(self, shutdown: callable) -> None:
+ self._proxy_server.register(
+ name="shutdown", action=ProxyServer.ack(msg="Bye bye!")(shutdown)
+ )
+
+ @logger.catch(reraise=True)
+ @overrides
+ def run(self) -> None:
+ logger.debug("started running!!!")
+ self._proxy_server.serve()
+
+ @logger.catch(reraise=True)
+ @overrides
+ def stop(self):
+ self._proxy_server.graceful_shutdown()
+ self._proxy_server.wait()
+
+ @property
+ def proxy_server(self):
+ return self._proxy_server
+
+ def register_actor_command_handler(self, handler: ActorCommandHandler) -> None:
+ self._handlers[handler.cmd] = handler
+
+ def look_up(self, cmd: ActorCommand) -> ActorCommandHandler:
+ logger.debug(cmd)
+ return self._handlers[type(cmd)]
diff --git a/amber/src/main/python/core/runnables/network_sender.py b/amber/src/main/python/core/runnables/network_sender.py
new file mode 100644
index 00000000000..d8e3889ac11
--- /dev/null
+++ b/amber/src/main/python/core/runnables/network_sender.py
@@ -0,0 +1,126 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+from loguru import logger
+from overrides import overrides
+from typing import Optional
+
+from core.models import DataPayload, InternalQueue, DataFrame, State, StateFrame
+from core.models.internal_queue import (
+ InternalQueueElement,
+ DataElement,
+ DCMElement,
+ ECMElement,
+)
+from core.proxy import ProxyClient
+from core.util import StoppableQueueBlockingRunnable
+from proto.org.apache.texera.amber.core import ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage
+from proto.org.apache.texera.amber.engine.common import (
+ DirectControlMessagePayloadV2,
+ PythonControlMessage,
+ PythonDataHeader,
+)
+
+
+class NetworkSender(StoppableQueueBlockingRunnable):
+ """
+ Serialize and send messages.
+ """
+
+ def __init__(
+ self,
+ shared_queue: InternalQueue,
+ host: str,
+ port: int,
+ handshake_port: Optional[int] = None,
+ ):
+ super().__init__(self.__class__.__name__, queue=shared_queue)
+ self._proxy_client = ProxyClient(
+ host=host, port=port, handshake_port=handshake_port
+ )
+
+ @overrides(check_signature=False)
+ def receive(self, next_entry: InternalQueueElement):
+ if isinstance(next_entry, DataElement):
+ self._send_data(next_entry.tag, next_entry.payload)
+ elif isinstance(next_entry, DCMElement):
+ self._send_control(next_entry.tag, next_entry.payload)
+ elif isinstance(next_entry, ECMElement):
+ self._send_ecm(next_entry.tag, next_entry.payload)
+ else:
+ raise TypeError(f"Unexpected entry {next_entry}")
+
+ @logger.catch(reraise=True)
+ def _send_ecm(self, to: ChannelIdentity, ecm: EmbeddedControlMessage) -> None:
+ """
+ Sends an ECM to the specified channel.
+
+ Args:
+ to (ChannelIdentity): The target channel to which the ECM should be sent.
+ ecm (EmbeddedControlMessage): The ECM to send.
+
+ This function constructs a `PythonDataHeader` with the appropriate metadata,
+ serializes the payload into an Arrow table, and sends it using the proxy client.
+ """
+ data_header = PythonDataHeader(tag=to, payload_type="ECM")
+ schema = pa.schema([("payload", pa.binary())])
+ data = [pa.array([bytes(ecm)])]
+ table = pa.Table.from_arrays(data, schema=schema)
+ self._proxy_client.send_data(bytes(data_header), table)
+
+ @logger.catch(reraise=True)
+ def _send_data(self, to: ChannelIdentity, data_payload: DataPayload) -> None:
+ """
+ Send data payload to the given target actor. This method is to be used
+ internally only.
+
+ :param to: The target ChannelIdentity
+ :param data_payload: The data payload to be sent in DataFrame
+ """
+
+ if isinstance(data_payload, DataFrame):
+ data_header = PythonDataHeader(tag=to, payload_type="Data")
+ self._proxy_client.send_data(bytes(data_header), data_payload.frame)
+ elif isinstance(data_payload, StateFrame):
+ data_header = PythonDataHeader(tag=to, payload_type="State")
+ table = pa.Table.from_pydict(
+ {State.CONTENT: [data_payload.frame.to_json()]},
+ schema=State.SCHEMA.as_arrow_schema(),
+ )
+ self._proxy_client.send_data(bytes(data_header), table)
+ else:
+ raise TypeError(f"Unexpected payload {data_payload}")
+
+ @logger.catch(reraise=True)
+ def _send_control(
+ self, to: ChannelIdentity, control_payload: DirectControlMessagePayloadV2
+ ) -> None:
+ """
+ Send the control payload to the given target actor. This method is to be used
+ internally only.
+
+ :param to: The target ChannelIdentity
+ :param control_payload: The control payload to be sent, can be either
+ ControlInvocation or ReturnInvocation.
+ """
+ python_control_message = PythonControlMessage(tag=to, payload=control_payload)
+ int.from_bytes(
+ self._proxy_client.call_action("control", bytes(python_control_message)),
+ byteorder="little",
+ ) # returned credits
diff --git a/amber/src/main/python/core/runnables/test_console_message.py b/amber/src/main/python/core/runnables/test_console_message.py
new file mode 100644
index 00000000000..fbe9041f96d
--- /dev/null
+++ b/amber/src/main/python/core/runnables/test_console_message.py
@@ -0,0 +1,90 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import pytest
+
+from core.models.internal_queue import InternalQueue
+from core.util import set_one_of
+from core.util.buffer.timed_buffer import TimedBuffer
+from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ControlInvocation,
+ ControlRequest,
+ ConsoleMessage,
+ ConsoleMessageType,
+)
+from proto.org.apache.texera.amber.engine.common import (
+ DirectControlMessagePayloadV2,
+ PythonControlMessage,
+)
+
+
+class TestConsoleMessage:
+ @pytest.fixture
+ def internal_queue(self):
+ return InternalQueue()
+
+ @pytest.fixture
+ def timed_buffer(self):
+ return TimedBuffer()
+
+ @pytest.fixture
+ def console_message(self):
+ return ConsoleMessage(
+ worker_id="0",
+ timestamp=datetime.datetime.now(),
+ msg_type=ConsoleMessageType.PRINT,
+ source="pytest",
+ title="Test Message",
+ message="Test Message",
+ )
+
+ @pytest.fixture
+ def mock_controller_channel(self):
+ return ChannelIdentity(
+ ActorVirtualIdentity("CONTROLLER"), ActorVirtualIdentity("test"), True
+ )
+
+ @pytest.mark.timeout(2)
+ def test_console_message_serialization(
+ self, mock_controller_channel, console_message
+ ):
+ """
+ Test the serialization of the console message
+ :param mock_controller_channel: the mock control channel id
+ :param console_message: the test message
+ """
+ # below statements wrap the console message as the python control message
+ command = set_one_of(ControlRequest, console_message)
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="ConsoleMessageTriggered", command_id=1, command=command
+ ),
+ )
+ python_control_message = PythonControlMessage(
+ tag=mock_controller_channel, payload=payload
+ )
+ # serialize the python control message to bytes
+ python_control_message_bytes = bytes(python_control_message)
+ # deserialize the control message from bytes
+ parsed_python_control_message = PythonControlMessage().parse(
+ python_control_message_bytes
+ )
+ # deserialized one should equal to the original one
+ assert python_control_message == parsed_python_control_message
diff --git a/amber/src/main/python/core/runnables/test_main_loop.py b/amber/src/main/python/core/runnables/test_main_loop.py
new file mode 100644
index 00000000000..c9daa633f55
--- /dev/null
+++ b/amber/src/main/python/core/runnables/test_main_loop.py
@@ -0,0 +1,1565 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import inspect
+import pandas
+import pickle
+import pyarrow
+import pytest
+import time
+from threading import Thread
+
+from core.models import (
+ DataFrame,
+ InternalQueue,
+ State,
+ StateFrame,
+ Tuple,
+)
+from core.models.internal_queue import (
+ DataElement,
+ DCMElement,
+ ECMElement,
+)
+from core.runnables import MainLoop
+from core.util import set_one_of
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ PhysicalLink,
+ PhysicalOpIdentity,
+ OperatorIdentity,
+ ChannelIdentity,
+ PortIdentity,
+ OpExecWithCode,
+ OpExecInitInfo,
+ EmbeddedControlMessageIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ControlRequest,
+ AssignPortRequest,
+ ControlInvocation,
+ AddInputChannelRequest,
+ InitializeExecutorRequest,
+ EmptyReturn,
+ ReturnInvocation,
+ ControlReturn,
+ WorkerMetricsResponse,
+ AddPartitioningRequest,
+ EmptyRequest,
+ PortCompletedRequest,
+ AsyncRpcContext,
+ WorkerStateResponse,
+ EmbeddedControlMessageType,
+ EmbeddedControlMessage,
+)
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ OneToOnePartitioning,
+ Partitioning,
+)
+from proto.org.apache.texera.amber.engine.architecture.worker import (
+ WorkerMetrics,
+ WorkerState,
+ WorkerStatistics,
+ PortTupleMetricsMapping,
+ TupleMetrics,
+)
+from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2
+from pytexera.udf.examples.count_batch_operator import CountBatchOperator
+from pytexera.udf.examples.echo_operator import EchoOperator
+
+
+class TestMainLoop:
+ @pytest.fixture
+ def command_sequence(self):
+ return 1
+
+ @pytest.fixture
+ def mock_link(self):
+ return PhysicalLink(
+ from_op_id=PhysicalOpIdentity(OperatorIdentity("from"), "from"),
+ from_port_id=PortIdentity(0, internal=False),
+ to_op_id=PhysicalOpIdentity(OperatorIdentity("to"), "to"),
+ to_port_id=PortIdentity(0, internal=False),
+ )
+
+ @pytest.fixture
+ def mock_tuple(self):
+ return Tuple({"test-1": "hello", "test-2": 10})
+
+ @pytest.fixture
+ def mock_binary_tuple(self):
+ return Tuple({"test-1": [1, 2, 3, 4], "test-2": 10})
+
+ @pytest.fixture
+ def mock_batch(self):
+ batch_list = []
+ for i in range(57):
+ batch_list.append(Tuple({"test-1": "hello", "test-2": i}))
+ return batch_list
+
+ @pytest.fixture
+ def mock_sender_actor(self):
+ return ActorVirtualIdentity("sender")
+
+ @pytest.fixture
+ def mock_data_input_channel(self):
+ return ChannelIdentity(
+ ActorVirtualIdentity("sender"),
+ ActorVirtualIdentity("dummy_worker_id"),
+ False,
+ )
+
+ @pytest.fixture
+ def mock_data_output_channel(self):
+ return ChannelIdentity(
+ ActorVirtualIdentity("dummy_worker_id"),
+ ActorVirtualIdentity("dummy_worker_id"),
+ False,
+ )
+
+ @pytest.fixture
+ def mock_control_input_channel(self):
+ return ChannelIdentity(
+ ActorVirtualIdentity("CONTROLLER"),
+ ActorVirtualIdentity("dummy_worker_id"),
+ True,
+ )
+
+ @pytest.fixture
+ def mock_control_output_channel(self):
+ return ChannelIdentity(
+ ActorVirtualIdentity("dummy_worker_id"),
+ ActorVirtualIdentity("CONTROLLER"),
+ True,
+ )
+
+ @pytest.fixture
+ def mock_receiver_actor(self):
+ return ActorVirtualIdentity("dummy_worker_id")
+
+ @pytest.fixture
+ def mock_data_element(self, mock_tuple, mock_data_input_channel):
+ return DataElement(
+ tag=mock_data_input_channel,
+ payload=DataFrame(
+ frame=pyarrow.Table.from_pandas(
+ pandas.DataFrame([mock_tuple.as_dict()])
+ )
+ ),
+ )
+
+ @pytest.fixture
+ def mock_state_data_elements(self, mock_data_input_channel):
+ elements = []
+ for value in (1, 2, 3, 4):
+ state = State({"value": value})
+ elements.append(
+ DataElement(
+ tag=mock_data_input_channel,
+ payload=StateFrame(frame=state),
+ )
+ )
+ return elements
+
+ @pytest.fixture
+ def state_processing_executor(self):
+ # In-process executor for the state-pipeline tests. Tags processed
+ # states with `processed_marker` and emits a finish-marker state
+ # from `produce_state_on_finish` so EndChannel handling can be
+ # observed.
+ class StateProcessingExecutor:
+ @staticmethod
+ def process_tuple(tuple_, port):
+ yield tuple_
+
+ @staticmethod
+ def process_state(state: State, port: int) -> State:
+ new_state = State(
+ {key: value for key, value in state.items() if key != "schema"}
+ )
+ new_state["processed_marker"] = "executed"
+ new_state["port"] = port
+ return new_state
+
+ @staticmethod
+ def produce_state_on_finish(port: int) -> State:
+ return State({"finish_marker": "produce_state_on_finish_ran"})
+
+ @staticmethod
+ def on_finish(port):
+ yield
+
+ @staticmethod
+ def close():
+ pass
+
+ return StateProcessingExecutor()
+
+ @pytest.fixture
+ def mock_binary_data_element(self, mock_binary_tuple, mock_data_input_channel):
+ return DataElement(
+ tag=mock_data_input_channel,
+ payload=DataFrame(
+ frame=pyarrow.Table.from_pandas(
+ pandas.DataFrame([mock_binary_tuple.as_dict()])
+ )
+ ),
+ )
+
+ @pytest.fixture
+ def mock_batch_data_elements(self, mock_batch, mock_data_input_channel):
+ data_elements = []
+ for i in range(57):
+ mock_tuple = Tuple({"test-1": "hello", "test-2": i})
+ data_elements.append(
+ DataElement(
+ tag=mock_data_input_channel,
+ payload=DataFrame(
+ frame=pyarrow.Table.from_pandas(
+ pandas.DataFrame([mock_tuple.as_dict()])
+ )
+ ),
+ )
+ )
+
+ return data_elements
+
+ @pytest.fixture
+ def mock_end_of_upstream(self, mock_tuple, mock_data_input_channel):
+ return ECMElement(
+ tag=mock_data_input_channel,
+ payload=EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity("EndChannel"),
+ EmbeddedControlMessageType.PORT_ALIGNMENT,
+ [],
+ {
+ mock_data_input_channel.to_worker_id.name: ControlInvocation(
+ "EndChannel",
+ ControlRequest(empty_request=EmptyRequest()),
+ AsyncRpcContext(ActorVirtualIdentity(), ActorVirtualIdentity()),
+ -1,
+ )
+ },
+ ),
+ )
+
+ @pytest.fixture
+ def input_queue(self):
+ return InternalQueue()
+
+ @pytest.fixture
+ def output_queue(self):
+ return InternalQueue()
+
+ @pytest.fixture
+ def mock_assign_input_port(
+ self, mock_raw_schema, mock_control_input_channel, mock_link, command_sequence
+ ):
+ command = set_one_of(
+ ControlRequest,
+ AssignPortRequest(
+ port_id=mock_link.to_port_id, input=True, schema=mock_raw_schema
+ ),
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="AssignPort", command_id=command_sequence, command=command
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_assign_output_port(
+ self, mock_raw_schema, mock_control_input_channel, command_sequence
+ ):
+ command = set_one_of(
+ ControlRequest,
+ AssignPortRequest(
+ port_id=PortIdentity(id=0), input=False, schema=mock_raw_schema
+ ),
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="AssignPort", command_id=command_sequence, command=command
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_assign_input_port_binary(
+ self,
+ mock_binary_raw_schema,
+ mock_control_input_channel,
+ mock_link,
+ command_sequence,
+ ):
+ command = set_one_of(
+ ControlRequest,
+ AssignPortRequest(
+ port_id=mock_link.to_port_id, input=True, schema=mock_binary_raw_schema
+ ),
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="AssignPort", command_id=command_sequence, command=command
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_assign_output_port_binary(
+ self, mock_binary_raw_schema, mock_control_input_channel, command_sequence
+ ):
+ command = set_one_of(
+ ControlRequest,
+ AssignPortRequest(
+ port_id=PortIdentity(id=0), input=False, schema=mock_binary_raw_schema
+ ),
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="AssignPort", command_id=command_sequence, command=command
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_add_input_channel(
+ self,
+ mock_control_input_channel,
+ mock_sender_actor,
+ mock_receiver_actor,
+ mock_link,
+ command_sequence,
+ ):
+ command = set_one_of(
+ ControlRequest,
+ AddInputChannelRequest(
+ ChannelIdentity(
+ from_worker_id=mock_sender_actor,
+ to_worker_id=mock_receiver_actor,
+ is_control=False,
+ ),
+ port_id=mock_link.to_port_id,
+ ),
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="AddInputChannel",
+ command_id=command_sequence,
+ command=command,
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_raw_schema(self):
+ return {"test-1": "STRING", "test-2": "INTEGER"}
+
+ @pytest.fixture
+ def mock_binary_raw_schema(self):
+ return {"test-1": "BINARY", "test-2": "INTEGER"}
+
+ @pytest.fixture
+ def mock_initialize_executor(
+ self,
+ mock_control_input_channel,
+ mock_sender_actor,
+ mock_link,
+ command_sequence,
+ mock_raw_schema,
+ ):
+ operator_code = "from pytexera import *\n" + inspect.getsource(EchoOperator)
+ command = set_one_of(
+ ControlRequest,
+ InitializeExecutorRequest(
+ op_exec_init_info=set_one_of(
+ OpExecInitInfo, OpExecWithCode(operator_code, "python")
+ ),
+ is_source=False,
+ ),
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="InitializeExecutor",
+ command_id=command_sequence,
+ command=command,
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_initialize_batch_count_executor(
+ self,
+ mock_control_input_channel,
+ mock_sender_actor,
+ mock_link,
+ command_sequence,
+ mock_raw_schema,
+ ):
+ operator_code = "from pytexera import *\n" + inspect.getsource(
+ CountBatchOperator
+ )
+ command = set_one_of(
+ ControlRequest,
+ InitializeExecutorRequest(
+ op_exec_init_info=set_one_of(
+ OpExecInitInfo, OpExecWithCode(operator_code, "python")
+ ),
+ is_source=False,
+ ),
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="InitializeExecutor",
+ command_id=command_sequence,
+ command=command,
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_add_partitioning(
+ self,
+ mock_control_input_channel,
+ mock_receiver_actor,
+ command_sequence,
+ mock_link,
+ ):
+ command = set_one_of(
+ ControlRequest,
+ AddPartitioningRequest(
+ tag=mock_link,
+ partitioning=set_one_of(
+ Partitioning,
+ OneToOnePartitioning(
+ batch_size=1,
+ channels=[
+ ChannelIdentity(
+ from_worker_id=ActorVirtualIdentity("dummy_worker_id"),
+ to_worker_id=mock_receiver_actor,
+ is_control=False,
+ )
+ ],
+ ),
+ ),
+ ),
+ )
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="AddPartitioning",
+ command_id=command_sequence,
+ command=command,
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_query_statistics(
+ self, mock_control_input_channel, mock_sender_actor, command_sequence
+ ):
+ command = set_one_of(ControlRequest, EmptyRequest())
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="QueryStatistics",
+ command_id=command_sequence,
+ command=command,
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_pause(
+ self, mock_control_input_channel, mock_sender_actor, command_sequence
+ ):
+ command = set_one_of(ControlRequest, EmptyRequest())
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="PauseWorker", command_id=command_sequence, command=command
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def mock_resume(
+ self, mock_control_input_channel, mock_sender_actor, command_sequence
+ ):
+ command = set_one_of(ControlRequest, EmptyRequest())
+ payload = set_one_of(
+ DirectControlMessagePayloadV2,
+ ControlInvocation(
+ method_name="ResumeWorker", command_id=command_sequence, command=command
+ ),
+ )
+ return DCMElement(tag=mock_control_input_channel, payload=payload)
+
+ @pytest.fixture
+ def main_loop(self, input_queue, output_queue, mock_link):
+ main_loop = MainLoop("dummy_worker_id", input_queue, output_queue)
+ yield main_loop
+ main_loop.stop()
+
+ @pytest.fixture
+ def main_loop_thread(self, main_loop, reraise):
+ def wrapper():
+ with reraise:
+ main_loop.run()
+
+ main_loop_thread = Thread(target=wrapper, name="main_loop_thread")
+ yield main_loop_thread
+
+ @staticmethod
+ def check_batch_rank_sum(
+ executor,
+ input_queue,
+ mock_batch_data_elements,
+ output_data_elements,
+ output_queue,
+ mock_batch,
+ start,
+ end,
+ count,
+ ):
+ # Checking the rank sum of each batch to make sure the accuracy
+ for i in range(start, end):
+ input_queue.put(mock_batch_data_elements[i])
+ rank_sum_real = 0
+ rank_sum_suppose = 0
+ for i in range(start, end):
+ output_data_elements.append(output_queue.get())
+ rank_sum_real += output_data_elements[i].payload.frame[0]["test-2"]
+ rank_sum_suppose += mock_batch[i]["test-2"]
+ assert executor.count == count
+ assert rank_sum_real == rank_sum_suppose
+
+ @pytest.mark.timeout(2)
+ def test_main_loop_thread_can_start(self, main_loop_thread):
+ main_loop_thread.start()
+ assert main_loop_thread.is_alive()
+
+ @pytest.mark.timeout(2)
+ def test_main_loop_thread_can_process_messages(
+ self,
+ mock_link,
+ mock_data_input_channel,
+ mock_data_output_channel,
+ mock_control_input_channel,
+ mock_control_output_channel,
+ input_queue,
+ output_queue,
+ mock_data_element,
+ main_loop_thread,
+ mock_assign_input_port,
+ mock_assign_output_port,
+ mock_add_input_channel,
+ mock_add_partitioning,
+ mock_initialize_executor,
+ mock_end_of_upstream,
+ mock_query_statistics,
+ mock_tuple,
+ command_sequence,
+ reraise,
+ ):
+ main_loop_thread.start()
+
+ # can process AssignPort
+ input_queue.put(mock_assign_input_port)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+ input_queue.put(mock_assign_output_port)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process AddInputChannel
+ input_queue.put(mock_add_input_channel)
+
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process AddPartitioning
+ input_queue.put(mock_add_partitioning)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process InitializeExecutor
+ input_queue.put(mock_initialize_executor)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process a DataFrame
+ input_queue.put(mock_data_element)
+
+ output_data_element: DataElement = output_queue.get()
+ assert output_data_element.tag == mock_data_output_channel
+ assert isinstance(output_data_element.payload, DataFrame)
+ data_frame: DataFrame = output_data_element.payload
+ assert len(data_frame.frame) == 1
+ assert Tuple(data_frame.frame.to_pylist()[0]) == mock_tuple
+
+ # can process QueryStatistics
+ input_queue.put(mock_query_statistics)
+ elem = output_queue.get()
+ stats_invocation = elem.payload.return_invocation
+ worker_metrics_response = stats_invocation.return_value.worker_metrics_response
+ stats = worker_metrics_response.metrics.worker_statistics
+
+ metrics = WorkerMetrics(
+ worker_state=WorkerState.RUNNING,
+ worker_statistics=WorkerStatistics(
+ input_tuple_metrics=[
+ PortTupleMetricsMapping(
+ PortIdentity(0),
+ TupleMetrics(
+ 1,
+ stats.input_tuple_metrics[0].tuple_metrics.size,
+ ),
+ )
+ ],
+ output_tuple_metrics=[
+ PortTupleMetricsMapping(
+ PortIdentity(0),
+ TupleMetrics(
+ 1,
+ stats.output_tuple_metrics[0].tuple_metrics.size,
+ ),
+ )
+ ],
+ data_processing_time=stats.data_processing_time,
+ control_processing_time=stats.control_processing_time,
+ idle_time=stats.idle_time,
+ ),
+ )
+
+ assert elem == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=1,
+ return_value=ControlReturn(
+ worker_metrics_response=WorkerMetricsResponse(metrics=metrics),
+ ),
+ ),
+ ),
+ )
+
+ input_queue.put(mock_end_of_upstream)
+ output_queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE)
+ # the input port should complete
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ control_invocation=ControlInvocation(
+ method_name="PortCompleted",
+ command_id=0,
+ context=AsyncRpcContext(
+ sender=ActorVirtualIdentity(name="dummy_worker_id"),
+ receiver=ActorVirtualIdentity(name="CONTROLLER"),
+ ),
+ command=ControlRequest(
+ port_completed_request=PortCompletedRequest(
+ port_id=mock_link.to_port_id, input=True
+ )
+ ),
+ )
+ ),
+ )
+
+ # the output port should complete
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ control_invocation=ControlInvocation(
+ method_name="PortCompleted",
+ command_id=1,
+ context=AsyncRpcContext(
+ sender=ActorVirtualIdentity(name="dummy_worker_id"),
+ receiver=ActorVirtualIdentity(name="CONTROLLER"),
+ ),
+ command=ControlRequest(
+ port_completed_request=PortCompletedRequest(
+ port_id=PortIdentity(id=0), input=False
+ )
+ ),
+ )
+ ),
+ )
+
+ # WorkerExecutionCompletedV2 should be triggered when workflow finishes
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ control_invocation=ControlInvocation(
+ method_name="WorkerExecutionCompleted",
+ command_id=2,
+ context=AsyncRpcContext(
+ sender=ActorVirtualIdentity(name="dummy_worker_id"),
+ receiver=ActorVirtualIdentity(name="CONTROLLER"),
+ ),
+ command=ControlRequest(empty_request=EmptyRequest()),
+ )
+ ),
+ )
+
+ output_queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE)
+ assert output_queue.get() == ECMElement(
+ tag=mock_data_output_channel,
+ payload=EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity("EndChannel"),
+ EmbeddedControlMessageType.PORT_ALIGNMENT,
+ [],
+ {
+ mock_data_output_channel.to_worker_id.name: ControlInvocation(
+ "EndChannel",
+ ControlRequest(empty_request=EmptyRequest()),
+ AsyncRpcContext(ActorVirtualIdentity(), ActorVirtualIdentity()),
+ -1,
+ )
+ },
+ ),
+ )
+
+ # can process ReturnInvocation
+ input_queue.put(
+ DCMElement(
+ tag=mock_control_input_channel,
+ payload=set_one_of(
+ DirectControlMessagePayloadV2,
+ ReturnInvocation(
+ command_id=0,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ ),
+ ),
+ )
+ )
+
+ reraise()
+
+ @pytest.mark.timeout(5)
+ def test_batch_dp_thread_can_process_batch(
+ self,
+ mock_control_input_channel,
+ mock_control_output_channel,
+ mock_data_input_channel,
+ mock_data_output_channel,
+ mock_link,
+ input_queue,
+ output_queue,
+ mock_receiver_actor,
+ main_loop,
+ main_loop_thread,
+ mock_query_statistics,
+ mock_assign_input_port,
+ mock_assign_output_port,
+ mock_add_input_channel,
+ mock_add_partitioning,
+ mock_pause,
+ mock_resume,
+ mock_initialize_batch_count_executor,
+ mock_batch,
+ mock_batch_data_elements,
+ mock_end_of_upstream,
+ command_sequence,
+ reraise,
+ ):
+ main_loop_thread.start()
+
+ # can process AssignPort
+ input_queue.put(mock_assign_input_port)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+ input_queue.put(mock_assign_output_port)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process AddInputChannel
+ input_queue.put(mock_add_input_channel)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process AddPartitioning
+ input_queue.put(mock_add_partitioning)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process InitializeExecutor
+ input_queue.put(mock_initialize_batch_count_executor)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+ executor = main_loop.context.executor_manager.executor
+ output_data_elements = []
+
+ # can process a DataFrame
+ executor.BATCH_SIZE = 10
+ for i in range(13):
+ input_queue.put(mock_batch_data_elements[i])
+ for i in range(10):
+ output_data_elements.append(output_queue.get())
+
+ self.send_pause(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_pause,
+ output_queue,
+ )
+ # input queue 13, output queue 10, batch_buffer 3
+ assert executor.count == 1
+ executor.BATCH_SIZE = 20
+ self.send_resume(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_resume,
+ output_queue,
+ )
+
+ for i in range(13, 41):
+ input_queue.put(mock_batch_data_elements[i])
+ for i in range(20):
+ output_data_elements.append(output_queue.get())
+
+ self.send_pause(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_pause,
+ output_queue,
+ )
+ # input queue 41, output queue 30, batch_buffer 11
+ assert executor.count == 2
+ executor.BATCH_SIZE = 5
+ self.send_resume(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_resume,
+ output_queue,
+ )
+
+ input_queue.put(mock_batch_data_elements[41])
+ input_queue.put(mock_batch_data_elements[42])
+ for i in range(10):
+ output_data_elements.append(output_queue.get())
+
+ self.send_pause(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_pause,
+ output_queue,
+ )
+ # input queue 43, output queue 40, batch_buffer 3
+ assert executor.count == 4
+ self.send_resume(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_resume,
+ output_queue,
+ )
+
+ for i in range(43, 57):
+ input_queue.put(mock_batch_data_elements[i])
+ for i in range(15):
+ output_data_elements.append(output_queue.get())
+
+ self.send_pause(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_pause,
+ output_queue,
+ )
+ # input queue 57, output queue 55, batch_buffer 2
+ assert executor.count == 7
+ self.send_resume(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_resume,
+ output_queue,
+ )
+
+ input_queue.put(mock_end_of_upstream)
+ for i in range(2):
+ output_data_elements.append(output_queue.get())
+
+ # check the batch count
+ assert main_loop.context.executor_manager.executor.count == 8
+
+ assert output_data_elements[0].tag == mock_data_output_channel
+ assert isinstance(output_data_elements[0].payload, DataFrame)
+ data_frame: DataFrame = output_data_elements[0].payload
+ assert len(data_frame.frame) == 1
+ assert Tuple(data_frame.frame.to_pylist()[0]) == Tuple(mock_batch[0])
+
+ reraise()
+
+ @pytest.mark.timeout(5)
+ def test_main_loop_thread_can_process_single_tuple_with_binary(
+ self,
+ mock_link,
+ mock_data_input_channel,
+ mock_data_output_channel,
+ mock_control_output_channel,
+ mock_control_input_channel,
+ input_queue,
+ output_queue,
+ mock_binary_tuple,
+ mock_binary_data_element,
+ main_loop_thread,
+ mock_assign_input_port_binary,
+ mock_assign_output_port_binary,
+ mock_add_input_channel,
+ mock_add_partitioning,
+ mock_initialize_executor,
+ mock_end_of_upstream,
+ mock_query_statistics,
+ command_sequence,
+ reraise,
+ ):
+ main_loop_thread.start()
+
+ # can process AssignPort
+ input_queue.put(mock_assign_input_port_binary)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+ input_queue.put(mock_assign_output_port_binary)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process AddInputChannel
+ input_queue.put(mock_add_input_channel)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process AddPartitioning
+ input_queue.put(mock_add_partitioning)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process InitializeExecutor
+ input_queue.put(mock_initialize_executor)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ input_queue.put(mock_binary_data_element)
+ output_data_element: DataElement = output_queue.get()
+ assert output_data_element.tag == mock_data_output_channel
+ assert isinstance(output_data_element.payload, DataFrame)
+ data_frame: DataFrame = output_data_element.payload
+
+ assert len(data_frame.frame) == 1
+ assert data_frame.frame.to_pylist()[0][
+ "test-1"
+ ] == b"pickle " + pickle.dumps(mock_binary_tuple["test-1"])
+
+ reraise()
+
+ @staticmethod
+ def send_pause(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_pause,
+ output_queue,
+ ):
+ input_queue.put(mock_pause)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(
+ worker_state_response=WorkerStateResponse(WorkerState.PAUSED)
+ ),
+ )
+ ),
+ )
+
+ @staticmethod
+ def send_resume(
+ command_sequence,
+ input_queue,
+ mock_control_output_channel,
+ mock_resume,
+ output_queue,
+ ):
+ input_queue.put(mock_resume)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(
+ worker_state_response=WorkerStateResponse(WorkerState.RUNNING)
+ ),
+ )
+ ),
+ )
+
+ @pytest.mark.timeout(2)
+ def test_process_state_can_emit_consecutive_states(
+ self,
+ main_loop,
+ output_queue,
+ mock_data_output_channel,
+ monkeypatch,
+ ):
+ class DummyExecutor:
+ @staticmethod
+ def process_state(state, port: int):
+ return State({"value": state["value"] + 1, "port": port})
+
+ main_loop.context.executor_manager.executor = DummyExecutor()
+ monkeypatch.setattr(main_loop, "_check_and_process_control", lambda: None)
+ monkeypatch.setattr(
+ main_loop.context.output_manager,
+ "emit_state",
+ lambda state: [(mock_data_output_channel.to_worker_id, StateFrame(state))],
+ )
+
+ def fake_switch_context():
+ current_input_state = (
+ main_loop.context.state_processing_manager.current_input_state
+ )
+ if current_input_state is not None:
+ main_loop.context.state_processing_manager.current_output_state = (
+ DummyExecutor.process_state(current_input_state, 0)
+ )
+
+ monkeypatch.setattr(main_loop, "_switch_context", fake_switch_context)
+
+ first_state = State({"value": 1})
+ second_state = State({"value": 41})
+
+ main_loop._process_state(first_state)
+ main_loop._process_state(second_state)
+
+ first_output: DataElement = output_queue.get()
+ second_output: DataElement = output_queue.get()
+
+ assert first_output.tag == mock_data_output_channel
+ assert isinstance(first_output.payload, StateFrame)
+ assert first_output.payload.frame["value"] == 2
+ assert first_output.payload.frame["port"] == 0
+
+ assert second_output.tag == mock_data_output_channel
+ assert isinstance(second_output.payload, StateFrame)
+ assert second_output.payload.frame["value"] == 42
+ assert second_output.payload.frame["port"] == 0
+
+ @pytest.mark.timeout(5)
+ def test_main_loop_thread_can_align_ecm(
+ self,
+ mock_link,
+ mock_data_input_channel,
+ mock_data_output_channel,
+ mock_control_output_channel,
+ mock_control_input_channel,
+ input_queue,
+ output_queue,
+ mock_binary_tuple,
+ mock_binary_data_element,
+ main_loop_thread,
+ mock_assign_input_port_binary,
+ mock_assign_output_port_binary,
+ mock_add_input_channel,
+ mock_add_partitioning,
+ mock_initialize_executor,
+ mock_end_of_upstream,
+ mock_query_statistics,
+ command_sequence,
+ reraise,
+ ):
+ main_loop_thread.start()
+
+ # can process AssignPort
+ input_queue.put(mock_assign_input_port_binary)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+ input_queue.put(mock_assign_output_port_binary)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process AddInputChannel
+ input_queue.put(mock_add_input_channel)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process AddPartitioning
+ input_queue.put(mock_add_partitioning)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # can process InitializeExecutor
+ input_queue.put(mock_initialize_executor)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ scope = [mock_control_input_channel, mock_data_input_channel]
+ command_mapping = {
+ mock_control_input_channel.to_worker_id.name: ControlInvocation(
+ "NoOperation", EmptyRequest(), AsyncRpcContext(), 98
+ )
+ }
+ test_ecm = EmbeddedControlMessage(
+ "test_ecm", EmbeddedControlMessageType.ALL_ALIGNMENT, scope, command_mapping
+ )
+ input_queue.put(ECMElement(tag=mock_control_input_channel, payload=test_ecm))
+ input_queue.put(mock_binary_data_element)
+ input_queue.put(ECMElement(tag=mock_data_input_channel, payload=test_ecm))
+
+ # The two outputs land on different channel sub-queues:
+ # - DataElement on the data channel to the downstream worker
+ # - DCMElement (NoOperation reply) on the control channel back to "sender"
+ # output_queue is a priority multi-queue. With both items present,
+ # the control sub-queue (priority 1) outranks the data sub-queue
+ # (priority 2), so the control reply must come out first. Wait for
+ # both channels to have their item before popping, so the priority
+ # guarantee is what we're actually testing — see #4524.
+ control_reply_channel = ChannelIdentity(
+ ActorVirtualIdentity("dummy_worker_id"),
+ ActorVirtualIdentity("sender"),
+ is_control=True,
+ )
+
+ def channel_size(channel: ChannelIdentity) -> int:
+ # Sub-queues are added lazily on first put, so the channel may not
+ # exist in the LBMQ yet. Treat that as size zero.
+ if channel not in output_queue._queue.sub_queues:
+ return 0
+ return output_queue._queue.size(channel)
+
+ deadline = time.time() + 5.0
+ while channel_size(mock_data_output_channel) == 0 or (
+ channel_size(control_reply_channel) == 0
+ ):
+ if time.time() > deadline:
+ raise AssertionError(
+ f"timed out waiting for outputs on both channels; "
+ f"data={channel_size(mock_data_output_channel)}, "
+ f"control={channel_size(control_reply_channel)}"
+ )
+ time.sleep(0.001)
+
+ # Priority pulls control before data when both are queued.
+ output_control_element = output_queue.get()
+ assert isinstance(output_control_element, DCMElement), (
+ f"expected control reply first (priority), got {type(output_control_element).__name__}"
+ )
+ assert output_control_element.tag == control_reply_channel
+ assert output_control_element.payload.return_invocation.command_id == 98
+ assert (
+ output_control_element.payload.return_invocation.return_value
+ == ControlReturn(empty_return=EmptyReturn())
+ )
+
+ output_data_element = output_queue.get()
+ assert isinstance(output_data_element, DataElement), (
+ f"expected data element second, got {type(output_data_element).__name__}"
+ )
+ assert output_data_element.tag == mock_data_output_channel
+ assert isinstance(output_data_element.payload, DataFrame)
+ data_frame: DataFrame = output_data_element.payload
+ assert len(data_frame.frame) == 1
+ assert data_frame.frame.to_pylist()[0][
+ "test-1"
+ ] == b"pickle " + pickle.dumps(mock_binary_tuple["test-1"])
+ reraise()
+
+ @pytest.mark.timeout(2)
+ def test_process_state_can_emit_multiple_states(
+ self,
+ main_loop,
+ output_queue,
+ mock_data_output_channel,
+ monkeypatch,
+ ):
+ # Stub-level coverage of the single-switch state handshake. Each
+ # call to the (stubbed) _switch_context simulates DataProc
+ # consuming the queued input state and writing
+ # current_output_state, mirroring what real DataProc.process_state
+ # does between MainLoop's switches.
+ class DummyExecutor:
+ @staticmethod
+ def process_state(state: State, port: int) -> State:
+ return State({"value": state["value"] + 1, "port": port})
+
+ main_loop.context.executor_manager.executor = DummyExecutor()
+ monkeypatch.setattr(main_loop, "_check_and_process_control", lambda: None)
+ monkeypatch.setattr(
+ main_loop.context.output_manager,
+ "emit_state",
+ lambda state: [(mock_data_output_channel.to_worker_id, StateFrame(state))],
+ )
+
+ def fake_switch_context():
+ current_input_state = (
+ main_loop.context.state_processing_manager.current_input_state
+ )
+ if current_input_state is not None:
+ main_loop.context.state_processing_manager.current_output_state = (
+ DummyExecutor.process_state(current_input_state, 0)
+ )
+
+ monkeypatch.setattr(main_loop, "_switch_context", fake_switch_context)
+
+ first_state = State({"value": 1})
+ second_state = State({"value": 41})
+
+ main_loop._process_state(first_state)
+ main_loop._process_state(second_state)
+
+ first_output: DataElement = output_queue.get()
+ second_output: DataElement = output_queue.get()
+
+ assert first_output.tag == mock_data_output_channel
+ assert isinstance(first_output.payload, StateFrame)
+ assert first_output.payload.frame["value"] == 2
+ assert first_output.payload.frame["port"] == 0
+
+ assert second_output.tag == mock_data_output_channel
+ assert isinstance(second_output.payload, StateFrame)
+ assert second_output.payload.frame["value"] == 42
+ assert second_output.payload.frame["port"] == 0
+
+ @pytest.mark.timeout(2)
+ def test_main_loop_thread_can_process_state(
+ self,
+ mock_data_output_channel,
+ mock_control_output_channel,
+ input_queue,
+ output_queue,
+ main_loop,
+ main_loop_thread,
+ mock_assign_input_port,
+ mock_assign_output_port,
+ mock_add_input_channel,
+ mock_add_partitioning,
+ mock_initialize_executor,
+ mock_state_data_elements,
+ mock_end_of_upstream,
+ state_processing_executor,
+ command_sequence,
+ reraise,
+ ):
+ # End-to-end coverage of the state-processing path through the real
+ # MainLoop + DataProcessor threads. The single-switch state handshake
+ # in MainLoop.process_input_state means each state is emitted in its
+ # own cycle (no lag), and an EndChannel ECM after the last state
+ # produces an additional output via produce_state_on_finish.
+ main_loop_thread.start()
+
+ for setup_msg in [
+ mock_assign_input_port,
+ mock_assign_output_port,
+ mock_add_input_channel,
+ mock_add_partitioning,
+ mock_initialize_executor,
+ ]:
+ input_queue.put(setup_msg)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ # Going through the InitializeExecutor RPC above sets up the rest of
+ # the worker state (output schema, partitioning bookkeeping). Swap
+ # the executor instance with the test helper here so the test can
+ # assert the executor's process_state and produce_state_on_finish
+ # actually ran, without depending on Python's cross-test module
+ # caching for operator classes loaded via OpExecWithCode.
+ main_loop.context.executor_manager.executor = state_processing_executor
+
+ # Send four states. With the lag-free state pipeline we expect each
+ # state to produce its own output in order.
+ for state_element in mock_state_data_elements:
+ input_queue.put(state_element)
+
+ for expected_value in (1, 2, 3, 4):
+ output_data_element: DataElement = output_queue.get()
+ assert output_data_element.tag == mock_data_output_channel
+ assert isinstance(output_data_element.payload, StateFrame), (
+ f"expected StateFrame for value={expected_value}, got "
+ f"{type(output_data_element.payload).__name__}"
+ )
+ output_state = output_data_element.payload.frame
+ assert output_state["value"] == expected_value, (
+ f"state outputs arrived out of order: expected value="
+ f"{expected_value}, got value={output_state['value']}"
+ )
+ assert output_state["processed_marker"] == "executed"
+ assert output_state["port"] == 0
+
+ # Send EndChannel to drive _process_end_channel. The executor's
+ # produce_state_on_finish writes a finish-marker state into
+ # current_output_state inside DataProc's process_internal_marker;
+ # MainLoop's process_input_state then emits it.
+ input_queue.put(mock_end_of_upstream)
+
+ # Drain the control reply messages so the next data
+ # output_queue.get() returns the post-EndChannel data emission.
+ output_queue.disable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE)
+ for _ in range(3):
+ control_reply = output_queue.get()
+ assert isinstance(control_reply, DCMElement), (
+ f"expected DCMElement during EndChannel teardown, got "
+ f"{type(control_reply).__name__}"
+ )
+ output_queue.enable_data(InternalQueue.DisableType.DISABLE_BY_PAUSE)
+
+ end_channel_state_output: DataElement = output_queue.get()
+ assert end_channel_state_output.tag == mock_data_output_channel
+ assert isinstance(end_channel_state_output.payload, StateFrame), (
+ f"expected StateFrame for the EndChannel-driven emission, got "
+ f"{type(end_channel_state_output.payload).__name__}"
+ )
+ end_channel_state = end_channel_state_output.payload.frame
+ assert "finish_marker" in end_channel_state, (
+ f"EndChannel emission should be the finish-marker state from "
+ f"produce_state_on_finish, got {end_channel_state!r}"
+ )
+ assert end_channel_state["finish_marker"] == "produce_state_on_finish_ran"
+
+ reraise()
+
+ @pytest.mark.timeout(2)
+ def test_main_loop_thread_can_process_state_after_tuple(
+ self,
+ mock_data_output_channel,
+ mock_control_output_channel,
+ input_queue,
+ output_queue,
+ main_loop,
+ main_loop_thread,
+ mock_assign_input_port,
+ mock_assign_output_port,
+ mock_add_input_channel,
+ mock_add_partitioning,
+ mock_initialize_executor,
+ mock_data_element,
+ mock_state_data_elements,
+ state_processing_executor,
+ command_sequence,
+ reraise,
+ ):
+ # Coverage for the mixed (tuple, then state) input sequence: a
+ # tuple followed by several state DataElements should still emit
+ # every state's processed output in order.
+ main_loop_thread.start()
+
+ for setup_msg in [
+ mock_assign_input_port,
+ mock_assign_output_port,
+ mock_add_input_channel,
+ mock_add_partitioning,
+ mock_initialize_executor,
+ ]:
+ input_queue.put(setup_msg)
+ assert output_queue.get() == DCMElement(
+ tag=mock_control_output_channel,
+ payload=DirectControlMessagePayloadV2(
+ return_invocation=ReturnInvocation(
+ command_id=command_sequence,
+ return_value=ControlReturn(empty_return=EmptyReturn()),
+ )
+ ),
+ )
+
+ main_loop.context.executor_manager.executor = state_processing_executor
+
+ # Tuple first, then four states.
+ input_queue.put(mock_data_element)
+ warmup_output: DataElement = output_queue.get()
+ assert warmup_output.tag == mock_data_output_channel
+ assert isinstance(warmup_output.payload, DataFrame)
+
+ for state_element in mock_state_data_elements:
+ input_queue.put(state_element)
+
+ for expected_value in (1, 2, 3, 4):
+ output_data_element: DataElement = output_queue.get()
+ assert output_data_element.tag == mock_data_output_channel
+ assert isinstance(output_data_element.payload, StateFrame), (
+ f"expected StateFrame for value={expected_value}, got "
+ f"{type(output_data_element.payload).__name__}"
+ )
+ output_state = output_data_element.payload.frame
+ assert output_state["value"] == expected_value, (
+ f"state outputs after a tuple arrived out of order: "
+ f"expected value={expected_value}, "
+ f"got value={output_state['value']}"
+ )
+ assert output_state["processed_marker"] == "executed"
+
+ reraise()
diff --git a/amber/src/main/python/core/runnables/test_network_receiver.py b/amber/src/main/python/core/runnables/test_network_receiver.py
new file mode 100644
index 00000000000..bf890e4a2f0
--- /dev/null
+++ b/amber/src/main/python/core/runnables/test_network_receiver.py
@@ -0,0 +1,236 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+import threading
+from pyarrow import Table
+
+from core.models.internal_queue import (
+ InternalQueue,
+ DCMElement,
+ DataElement,
+ ECMElement,
+)
+from core.models.payload import DataFrame, StateFrame
+from core.models.state import State
+from core.proxy import ProxyClient
+from core.runnables.network_receiver import NetworkReceiver
+from core.runnables.network_sender import NetworkSender
+from core.util.proto import set_one_of
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ControlInvocation,
+ EmbeddedControlMessage,
+ EmbeddedControlMessageType,
+ EmptyRequest,
+ AsyncRpcContext,
+ ControlRequest,
+)
+from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2
+
+
+class TestNetworkReceiver:
+ @pytest.fixture
+ def input_queue(self):
+ return InternalQueue()
+
+ @pytest.fixture
+ def output_queue(self):
+ return InternalQueue()
+
+ @pytest.fixture
+ def network_receiver(self, output_queue):
+ network_receiver = NetworkReceiver(output_queue, host="localhost", port=5555)
+ yield network_receiver
+ network_receiver.stop()
+
+ class MockFlightMetadataReader:
+ """
+ MockFlightMetadataReader is a mocked FlightMetadataReader class to ultimately
+ mock a credit value to be returned from Scala server to Python client
+ """
+
+ class MockBuffer:
+ def to_pybytes(self):
+ dummy_credit = 31
+ return dummy_credit.to_bytes(8, "little")
+
+ def read(self):
+ return self.MockBuffer()
+
+ @pytest.fixture
+ def network_sender_thread(self, input_queue):
+ network_sender = NetworkSender(input_queue, host="localhost", port=5555)
+
+ # mocking do_put, read, to_pybytes to return fake credit values
+ def mock_do_put(
+ self,
+ FlightDescriptor_descriptor,
+ Schema_schema,
+ FlightCallOptions_options=None,
+ ):
+ """
+ Mocking FlightClient.do_put that is called in ProxyClient to return
+ a MockFlightMetadataReader instead of a FlightMetadataReader
+
+ :param self: an instance of FlightClient (would be ProxyClient in this case)
+ :param FlightDescriptor_descriptor: descriptor
+ :param Schema_schema: schema
+ :param FlightCallOptions_options: options, None by default
+ :return: writer : FlightStreamWriter, reader : MockFlightMetadataReader
+ """
+ writer, _ = super(ProxyClient, self).do_put(
+ FlightDescriptor_descriptor, Schema_schema, FlightCallOptions_options
+ )
+ reader = TestNetworkReceiver.MockFlightMetadataReader()
+ return writer, reader
+
+ mock_proxy_client = network_sender._proxy_client
+ mock_proxy_client.do_put = mock_do_put.__get__(
+ mock_proxy_client, ProxyClient
+ ) # override do_put with mock_do_put
+
+ network_sender_thread = threading.Thread(target=network_sender.run)
+ yield network_sender_thread
+ network_sender.stop()
+
+ @pytest.fixture
+ def data_payload(self):
+ return DataFrame(
+ frame=Table.from_pydict(
+ {
+ "Brand": ["Honda Civic", "Toyota Corolla", "Ford Focus", "Audi A4"],
+ "Price": [22000, 25000, 27000, 35000],
+ }
+ )
+ )
+
+ @pytest.mark.timeout(10)
+ def test_network_receiver_can_receive_data_messages(
+ self,
+ data_payload,
+ output_queue,
+ input_queue,
+ network_receiver,
+ network_sender_thread,
+ ):
+ network_sender_thread.start()
+ worker_id = ActorVirtualIdentity(name="test")
+ channel_id = ChannelIdentity(worker_id, worker_id, False)
+ input_queue.put(DataElement(tag=channel_id, payload=data_payload))
+ element: DataElement = output_queue.get()
+ assert len(element.payload.frame) == len(data_payload.frame)
+ assert element.tag == channel_id
+
+ @pytest.mark.timeout(10)
+ def test_network_receiver_can_receive_consecutive_state_messages(
+ self,
+ output_queue,
+ input_queue,
+ network_receiver,
+ network_sender_thread,
+ ):
+ network_sender_thread.start()
+ worker_id = ActorVirtualIdentity(name="test")
+ channel_id = ChannelIdentity(worker_id, worker_id, False)
+
+ input_queue.put(
+ DataElement(
+ tag=channel_id,
+ payload=StateFrame(State({"loop_counter": 0, "i": 1})),
+ )
+ )
+ input_queue.put(
+ DataElement(
+ tag=channel_id,
+ payload=StateFrame(State({"loop_counter": 1, "i": 2})),
+ )
+ )
+
+ first_element: DataElement = output_queue.get()
+ second_element: DataElement = output_queue.get()
+
+ assert isinstance(first_element.payload, StateFrame)
+ assert first_element.payload.frame == {"loop_counter": 0, "i": 1}
+ assert first_element.tag == channel_id
+
+ assert isinstance(second_element.payload, StateFrame)
+ assert second_element.payload.frame == {"loop_counter": 1, "i": 2}
+ assert second_element.tag == channel_id
+
+ @pytest.mark.timeout(10)
+ def test_network_receiver_can_receive_control_messages(
+ self,
+ data_payload,
+ output_queue,
+ input_queue,
+ network_receiver,
+ network_sender_thread,
+ ):
+ worker_id = ActorVirtualIdentity(name="test")
+ control_payload = set_one_of(DirectControlMessagePayloadV2, ControlInvocation())
+ channel_id = ChannelIdentity(worker_id, worker_id, False)
+ input_queue.put(DCMElement(tag=channel_id, payload=control_payload))
+ network_sender_thread.start()
+ element: DCMElement = output_queue.get()
+ assert element.payload == control_payload
+ assert element.tag == channel_id
+
+ @pytest.mark.timeout(10)
+ def test_network_receiver_can_receive_ecm(
+ self,
+ output_queue,
+ input_queue,
+ network_receiver,
+ network_sender_thread,
+ ):
+ network_sender_thread.start()
+ worker_id = ActorVirtualIdentity(name="test")
+ channel_id = ChannelIdentity(worker_id, worker_id, False)
+ ecm_id = EmbeddedControlMessageIdentity("test_ecm")
+ scope = [channel_id]
+ rpc_context = AsyncRpcContext(worker_id, worker_id)
+ command_mapping = {
+ str(worker_id): ControlInvocation(
+ "NoOperation",
+ ControlRequest(empty_request=EmptyRequest()),
+ rpc_context,
+ 12,
+ )
+ }
+ input_queue.put(
+ ECMElement(
+ tag=channel_id,
+ payload=EmbeddedControlMessage(
+ ecm_id,
+ EmbeddedControlMessageType.ALL_ALIGNMENT,
+ scope,
+ command_mapping,
+ ),
+ )
+ )
+ element: DataElement = output_queue.get()
+ assert isinstance(element.payload, EmbeddedControlMessage)
+ assert element.payload.ecm_type == EmbeddedControlMessageType.ALL_ALIGNMENT
+ assert element.payload.id == ecm_id
+ assert element.payload.command_mapping == command_mapping
+ assert element.payload.scope == scope
+ assert element.tag == channel_id
diff --git a/amber/src/main/python/core/runnables/test_network_sender.py b/amber/src/main/python/core/runnables/test_network_sender.py
new file mode 100644
index 00000000000..529cd19d33e
--- /dev/null
+++ b/amber/src/main/python/core/runnables/test_network_sender.py
@@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+import threading
+from time import sleep
+
+from core.models.internal_queue import InternalQueue
+from core.runnables.network_receiver import NetworkReceiver
+from core.runnables.network_sender import NetworkSender
+
+
+class TestNetworkSender:
+ @pytest.fixture
+ def network_receiver(self):
+ network_receiver = NetworkReceiver(InternalQueue(), host="localhost", port=5555)
+ yield network_receiver
+ network_receiver.stop()
+
+ @pytest.fixture
+ def network_receiver_thread(self, network_receiver):
+ network_receiver_thread = threading.Thread(target=network_receiver.run)
+ yield network_receiver_thread
+
+ @pytest.fixture
+ def network_sender(self):
+ network_sender = NetworkSender(InternalQueue(), host="localhost", port=5555)
+ yield network_sender
+ network_sender.stop()
+
+ @pytest.fixture
+ def network_sender_thread(self, network_sender):
+ network_sender_thread = threading.Thread(target=network_sender.run)
+ yield network_sender_thread
+
+ @pytest.mark.timeout(2)
+ def test_network_sender_can_stop(
+ self,
+ network_receiver,
+ network_receiver_thread,
+ network_sender,
+ network_sender_thread,
+ ):
+ network_receiver_thread.start()
+ network_sender_thread.start()
+ assert network_receiver_thread.is_alive()
+ assert network_sender_thread.is_alive()
+ sleep(0.1)
+ network_receiver.stop()
+ network_sender.stop()
+ sleep(0.1)
+ assert not network_receiver_thread.is_alive()
+ assert not network_sender_thread.is_alive()
+ network_receiver_thread.join()
+ network_sender_thread.join()
diff --git a/amber/src/main/python/core/storage/__init__.py b/amber/src/main/python/core/storage/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/storage/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/storage/document_factory.py b/amber/src/main/python/core/storage/document_factory.py
new file mode 100644
index 00000000000..9b686ab66b6
--- /dev/null
+++ b/amber/src/main/python/core/storage/document_factory.py
@@ -0,0 +1,126 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import typing
+import urllib
+from typing import Optional
+from urllib.parse import urlparse
+
+from core.models import Schema, Tuple
+from core.storage.iceberg.iceberg_catalog_instance import IcebergCatalogInstance
+from core.storage.iceberg.iceberg_document import IcebergDocument
+from core.storage.iceberg.iceberg_utils import (
+ create_table,
+ amber_schema_to_iceberg_schema,
+ amber_tuples_to_arrow_table,
+ arrow_table_to_amber_tuples,
+ load_table_metadata,
+)
+from core.storage.model.virtual_document import VirtualDocument
+from core.storage.storage_config import StorageConfig
+from core.storage.vfs_uri_factory import VFSURIFactory, VFSResourceType
+
+
+class DocumentFactory:
+ """
+ Factory class to create and open documents.
+ Currently only iceberg documents are supported.
+ """
+
+ ICEBERG = "iceberg"
+
+ @staticmethod
+ def sanitize_uri_path(uri):
+ """
+ Matches the same implementation in our Scala codebase.
+ urllib.parse.urlparse does not automatically unquote the URI, while
+ java.net.URI.getPath does. Hence we need to explicitly
+ unquote to decode percent-encoded characters, then sanitize.
+ :param uri: Result of urllib.parse.urlparse(). Could be quoted.
+ :return: Unquoted and sanitized format of uri.
+ """
+ return urllib.parse.unquote(uri.path).lstrip("/").replace("/", "_")
+
+ @staticmethod
+ def create_document(uri: str, schema: Schema) -> VirtualDocument:
+ parsed_uri = urlparse(uri)
+ if parsed_uri.scheme == VFSURIFactory.VFS_FILE_URI_SCHEME:
+ _, _, _, resource_type = VFSURIFactory.decode_uri(uri)
+
+ if resource_type in {VFSResourceType.RESULT}:
+ storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
+
+ # Convert Amber Schema to Iceberg Schema with LARGE_BINARY
+ # field name encoding
+ iceberg_schema = amber_schema_to_iceberg_schema(schema)
+
+ create_table(
+ IcebergCatalogInstance.get_instance(),
+ StorageConfig.ICEBERG_TABLE_RESULT_NAMESPACE,
+ storage_key,
+ iceberg_schema,
+ override_if_exists=True,
+ )
+
+ return IcebergDocument[Tuple](
+ StorageConfig.ICEBERG_TABLE_RESULT_NAMESPACE,
+ storage_key,
+ iceberg_schema,
+ amber_tuples_to_arrow_table,
+ arrow_table_to_amber_tuples,
+ )
+ else:
+ raise ValueError(f"Resource type {resource_type} is not supported")
+ else:
+ raise NotImplementedError(
+ f"Unsupported URI scheme: {parsed_uri.scheme} for creating the document"
+ )
+
+ @staticmethod
+ def open_document(uri: str) -> typing.Tuple[VirtualDocument, Optional[Schema]]:
+ parsed_uri = urlparse(uri)
+ if parsed_uri.scheme == "vfs":
+ _, _, _, resource_type = VFSURIFactory.decode_uri(uri)
+
+ if resource_type in {VFSResourceType.RESULT}:
+ storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
+
+ table = load_table_metadata(
+ IcebergCatalogInstance.get_instance(),
+ StorageConfig.ICEBERG_TABLE_RESULT_NAMESPACE,
+ storage_key,
+ )
+
+ if table is None:
+ raise ValueError("No storage is found for the given URI")
+
+ amber_schema = Schema(table.schema().as_arrow())
+
+ document = IcebergDocument(
+ StorageConfig.ICEBERG_TABLE_RESULT_NAMESPACE,
+ storage_key,
+ table.schema(),
+ amber_tuples_to_arrow_table,
+ arrow_table_to_amber_tuples,
+ )
+ return document, amber_schema
+ else:
+ raise ValueError(f"Resource type {resource_type} is not supported")
+ else:
+ raise NotImplementedError(
+ f"Unsupported URI scheme: {parsed_uri.scheme} for opening the document"
+ )
diff --git a/amber/src/main/python/core/storage/iceberg/__init__.py b/amber/src/main/python/core/storage/iceberg/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/storage/iceberg/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/storage/iceberg/iceberg_catalog_instance.py b/amber/src/main/python/core/storage/iceberg/iceberg_catalog_instance.py
new file mode 100644
index 00000000000..0059808f9f8
--- /dev/null
+++ b/amber/src/main/python/core/storage/iceberg/iceberg_catalog_instance.py
@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from pyiceberg.catalog import Catalog
+from typing import Optional
+
+from core.storage.iceberg.iceberg_utils import (
+ create_postgres_catalog,
+ create_rest_catalog,
+)
+from core.storage.storage_config import StorageConfig
+
+
+class IcebergCatalogInstance:
+ """
+ IcebergCatalogInstance is a singleton that manages the Iceberg catalog instance.
+ Supports postgres SQL catalog and REST catalog.
+ - Provides a single shared catalog for all Iceberg table-related operations.
+ - Lazily initializes the catalog on first access.
+ - Supports replacing the catalog instance for testing or reconfiguration.
+ """
+
+ _instance: Optional[Catalog] = None
+
+ @classmethod
+ def get_instance(cls):
+ """
+ Retrieves the singleton Iceberg catalog instance.
+ - If the catalog is not initialized, it is lazily created using the configured
+ properties.
+ - Supports "postgres" and "rest" catalog types.
+ :return: the Iceberg catalog instance.
+ """
+ if cls._instance is None:
+ catalog_type = StorageConfig.ICEBERG_CATALOG_TYPE
+ if catalog_type == "postgres":
+ cls._instance = create_postgres_catalog(
+ "texera_iceberg",
+ StorageConfig.ICEBERG_FILE_STORAGE_DIRECTORY_PATH,
+ StorageConfig.ICEBERG_POSTGRES_CATALOG_URI_WITHOUT_SCHEME,
+ StorageConfig.ICEBERG_POSTGRES_CATALOG_USERNAME,
+ StorageConfig.ICEBERG_POSTGRES_CATALOG_PASSWORD,
+ )
+ elif catalog_type == "rest":
+ cls._instance = create_rest_catalog(
+ "texera_iceberg",
+ StorageConfig.ICEBERG_REST_CATALOG_WAREHOUSE_NAME,
+ StorageConfig.ICEBERG_REST_CATALOG_URI,
+ StorageConfig.S3_ENDPOINT,
+ StorageConfig.S3_REGION,
+ StorageConfig.S3_AUTH_USERNAME,
+ StorageConfig.S3_AUTH_PASSWORD,
+ )
+ else:
+ raise ValueError(f"Unsupported catalog type: {catalog_type}")
+ return cls._instance
+
+ @classmethod
+ def replace_instance(cls, catalog: Catalog):
+ """
+ Replaces the existing Iceberg catalog instance.
+ - This method is useful for testing or dynamically updating the catalog.
+ :param catalog: the new Iceberg catalog instance to replace the current one.
+ """
+ cls._instance = catalog
diff --git a/amber/src/main/python/core/storage/iceberg/iceberg_document.py b/amber/src/main/python/core/storage/iceberg/iceberg_document.py
new file mode 100644
index 00000000000..997ab9b5b70
--- /dev/null
+++ b/amber/src/main/python/core/storage/iceberg/iceberg_document.py
@@ -0,0 +1,276 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+from itertools import islice
+from pyiceberg.catalog import Catalog
+from pyiceberg.schema import Schema
+from pyiceberg.table import Table, FileScanTask
+from readerwriterlock import rwlock
+from threading import RLock
+from typing import Iterator, Optional, Callable, Iterable
+from typing import TypeVar
+from urllib.parse import ParseResult, urlparse
+
+from core.storage.iceberg.iceberg_catalog_instance import IcebergCatalogInstance
+from core.storage.iceberg.iceberg_table_writer import IcebergTableWriter
+from core.storage.iceberg.iceberg_utils import (
+ load_table_metadata,
+ read_data_file_as_arrow_table,
+)
+from core.storage.model.virtual_document import VirtualDocument
+
+# Define a type variable
+T = TypeVar("T")
+
+
+class IcebergDocument(VirtualDocument[T]):
+ """
+ IcebergDocument is used to read and write a set of T as an Iceberg table.
+ It provides iterator-based read methods and supports multiple writers to write to
+ the same table.
+
+ - On construction, the table will be created if it does not exist.
+ - If the table exists, it will be overridden.
+
+ :param table_namespace: Namespace of the table.
+ :param table_name: Name of the table.
+ :param table_schema: Schema of the table.
+ :param serde: A function to convert a T iterable into a pyarrow Table. Note the
+ conversion is not based on a single T item (unlike Texera's Java IcebergDocument.)
+ :param deserde: A function to convert a pyarrow Table back into a T iterable.
+ """
+
+ def __init__(
+ self,
+ table_namespace: str,
+ table_name: str,
+ table_schema: Schema,
+ serde: Callable[[Schema, Iterable[T]], pa.Table],
+ deserde: Callable[[Schema, pa.Table], Iterable[T]],
+ ):
+ self.table_namespace = table_namespace
+ self.table_name = table_name
+ self.table_schema = table_schema
+ self.serde = serde
+ self.deserde = deserde
+
+ self.lock = rwlock.RWLockFair()
+ self.catalog = IcebergCatalogInstance.get_instance()
+
+ def get_uri(self) -> ParseResult:
+ """Returns the URI of the table location."""
+ table = load_table_metadata(self.catalog, self.table_namespace, self.table_name)
+ if not table:
+ raise Exception(
+ f"table {self.table_namespace}.{self.table_name} doesn't exist."
+ )
+ return urlparse(table.location())
+
+ def clear(self):
+ """Deletes the table and clears its contents."""
+ with self.lock.gen_wlock():
+ table_identifier = f"{self.table_namespace}.{self.table_name}"
+ if self.catalog.table_exists(table_identifier):
+ self.catalog.drop_table(table_identifier)
+
+ def get(self) -> Iterator[T]:
+ """Get an iterator for reading all records from the table."""
+ return self._get_using_file_sequence_order(0, None)
+
+ def get_range(self, from_index: int, until_index: int) -> Iterator[T]:
+ """Get records within a specified range [from, until)."""
+ return self._get_using_file_sequence_order(from_index, until_index)
+
+ def get_after(self, offset: int) -> Iterator[T]:
+ """Get records starting after a specified offset."""
+ return self._get_using_file_sequence_order(offset, None)
+
+ def get_count(self) -> int:
+ """Get the total count of records in the table."""
+ table = load_table_metadata(self.catalog, self.table_namespace, self.table_name)
+ if not table:
+ return 0
+ return sum(f.file.record_count for f in table.scan().plan_files())
+
+ def writer(self, writer_identifier: str):
+ """
+ Creates a BufferedItemWriter for writing data to the table.
+ :param writer_identifier: The writer's ID. It should be unique within the same
+ table, as each writer will use it as the prefix of the files they append
+ :return: An IcebergTableWriter
+ """
+ return IcebergTableWriter[T](
+ writer_identifier=writer_identifier,
+ catalog=self.catalog,
+ table_namespace=self.table_namespace,
+ table_name=self.table_name,
+ table_schema=self.table_schema,
+ serde=self.serde,
+ )
+
+ def _get_using_file_sequence_order(
+ self, from_index: int, until_index: Optional[int]
+ ) -> Iterator[T]:
+ """Utility to get records within a specified range."""
+ with self.lock.gen_rlock():
+ return IcebergIterator[T](
+ from_index,
+ until_index,
+ self.catalog,
+ self.table_namespace,
+ self.table_name,
+ self.table_schema,
+ self.deserde,
+ )
+
+
+class IcebergIterator(Iterator[T]):
+ """
+ A custom iterator class to read items from an iceberg table based on an index range.
+ """
+
+ def __init__(
+ self,
+ from_index: int,
+ until_index: int,
+ catalog: Catalog,
+ table_namespace: str,
+ table_name: str,
+ table_schema: Schema,
+ deserde: Callable[[Schema, pa.Table], Iterable[T]],
+ ):
+ self.from_index = from_index
+ self.until_index = until_index
+ self.catalog = catalog
+ self.table_namespace = table_namespace
+ self.table_name = table_name
+ self.table_schema = table_schema
+ self.deserde = deserde
+ self.lock = RLock()
+ # Counter for how many records have been skipped
+ self.num_of_skipped_records = 0
+ # Counter for how many records have been returned
+ self.num_of_returned_records = 0
+ # Total number of records to return, used for termination condition
+ self.total_records_to_return = (
+ self.until_index - self.from_index if until_index else float("inf")
+ )
+ # Load the table instance, initially the table instance may not exist
+ self.table = self._load_table_metadata()
+ # Iterator for usable file scan tasks
+ self.usable_file_iterator = self._seek_to_usable_file()
+ # Current record iterator for the active file
+ self.current_record_iterator = iter([])
+
+ def _load_table_metadata(self) -> Optional[Table]:
+ """Util function to load the table's metadata."""
+ return load_table_metadata(self.catalog, self.table_namespace, self.table_name)
+
+ def _seek_to_usable_file(self) -> Iterator[FileScanTask]:
+ """Find usable file scan tasks starting from the specified record index."""
+ with self.lock:
+ if self.num_of_skipped_records > self.from_index:
+ raise RuntimeError("seek operation should not be called")
+
+ # Load the table for the first time
+ if not self.table:
+ self.table = self._load_table_metadata()
+
+ # If the table still does not exist after loading, end iterator.
+ if self.table:
+ try:
+ self.table.refresh()
+ current_snapshot = self.table.current_snapshot()
+ if current_snapshot is None:
+ return iter([])
+ sorted_file_scan_tasks = self._extract_sorted_file_scan_tasks(
+ current_snapshot
+ )
+ # Skip records in files before the `from_index`
+ for task in sorted_file_scan_tasks:
+ record_count = task.file.record_count
+ if (
+ self.num_of_skipped_records + record_count
+ <= self.from_index
+ ):
+ self.num_of_skipped_records += record_count
+ continue
+ yield task
+ except Exception:
+ print("Could not read iceberg table:\n")
+ raise Exception
+ else:
+ return iter([])
+
+ def _extract_sorted_file_scan_tasks(self, current_snapshot):
+ """
+ As self.table.inspect.entries() does not work with java files, this method
+ implements the logic to find file_sequence_number for each data file ourselves
+ :param current_snapshot: The current snapshot of the table.
+ :return: The file scan tasks of the file sorted by file_sequence_number
+ """
+ file_sequence_map = {}
+ for manifest in current_snapshot.manifests(self.table.io):
+ for entry in manifest.fetch_manifest_entry(io=self.table.io):
+ file_sequence_map[entry.data_file.file_path] = entry.sequence_number
+ # Retrieve and sort the file scan tasks by file sequence number
+ file_scan_tasks = list(self.table.scan().plan_files())
+ # Sort files by their sequence number. Files without a sequence
+ # number will be read last.
+ sorted_file_scan_tasks = sorted(
+ file_scan_tasks,
+ key=lambda t: file_sequence_map.get(t.file.file_path, float("inf")),
+ )
+ return sorted_file_scan_tasks
+
+ def __iter__(self) -> Iterator[T]:
+ return self
+
+ def __next__(self) -> T:
+ if self.num_of_returned_records >= self.total_records_to_return:
+ raise StopIteration("No more records available")
+
+ while True:
+ try:
+ record = next(self.current_record_iterator)
+ self.num_of_returned_records += 1
+ return record
+ except StopIteration:
+ # current_record_iterator is exhausted, need to go to the next file
+ try:
+ next_file = next(self.usable_file_iterator)
+ arrow_table = read_data_file_as_arrow_table(next_file, self.table)
+ self.current_record_iterator = self.deserde(
+ self.table_schema, arrow_table
+ )
+ # Skip records within the file if necessary
+ records_to_skip_in_file = (
+ self.from_index - self.num_of_skipped_records
+ )
+ if records_to_skip_in_file > 0:
+ self.current_record_iterator = self._skip_records(
+ self.current_record_iterator, records_to_skip_in_file
+ )
+ self.num_of_skipped_records += records_to_skip_in_file
+ except StopIteration:
+ # no more files left in this table
+ raise StopIteration("No more records available")
+
+ @staticmethod
+ def _skip_records(iterator, count):
+ return islice(iterator, count, None)
diff --git a/amber/src/main/python/core/storage/iceberg/iceberg_table_writer.py b/amber/src/main/python/core/storage/iceberg/iceberg_table_writer.py
new file mode 100644
index 00000000000..dd2a42c2ca8
--- /dev/null
+++ b/amber/src/main/python/core/storage/iceberg/iceberg_table_writer.py
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+from pyiceberg.catalog import Catalog
+from pyiceberg.schema import Schema
+from pyiceberg.table import Table
+from tenacity import retry, stop_after_attempt, wait_random_exponential
+from typing import List, TypeVar, Callable, Iterable
+
+from core.storage.model.buffered_item_writer import BufferedItemWriter
+from core.storage.storage_config import StorageConfig
+
+# Define a type variable for the data type T
+T = TypeVar("T")
+
+
+class IcebergTableWriter(BufferedItemWriter[T]):
+ """
+ IcebergTableWriter writes data to the given Iceberg table in an append-only way.
+ - Each time the buffer is flushed, a new data file is created using pyarrow
+ - Iceberg data files are immutable once created. So each flush will create a
+ distinct file.
+
+ **Thread Safety**: This writer is NOT thread-safe, so only one thread should call
+ this writer.
+
+ :param writer_identifier: A unique identifier used to prefix the created files.
+ :param catalog: The Iceberg catalog to manage table metadata.
+ :param table_namespace: The namespace of the Iceberg table.
+ :param table_name: The name of the Iceberg table.
+ :param table_schema: The schema of the Iceberg table.
+ """
+
+ def __init__(
+ self,
+ writer_identifier: str,
+ catalog: Catalog,
+ table_namespace: str,
+ table_name: str,
+ table_schema: pa.Schema,
+ serde: Callable[[Schema, Iterable[T]], pa.Table],
+ ):
+ self.writer_identifier = writer_identifier
+ self.catalog = catalog
+ self.table_namespace = table_namespace
+ self.table_name = table_name
+ self.table_schema = table_schema
+ self.serde = serde
+ self.buffer_size = StorageConfig.ICEBERG_TABLE_COMMIT_BATCH_SIZE
+
+ # Internal state
+ self.buffer: List[T] = []
+
+ # Load the Iceberg table
+ self.table: Table = self.catalog.load_table(
+ f"{self.table_namespace}.{self.table_name}"
+ )
+
+ @property
+ def buffer_size(self) -> int:
+ return self._buffer_size
+
+ def open(self) -> None:
+ """Open the writer and clear the buffer."""
+ self.buffer.clear()
+
+ def put_one(self, item: T) -> None:
+ """Add a single item to the buffer."""
+ self.buffer.append(item)
+ if len(self.buffer) >= self.buffer_size:
+ self._flush_buffer()
+
+ def remove_one(self, item: T) -> None:
+ """Remove a single item from the buffer."""
+ self.buffer.remove(item)
+
+ def _flush_buffer(self) -> None:
+ """
+ Flush the current buffer to a new Iceberg data file. The buffer is first
+ converted to a pyarrow table, and then appended to the iceberg table as a
+ parquet file. Note in the case of concurrent writers, as iceberg uses
+ optimistic concurrency control, we use a random exponential backoff mechanism
+ when commit failure happens because currently pyiceberg does not natively
+ support retry.
+ """
+ if not self.buffer:
+ return
+ df = self.serde(self.table_schema, self.buffer)
+
+ def append_to_table_with_retry(pa_df: pa.Table) -> None:
+ @retry(
+ wait=wait_random_exponential(0.001, 10),
+ stop=stop_after_attempt(10),
+ reraise=True,
+ )
+ def append_with_retry():
+ self.table.refresh()
+ self.table.append(pa_df)
+
+ append_with_retry()
+
+ append_to_table_with_retry(df)
+ self.buffer.clear()
+
+ def close(self) -> None:
+ """Close the writer, ensuring any remaining buffered items are flushed."""
+ if self.buffer:
+ self._flush_buffer()
+
+ @buffer_size.setter
+ def buffer_size(self, value):
+ self._buffer_size = value
diff --git a/amber/src/main/python/core/storage/iceberg/iceberg_utils.py b/amber/src/main/python/core/storage/iceberg/iceberg_utils.py
new file mode 100644
index 00000000000..844ef3e00ff
--- /dev/null
+++ b/amber/src/main/python/core/storage/iceberg/iceberg_utils.py
@@ -0,0 +1,313 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+import pyiceberg.table
+from pyiceberg.catalog import Catalog, load_catalog
+from pyiceberg.catalog.sql import SqlCatalog
+from pyiceberg.expressions import AlwaysTrue
+from pyiceberg.io.pyarrow import ArrowScan
+from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC
+from pyiceberg.schema import Schema
+from pyiceberg.table import Table
+from typing import Optional, Iterable
+from pyiceberg import types as iceberg_types
+
+import core
+import core.models
+from core.models import ArrowTableTupleProvider, Tuple
+from core.models.schema.attribute_type import AttributeType, TO_ARROW_MAPPING
+
+# Suffix used to encode LARGE_BINARY fields in Iceberg (must match Scala IcebergUtil)
+LARGE_BINARY_FIELD_SUFFIX = "__texera_large_binary_ptr"
+
+# Type mappings
+_ICEBERG_TO_AMBER_TYPE_MAPPING = {
+ "string": "STRING",
+ "int": "INT",
+ "integer": "INT",
+ "long": "LONG",
+ "double": "DOUBLE",
+ "float": "DOUBLE",
+ "boolean": "BOOL",
+ "timestamp": "TIMESTAMP",
+ "binary": "BINARY",
+}
+
+_AMBER_TO_ICEBERG_TYPE_MAPPING = {
+ AttributeType.STRING: iceberg_types.StringType(),
+ AttributeType.INT: iceberg_types.IntegerType(),
+ AttributeType.LONG: iceberg_types.LongType(),
+ AttributeType.DOUBLE: iceberg_types.DoubleType(),
+ AttributeType.BOOL: iceberg_types.BooleanType(),
+ AttributeType.TIMESTAMP: iceberg_types.TimestampType(),
+ AttributeType.BINARY: iceberg_types.BinaryType(),
+ AttributeType.LARGE_BINARY: iceberg_types.StringType(),
+}
+
+
+def encode_large_binary_field_name(field_name: str, attr_type) -> str:
+ """Encodes LARGE_BINARY field names with suffix for Iceberg storage."""
+ if attr_type == AttributeType.LARGE_BINARY:
+ return f"{field_name}{LARGE_BINARY_FIELD_SUFFIX}"
+ return field_name
+
+
+def decode_large_binary_field_name(field_name: str) -> str:
+ """Decodes field names by removing LARGE_BINARY suffix if present."""
+ if field_name.endswith(LARGE_BINARY_FIELD_SUFFIX):
+ return field_name[: -len(LARGE_BINARY_FIELD_SUFFIX)]
+ return field_name
+
+
+def iceberg_schema_to_amber_schema(iceberg_schema: Schema):
+ """
+ Converts PyIceberg Schema to Amber Schema.
+ Decodes LARGE_BINARY field names and adds Arrow metadata.
+ """
+ arrow_fields = []
+ for field in iceberg_schema.fields:
+ decoded_name = decode_large_binary_field_name(field.name)
+ is_large_binary = field.name != decoded_name
+
+ if is_large_binary:
+ attr_type = AttributeType.LARGE_BINARY
+ else:
+ iceberg_type_str = str(field.field_type).lower()
+ attr_type_name = _ICEBERG_TO_AMBER_TYPE_MAPPING.get(
+ iceberg_type_str, "STRING"
+ )
+ attr_type = getattr(AttributeType, attr_type_name)
+
+ arrow_fields.append(
+ pa.field(
+ decoded_name,
+ TO_ARROW_MAPPING[attr_type],
+ metadata={b"texera_type": b"LARGE_BINARY"} if is_large_binary else None,
+ )
+ )
+
+ return core.models.Schema(pa.schema(arrow_fields))
+
+
+def amber_schema_to_iceberg_schema(amber_schema) -> Schema:
+ """
+ Converts Amber Schema to PyIceberg Schema.
+ Encodes LARGE_BINARY field names with suffix.
+ """
+ fields = [
+ iceberg_types.NestedField(
+ field_id=idx,
+ name=encode_large_binary_field_name(field_name, attr_type),
+ field_type=_AMBER_TO_ICEBERG_TYPE_MAPPING[attr_type],
+ required=False,
+ )
+ for idx, (field_name, attr_type) in enumerate(
+ amber_schema._name_type_mapping.items(), start=1
+ )
+ ]
+
+ return Schema(*fields)
+
+
+def create_postgres_catalog(
+ catalog_name: str,
+ warehouse_path: str,
+ uri_without_scheme: str,
+ username: str,
+ password: str,
+) -> SqlCatalog:
+ """
+ Creates a Postgres SQL catalog instance by connecting to the database named
+ "texera_iceberg_catalog".
+ - The only requirement of the database is that it already exists. Once pyiceberg
+ can connect to the database, it will handle the initializations.
+ :param catalog_name: the name of the catalog.
+ :param warehouse_path: the root path for the warehouse where the tables are stored.
+ :param uri_without_scheme: the uri of the postgres database but without
+ the scheme prefix since java and python use different schemes.
+ :param username: the username of the postgres database.
+ :param password: the password of the postgres database.
+ :return: a SQLCatalog instance.
+ """
+ return SqlCatalog(
+ catalog_name,
+ **{
+ "uri": f"postgresql+pg8000://{username}:{password}@{uri_without_scheme}",
+ "warehouse": warehouse_path,
+ },
+ )
+
+
+def create_rest_catalog(
+ catalog_name: str,
+ warehouse_name: str,
+ rest_uri: str,
+ s3_endpoint: str,
+ s3_region: str,
+ s3_username: str,
+ s3_password: str,
+) -> Catalog:
+ """
+ Creates a REST catalog instance by connecting to a REST endpoint.
+ - Configures the catalog to interact with a REST endpoint.
+ - The warehouse_name parameter specifies the warehouse identifier.
+ - Configures S3FileIO for MinIO/S3 storage backend.
+ :param catalog_name: the name of the catalog.
+ :param warehouse_name: the warehouse identifier.
+ :param rest_uri: the URI of the REST catalog endpoint.
+ :param s3_endpoint: the S3 endpoint URL.
+ :param s3_region: the S3 region.
+ :param s3_username: the S3 access key ID.
+ :param s3_password: the S3 secret access key.
+ :return: a Catalog instance (REST catalog).
+ """
+ return load_catalog(
+ catalog_name,
+ **{
+ "type": "rest",
+ "uri": rest_uri,
+ "warehouse": warehouse_name,
+ "s3.endpoint": s3_endpoint,
+ "s3.access-key-id": s3_username,
+ "s3.secret-access-key": s3_password,
+ "s3.region": s3_region,
+ "s3.path-style-access": "true",
+ },
+ )
+
+
+def create_table(
+ catalog: Catalog,
+ table_namespace: str,
+ table_name: str,
+ table_schema: Schema,
+ override_if_exists: bool = False,
+) -> Table:
+ """
+ Creates a new Iceberg table with the specified schema and properties.
+ - Drops the existing table if `override_if_exists` is true and the table already
+ exists.
+ - Creates an unpartitioned table with custom commit retry properties.
+
+ :param catalog: The Iceberg catalog to manage the table.
+ :param table_namespace: The namespace of the table.
+ :param table_name: The name of the table.
+ :param table_schema: The schema of the table.
+ :param override_if_exists: Whether to drop and recreate the table if it exists.
+ :return: The created Iceberg table.
+ """
+
+ identifier = f"{table_namespace}.{table_name}"
+
+ catalog.create_namespace_if_not_exists(table_namespace)
+
+ if catalog.table_exists(identifier) and override_if_exists:
+ catalog.drop_table(identifier)
+
+ table = catalog.create_table(
+ identifier=identifier,
+ schema=table_schema,
+ partition_spec=UNPARTITIONED_PARTITION_SPEC,
+ )
+
+ return table
+
+
+def load_table_metadata(
+ catalog: Catalog, table_namespace: str, table_name: str
+) -> Optional[Table]:
+ """
+ Loads metadata for an existing Iceberg table.
+ - Returns the table if it exists and is successfully loaded.
+ - Returns None if the table does not exist or cannot be loaded.
+
+ :param catalog: The Iceberg catalog to load the table from.
+ :param table_namespace: The namespace of the table.
+ :param table_name: The name of the table.
+ :return: The table if found, or None if not found.
+ """
+ identifier = f"{table_namespace}.{table_name}"
+ try:
+ return catalog.load_table(identifier)
+ except Exception:
+ return None
+
+
+def read_data_file_as_arrow_table(
+ planfile: pyiceberg.table.FileScanTask, iceberg_table: pyiceberg.table.Table
+) -> pa.Table:
+ """Reads a data file as a pyarrow table and returns an iterator over its records."""
+ arrow_table: pa.Table = ArrowScan(
+ iceberg_table.metadata,
+ iceberg_table.io,
+ iceberg_table.schema(),
+ AlwaysTrue(),
+ True,
+ ).to_table([planfile])
+ return arrow_table
+
+
+def amber_tuples_to_arrow_table(
+ iceberg_schema: Schema, tuple_list: Iterable[Tuple]
+) -> pa.Table:
+ """
+ Converts a list of amber tuples to a pyarrow table for serialization.
+ Handles LARGE_BINARY field name encoding and serialization.
+ """
+ from core.models.type.large_binary import largebinary
+
+ tuple_list = list(tuple_list) # Convert to list to allow multiple iterations
+ data_dict = {}
+ for encoded_name in iceberg_schema.as_arrow().names:
+ decoded_name = decode_large_binary_field_name(encoded_name)
+ data_dict[encoded_name] = [
+ (
+ t[decoded_name].uri
+ if isinstance(t[decoded_name], largebinary)
+ else t[decoded_name]
+ )
+ for t in tuple_list
+ ]
+
+ return pa.Table.from_pydict(data_dict, schema=iceberg_schema.as_arrow())
+
+
+def arrow_table_to_amber_tuples(
+ iceberg_schema: Schema, arrow_table: pa.Table
+) -> Iterable[Tuple]:
+ """
+ Converts an arrow table read from Iceberg to Amber tuples.
+ Properly handles LARGE_BINARY field name decoding and type detection.
+ """
+ amber_schema = iceberg_schema_to_amber_schema(iceberg_schema)
+ arrow_table_with_metadata = pa.Table.from_arrays(
+ [arrow_table.column(name) for name in arrow_table.column_names],
+ schema=amber_schema.as_arrow_schema(),
+ )
+
+ tuple_provider = ArrowTableTupleProvider(arrow_table_with_metadata)
+ return (
+ Tuple(
+ {
+ decode_large_binary_field_name(name): field_accessor
+ for name in arrow_table.column_names
+ },
+ schema=amber_schema,
+ )
+ for field_accessor in tuple_provider
+ )
diff --git a/amber/src/main/python/core/storage/iceberg/test_iceberg_document.py b/amber/src/main/python/core/storage/iceberg/test_iceberg_document.py
new file mode 100644
index 00000000000..9b374f7d5c7
--- /dev/null
+++ b/amber/src/main/python/core/storage/iceberg/test_iceberg_document.py
@@ -0,0 +1,319 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import pytest
+import random
+import uuid
+from concurrent.futures import as_completed
+from concurrent.futures.thread import ThreadPoolExecutor
+
+from core.models import Schema, Tuple
+from core.storage.document_factory import DocumentFactory
+from core.storage.storage_config import StorageConfig
+from core.storage.vfs_uri_factory import VFSURIFactory
+from proto.org.apache.texera.amber.core import (
+ WorkflowIdentity,
+ ExecutionIdentity,
+ OperatorIdentity,
+ PortIdentity,
+ GlobalPortIdentity,
+ PhysicalOpIdentity,
+)
+
+# Hardcoded storage config only for test purposes.
+StorageConfig.initialize(
+ catalog_type="postgres",
+ postgres_uri_without_scheme="localhost:5432/texera_iceberg_catalog",
+ postgres_username="texera",
+ postgres_password="password",
+ rest_catalog_uri="http://localhost:8181/catalog/",
+ rest_catalog_warehouse_name="texera",
+ table_result_namespace="operator-port-result",
+ directory_path="../../../../../../amber/user-resources/workflow-results",
+ commit_batch_size=4096,
+ s3_endpoint="http://localhost:9000",
+ s3_region="us-east-1",
+ s3_auth_username="minioadmin",
+ s3_auth_password="minioadmin",
+)
+
+
+class TestIcebergDocument:
+ @pytest.fixture
+ def amber_schema(self):
+ """Sample Amber schema"""
+ return Schema(
+ raw_schema={
+ "col-string": "STRING",
+ "col-int": "INTEGER",
+ "col-bool": "BOOLEAN",
+ "col-long": "LONG",
+ "col-double": "DOUBLE",
+ "col-timestamp": "TIMESTAMP",
+ "col-binary": "BINARY",
+ }
+ )
+
+ @pytest.fixture
+ def iceberg_document(self, amber_schema):
+ """
+ Creates an iceberg document of operator port results using the sample schema
+ with a random operator id
+ """
+ operator_uuid = str(uuid.uuid4()).replace("-", "")
+ uri = VFSURIFactory.create_result_uri(
+ WorkflowIdentity(id=0),
+ ExecutionIdentity(id=0),
+ GlobalPortIdentity(
+ op_id=PhysicalOpIdentity(
+ logical_op_id=OperatorIdentity(id=f"test_table_{operator_uuid}"),
+ layer_name="main",
+ ),
+ port_id=PortIdentity(id=0),
+ input=False,
+ ),
+ )
+ DocumentFactory.create_document(uri, amber_schema)
+ document, _ = DocumentFactory.open_document(uri)
+ return document
+
+ @pytest.fixture
+ def sample_items(self, amber_schema) -> [Tuple]:
+ """
+ Generates a list of sample tuples
+ """
+ base_tuples = [
+ Tuple(
+ {
+ "col-string": "Hello World",
+ "col-int": 42,
+ "col-bool": True,
+ "col-long": 1123213213213,
+ "col-double": 214214.9969346,
+ "col-timestamp": datetime.datetime.now(),
+ "col-binary": b"hello",
+ },
+ schema=amber_schema,
+ ),
+ Tuple(
+ {
+ "col-string": "",
+ "col-int": -1,
+ "col-bool": False,
+ "col-long": -98765432109876,
+ "col-double": -0.001,
+ "col-timestamp": datetime.datetime.fromtimestamp(100000000),
+ "col-binary": bytearray([255, 0, 0, 64]),
+ },
+ schema=amber_schema,
+ ),
+ Tuple(
+ {
+ "col-string": "Special Characters: \n\t\r",
+ "col-int": 2147483647,
+ "col-bool": True,
+ "col-long": 9223372036854775807,
+ "col-double": 1.7976931348623157e308,
+ "col-timestamp": datetime.datetime.fromtimestamp(1234567890),
+ "col-binary": bytearray([1, 2, 3, 4, 5]),
+ },
+ schema=amber_schema,
+ ),
+ ]
+
+ # Function to generate random binary data
+ def generate_random_binary(size):
+ return bytearray(random.getrandbits(8) for _ in range(size))
+
+ # Generate additional tuples
+ additional_tuples = [
+ Tuple(
+ {
+ "col-string": None if i % 7 == 0 else f"Generated String {i}",
+ "col-int": None if i % 5 == 0 else i,
+ "col-bool": None if i % 6 == 0 else i % 2 == 0,
+ "col-long": None if i % 4 == 0 else i * 1000000,
+ "col-double": None if i % 3 == 0 else i * 0.12345,
+ "col-timestamp": (
+ None
+ if i % 8 == 0
+ else datetime.datetime.fromtimestamp(
+ datetime.datetime.now().timestamp() + i
+ )
+ ),
+ "col-binary": None if i % 9 == 0 else generate_random_binary(10),
+ },
+ schema=amber_schema,
+ )
+ for i in range(1, 20001)
+ ]
+
+ return base_tuples + additional_tuples
+
+ def test_basic_read_and_write(self, iceberg_document, sample_items):
+ """
+ Create an iceberg document, write sample items, and read it back.
+ """
+ writer = iceberg_document.writer(str(uuid.uuid4()))
+ writer.open()
+ for item in sample_items:
+ writer.put_one(item)
+ writer.close()
+ retrieved_items = list(iceberg_document.get())
+ assert sample_items == retrieved_items
+
+ def test_clear_document(self, iceberg_document, sample_items):
+ """
+ Create an iceberg document, write sample items, and clear the document.
+ """
+ writer = iceberg_document.writer(str(uuid.uuid4()))
+ writer.open()
+ for item in sample_items:
+ writer.put_one(item)
+ writer.close()
+ assert len(list(iceberg_document.get())) > 0
+
+ iceberg_document.clear()
+ assert len(list(iceberg_document.get())) == 0
+
+ def test_handle_empty_read(self, iceberg_document):
+ """
+ The iceberg document should handle empty reads gracefully
+ """
+ retrieved_items = list(iceberg_document.get())
+ assert retrieved_items == []
+
+ def test_concurrent_writes_followed_by_read(self, iceberg_document, sample_items):
+ """
+ Tests multiple concurrent writers writing to the same iceberg document
+ """
+ all_items = sample_items
+ num_writers = 10
+ # Calculate the batch size and the remainder
+ batch_size = len(all_items) // num_writers
+ remainder = len(all_items) % num_writers
+ # Create writer's batches
+ item_batches = [
+ all_items[
+ i * batch_size + min(i, remainder) : i * batch_size
+ + min(i, remainder)
+ + batch_size
+ + (1 if i < remainder else 0)
+ ]
+ for i in range(num_writers)
+ ]
+
+ assert len(item_batches) == num_writers, (
+ f"Expected {num_writers} batches but got {len(item_batches)}"
+ )
+
+ # Perform concurrent writes
+ def write_batch(batch):
+ writer = iceberg_document.writer(str(uuid.uuid4()))
+ writer.open()
+ for item in batch:
+ writer.put_one(item)
+ writer.close()
+
+ with ThreadPoolExecutor(max_workers=num_writers) as executor:
+ futures = [executor.submit(write_batch, batch) for batch in item_batches]
+ for future in as_completed(futures):
+ future.result() # Wait for each future to complete
+
+ # Read all items back
+ retrieved_items = list(iceberg_document.get())
+ # Verify that the retrieved items match the original items
+ assert set(retrieved_items) == set(all_items), (
+ "All items should be read correctly after concurrent writes."
+ )
+
+ def test_read_using_range(self, iceberg_document, sample_items):
+ """
+ The iceberg document should read all items using rages correctly.
+ """
+ writer = iceberg_document.writer(str(uuid.uuid4()))
+ writer.open()
+ for item in sample_items:
+ writer.put_one(item)
+ writer.close()
+ # Read all items using ranges
+ batch_size = 1500
+ # Generate ranges
+ ranges = [
+ range(i, min(i + batch_size, len(sample_items)))
+ for i in range(0, len(sample_items), batch_size)
+ ]
+
+ # Retrieve items using ranges
+ retrieved_items = [
+ item for r in ranges for item in iceberg_document.get_range(r.start, r.stop)
+ ]
+
+ assert len(retrieved_items) == len(sample_items), (
+ "The number of retrieved items does not match the number of all items."
+ )
+
+ # Verify that the retrieved items match the original items
+ assert set(retrieved_items) == set(sample_items), (
+ "All items should be retrieved correctly using ranges."
+ )
+
+ def test_get_after(self, iceberg_document, sample_items):
+ """
+ The iceberg document should retrieve items correctly using get_after
+ """
+ writer = iceberg_document.writer(str(uuid.uuid4()))
+ writer.open()
+ for item in sample_items:
+ writer.put_one(item)
+ writer.close()
+ # Test get_after for various offsets
+ offsets = [0, len(sample_items) // 2, len(sample_items) - 1]
+ for offset in offsets:
+ if offset < len(sample_items):
+ expected_items = sample_items[offset:]
+ else:
+ expected_items = []
+
+ retrieved_items = list(iceberg_document.get_after(offset))
+ assert retrieved_items == expected_items, (
+ f"get_after({offset}) did not return the expected items. "
+ f"Expected: {expected_items}, Got: {retrieved_items}"
+ )
+
+ # Test get_after for an offset beyond the range
+ invalid_offset = len(sample_items)
+ retrieved_items = list(iceberg_document.get_after(invalid_offset))
+ assert not retrieved_items, (
+ f"get_after({invalid_offset}) should return "
+ f"an empty list, but got: {retrieved_items}"
+ )
+
+ def test_get_counts(self, iceberg_document, sample_items):
+ """
+ The iceberg document should correctly return the count of items.
+ """
+ writer = iceberg_document.writer(str(uuid.uuid4()))
+ writer.open()
+ for item in sample_items:
+ writer.put_one(item)
+ writer.close()
+
+ assert iceberg_document.get_count() == len(sample_items), (
+ "get_count should return the same number as the length of sample_items"
+ )
diff --git a/amber/src/main/python/core/storage/iceberg/test_iceberg_utils_catalog.py b/amber/src/main/python/core/storage/iceberg/test_iceberg_utils_catalog.py
new file mode 100644
index 00000000000..902829d44c4
--- /dev/null
+++ b/amber/src/main/python/core/storage/iceberg/test_iceberg_utils_catalog.py
@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from unittest.mock import patch
+
+from core.storage.iceberg import iceberg_utils
+from core.storage.iceberg.iceberg_utils import create_postgres_catalog
+
+
+class TestCreatePostgresCatalog:
+ """
+ Regression tests for `create_postgres_catalog`.
+
+ The Scala side (`IcebergUtil.createPostgresCatalog`) initializes the JDBC
+ catalog with a plain filesystem warehouse path (no URI scheme). PyIceberg
+ persists the `warehouse` property into table metadata, so if the Python
+ side registers the catalog with a `file://`-prefixed value, Iceberg tables
+ written from Python UDFs become unreadable from the Scala/Java engine
+ (and vice versa). These tests pin the Python side to the same plain-path
+ convention used on the Scala side.
+ """
+
+ def test_warehouse_is_passed_without_file_scheme(self):
+ """`warehouse` must be forwarded as-is, without a `file://` prefix."""
+ warehouse_path = "/tmp/texera/iceberg-warehouse"
+
+ with patch.object(iceberg_utils, "SqlCatalog") as mock_sql_catalog:
+ create_postgres_catalog(
+ catalog_name="texera_iceberg",
+ warehouse_path=warehouse_path,
+ uri_without_scheme="localhost:5432/texera_iceberg_catalog",
+ username="texera",
+ password="password",
+ )
+
+ assert mock_sql_catalog.call_count == 1
+ _, kwargs = mock_sql_catalog.call_args
+ assert kwargs["warehouse"] == warehouse_path
+ assert not kwargs["warehouse"].startswith("file://")
+
+ def test_windows_style_warehouse_is_passed_verbatim(self):
+ """
+ The Scala side strips the Windows drive colon (e.g. `C:/x` -> `C/x`)
+ before registering the catalog so PyArrow can parse the path. The
+ Python side should forward whatever it receives verbatim, so the two
+ runtimes agree on the warehouse string stored in Iceberg metadata.
+ """
+ warehouse_path = "C/Users/texera/iceberg-warehouse"
+
+ with patch.object(iceberg_utils, "SqlCatalog") as mock_sql_catalog:
+ create_postgres_catalog(
+ catalog_name="texera_iceberg",
+ warehouse_path=warehouse_path,
+ uri_without_scheme="localhost:5432/texera_iceberg_catalog",
+ username="texera",
+ password="password",
+ )
+
+ _, kwargs = mock_sql_catalog.call_args
+ assert kwargs["warehouse"] == warehouse_path
+ assert "file://" not in kwargs["warehouse"]
+
+ def test_postgres_uri_is_built_with_pg8000_scheme(self):
+ """The JDBC URI should be prefixed with `postgresql+pg8000://` and
+ include credentials; nothing about that should bleed into `warehouse`.
+ """
+ warehouse_path = "/var/lib/texera/warehouse"
+
+ with patch.object(iceberg_utils, "SqlCatalog") as mock_sql_catalog:
+ create_postgres_catalog(
+ catalog_name="texera_iceberg",
+ warehouse_path=warehouse_path,
+ uri_without_scheme="db.internal:5432/texera_iceberg_catalog",
+ username="texera",
+ password="s3cret",
+ )
+
+ args, kwargs = mock_sql_catalog.call_args
+ assert args == ("texera_iceberg",)
+ assert kwargs["uri"] == (
+ "postgresql+pg8000://texera:s3cret@db.internal:5432/texera_iceberg_catalog"
+ )
+ # And warehouse is still the plain path.
+ assert kwargs["warehouse"] == warehouse_path
diff --git a/amber/src/main/python/core/storage/iceberg/test_iceberg_utils_large_binary.py b/amber/src/main/python/core/storage/iceberg/test_iceberg_utils_large_binary.py
new file mode 100644
index 00000000000..c601d4f7190
--- /dev/null
+++ b/amber/src/main/python/core/storage/iceberg/test_iceberg_utils_large_binary.py
@@ -0,0 +1,230 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+from pyiceberg import types as iceberg_types
+from pyiceberg.schema import Schema as IcebergSchema
+from core.models import Schema, Tuple
+from core.models.schema.attribute_type import AttributeType
+from core.models.type.large_binary import largebinary
+from core.storage.iceberg.iceberg_utils import (
+ encode_large_binary_field_name,
+ decode_large_binary_field_name,
+ iceberg_schema_to_amber_schema,
+ amber_schema_to_iceberg_schema,
+ amber_tuples_to_arrow_table,
+ arrow_table_to_amber_tuples,
+)
+
+
+class TestIcebergUtilsLargeBinary:
+ def test_encode_large_binary_field_name(self):
+ """Test encoding LARGE_BINARY field names with suffix."""
+ assert (
+ encode_large_binary_field_name("my_field", AttributeType.LARGE_BINARY)
+ == "my_field__texera_large_binary_ptr"
+ )
+ assert (
+ encode_large_binary_field_name("my_field", AttributeType.STRING)
+ == "my_field"
+ )
+
+ def test_decode_large_binary_field_name(self):
+ """Test decoding LARGE_BINARY field names by removing suffix."""
+ assert (
+ decode_large_binary_field_name("my_field__texera_large_binary_ptr")
+ == "my_field"
+ )
+ assert decode_large_binary_field_name("my_field") == "my_field"
+ assert decode_large_binary_field_name("regular_field") == "regular_field"
+
+ def test_amber_schema_to_iceberg_schema_with_large_binary(self):
+ """Test converting Amber schema with LARGE_BINARY to Iceberg schema."""
+ amber_schema = Schema()
+ amber_schema.add("regular_field", AttributeType.STRING)
+ amber_schema.add("large_binary_field", AttributeType.LARGE_BINARY)
+ amber_schema.add("int_field", AttributeType.INT)
+
+ iceberg_schema = amber_schema_to_iceberg_schema(amber_schema)
+
+ # Check field names are encoded
+ field_names = [field.name for field in iceberg_schema.fields]
+ assert "regular_field" in field_names
+ assert "large_binary_field__texera_large_binary_ptr" in field_names
+ assert "int_field" in field_names
+
+ # Check types
+ large_binary_field = next(
+ f for f in iceberg_schema.fields if "large_binary" in f.name
+ )
+ assert isinstance(large_binary_field.field_type, iceberg_types.StringType)
+
+ def test_iceberg_schema_to_amber_schema_with_large_binary(self):
+ """Test converting Iceberg schema with LARGE_BINARY to Amber schema."""
+ iceberg_schema = IcebergSchema(
+ iceberg_types.NestedField(
+ 1, "regular_field", iceberg_types.StringType(), required=False
+ ),
+ iceberg_types.NestedField(
+ 2,
+ "large_binary_field__texera_large_binary_ptr",
+ iceberg_types.StringType(),
+ required=False,
+ ),
+ iceberg_types.NestedField(
+ 3, "int_field", iceberg_types.IntegerType(), required=False
+ ),
+ )
+
+ amber_schema = iceberg_schema_to_amber_schema(iceberg_schema)
+
+ assert amber_schema.get_attr_type("regular_field") == AttributeType.STRING
+ assert (
+ amber_schema.get_attr_type("large_binary_field")
+ == AttributeType.LARGE_BINARY
+ )
+ assert amber_schema.get_attr_type("int_field") == AttributeType.INT
+
+ # Check Arrow schema has metadata for LARGE_BINARY
+ arrow_schema = amber_schema.as_arrow_schema()
+ large_binary_field = arrow_schema.field("large_binary_field")
+ assert large_binary_field.metadata is not None
+ assert large_binary_field.metadata.get(b"texera_type") == b"LARGE_BINARY"
+
+ def test_amber_tuples_to_arrow_table_with_large_binary(self):
+ """Test converting Amber tuples with largebinary to Arrow table."""
+ amber_schema = Schema()
+ amber_schema.add("regular_field", AttributeType.STRING)
+ amber_schema.add("large_binary_field", AttributeType.LARGE_BINARY)
+
+ large_binary1 = largebinary("s3://bucket/path1")
+ large_binary2 = largebinary("s3://bucket/path2")
+
+ tuples = [
+ Tuple(
+ {"regular_field": "value1", "large_binary_field": large_binary1},
+ schema=amber_schema,
+ ),
+ Tuple(
+ {"regular_field": "value2", "large_binary_field": large_binary2},
+ schema=amber_schema,
+ ),
+ ]
+
+ iceberg_schema = amber_schema_to_iceberg_schema(amber_schema)
+ arrow_table = amber_tuples_to_arrow_table(iceberg_schema, tuples)
+
+ # Check that largebinary values are converted to URI strings
+ regular_values = arrow_table.column("regular_field").to_pylist()
+ large_binary_values = arrow_table.column(
+ "large_binary_field__texera_large_binary_ptr"
+ ).to_pylist()
+
+ assert regular_values == ["value1", "value2"]
+ assert large_binary_values == ["s3://bucket/path1", "s3://bucket/path2"]
+
+ def test_arrow_table_to_amber_tuples_with_large_binary(self):
+ """Test converting Arrow table with LARGE_BINARY to Amber tuples."""
+ # Create Iceberg schema with encoded field name
+ iceberg_schema = IcebergSchema(
+ iceberg_types.NestedField(
+ 1, "regular_field", iceberg_types.StringType(), required=False
+ ),
+ iceberg_types.NestedField(
+ 2,
+ "large_binary_field__texera_large_binary_ptr",
+ iceberg_types.StringType(),
+ required=False,
+ ),
+ )
+
+ # Create Arrow table with URI strings
+ arrow_table = pa.Table.from_pydict(
+ {
+ "regular_field": ["value1", "value2"],
+ "large_binary_field__texera_large_binary_ptr": [
+ "s3://bucket/path1",
+ "s3://bucket/path2",
+ ],
+ }
+ )
+
+ tuples = list(arrow_table_to_amber_tuples(iceberg_schema, arrow_table))
+
+ assert len(tuples) == 2
+ assert tuples[0]["regular_field"] == "value1"
+ assert isinstance(tuples[0]["large_binary_field"], largebinary)
+ assert tuples[0]["large_binary_field"].uri == "s3://bucket/path1"
+
+ assert tuples[1]["regular_field"] == "value2"
+ assert isinstance(tuples[1]["large_binary_field"], largebinary)
+ assert tuples[1]["large_binary_field"].uri == "s3://bucket/path2"
+
+ def test_round_trip_large_binary_tuples(self):
+ """Test round-trip conversion of tuples with largebinary."""
+ amber_schema = Schema()
+ amber_schema.add("regular_field", AttributeType.STRING)
+ amber_schema.add("large_binary_field", AttributeType.LARGE_BINARY)
+
+ large_binary = largebinary("s3://bucket/path/to/object")
+ original_tuples = [
+ Tuple(
+ {"regular_field": "value1", "large_binary_field": large_binary},
+ schema=amber_schema,
+ ),
+ ]
+
+ # Convert to Iceberg and Arrow
+ iceberg_schema = amber_schema_to_iceberg_schema(amber_schema)
+ arrow_table = amber_tuples_to_arrow_table(iceberg_schema, original_tuples)
+
+ # Convert back to Amber tuples
+ retrieved_tuples = list(
+ arrow_table_to_amber_tuples(iceberg_schema, arrow_table)
+ )
+
+ assert len(retrieved_tuples) == 1
+ assert retrieved_tuples[0]["regular_field"] == "value1"
+ assert isinstance(retrieved_tuples[0]["large_binary_field"], largebinary)
+ assert retrieved_tuples[0]["large_binary_field"].uri == large_binary.uri
+
+ def test_arrow_table_to_amber_tuples_with_null_large_binary(self):
+ """Test converting Arrow table with null largebinary values."""
+ iceberg_schema = IcebergSchema(
+ iceberg_types.NestedField(
+ 1, "regular_field", iceberg_types.StringType(), required=False
+ ),
+ iceberg_types.NestedField(
+ 2,
+ "large_binary_field__texera_large_binary_ptr",
+ iceberg_types.StringType(),
+ required=False,
+ ),
+ )
+
+ arrow_table = pa.Table.from_pydict(
+ {
+ "regular_field": ["value1"],
+ "large_binary_field__texera_large_binary_ptr": [None],
+ }
+ )
+
+ tuples = list(arrow_table_to_amber_tuples(iceberg_schema, arrow_table))
+
+ assert len(tuples) == 1
+ assert tuples[0]["regular_field"] == "value1"
+ assert tuples[0]["large_binary_field"] is None
diff --git a/amber/src/main/python/core/storage/model/__init__.py b/amber/src/main/python/core/storage/model/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/storage/model/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/storage/model/buffered_item_writer.py b/amber/src/main/python/core/storage/model/buffered_item_writer.py
new file mode 100644
index 00000000000..2b50b34fcf7
--- /dev/null
+++ b/amber/src/main/python/core/storage/model/buffered_item_writer.py
@@ -0,0 +1,75 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+# Define a type variable
+T = TypeVar("T")
+
+
+class BufferedItemWriter(ABC, Generic[T]):
+ """
+ BufferedItemWriter provides an interface for writing items to a buffer and
+ performing I/O operations.
+ The items are buffered before being written to the underlying storage to
+ optimize performance.
+ """
+
+ @property
+ @abstractmethod
+ def buffer_size(self) -> int:
+ """
+ The size of the buffer.
+ :return: the buffer size.
+ """
+ pass
+
+ @abstractmethod
+ def open(self) -> None:
+ """
+ Open the writer, initializing any necessary resources.
+ This method should be called before any write operations.
+ """
+ pass
+
+ @abstractmethod
+ def close(self) -> None:
+ """
+ Close the writer, flushing any remaining items in the buffer
+ to the underlying storage and releasing any held resources.
+ """
+ pass
+
+ @abstractmethod
+ def put_one(self, item: T) -> None:
+ """
+ Put one item into the buffer. If the buffer is full, it should be flushed to
+ the underlying storage.
+ :param item: the data item to be written.
+ """
+ pass
+
+ @abstractmethod
+ def remove_one(self, item: T) -> None:
+ """
+ Remove one item from the buffer. If the item is not found in the buffer, an
+ appropriate action should be taken,
+ such as throwing an exception or ignoring the request.
+ :param item: the data item to be removed.
+ """
+ pass
diff --git a/amber/src/main/python/core/storage/model/readonly_virtual_document.py b/amber/src/main/python/core/storage/model/readonly_virtual_document.py
new file mode 100644
index 00000000000..986aaac77b1
--- /dev/null
+++ b/amber/src/main/python/core/storage/model/readonly_virtual_document.py
@@ -0,0 +1,85 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar, Iterator
+from urllib.parse import ParseResult
+
+# Define a type variable
+T = TypeVar("T")
+
+
+class ReadonlyVirtualDocument(ABC, Generic[T]):
+ """
+ ReadonlyVirtualDocument provides an abstraction for read operations over a single
+ resource.
+ This class can be implemented by resources that only need to support read-related
+ functionality.
+ """
+
+ @abstractmethod
+ def get_uri(self) -> ParseResult:
+ """
+ Get the URI of the corresponding document.
+ :return: the URI of the document as a ParseResult object
+ """
+ pass
+
+ @abstractmethod
+ def get_item(self, i: int) -> T:
+ """
+ Find the ith item and return.
+ :param i: index starting from 0
+ :return: data item of type T
+ """
+ pass
+
+ @abstractmethod
+ def get(self) -> Iterator[T]:
+ """
+ Get an iterator that iterates over all indexed items.
+ :return: an iterator that returns data items of type T
+ """
+ pass
+
+ @abstractmethod
+ def get_range(self, from_index: int, until_index: int) -> Iterator[T]:
+ """
+ Get an iterator of a sequence starting from index `from_index`, until index
+ `until`.
+ :param from_index: the starting index (inclusive)
+ :param until_index: the ending index (exclusive)
+ :return: an iterator that returns data items of type T
+ """
+ pass
+
+ @abstractmethod
+ def get_after(self, offset: int) -> Iterator[T]:
+ """
+ Get an iterator of all items after the specified index `offset`.
+ :param offset: the starting index (exclusive)
+ :return: an iterator that returns data items of type T
+ """
+ pass
+
+ @abstractmethod
+ def get_count(self) -> int:
+ """
+ Get the count of items in the document.
+ :return: the count of items
+ """
+ pass
diff --git a/amber/src/main/python/core/storage/model/virtual_document.py b/amber/src/main/python/core/storage/model/virtual_document.py
new file mode 100644
index 00000000000..e145c616138
--- /dev/null
+++ b/amber/src/main/python/core/storage/model/virtual_document.py
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import ABC, abstractmethod
+from overrides import overrides
+from typing import TypeVar, Iterator
+from urllib.parse import ParseResult
+
+from core.storage.model.buffered_item_writer import BufferedItemWriter
+from core.storage.model.readonly_virtual_document import ReadonlyVirtualDocument
+
+# Define a type variable
+T = TypeVar("T")
+
+
+class VirtualDocument(ReadonlyVirtualDocument[T], ABC):
+ """
+ VirtualDocument provides the abstraction of performing read/write/copy/delete
+ operations over a single resource.
+ Note that all methods have a default implementation. This is because one document
+ implementation may not be able to reasonably support all methods.
+ """
+
+ @overrides
+ def get_uri(self) -> ParseResult:
+ raise NotImplementedError("get_uri method is not implemented")
+
+ @overrides
+ def get_item(self, i: int) -> T:
+ raise NotImplementedError("get_item method is not implemented")
+
+ @overrides
+ def get(self) -> Iterator[T]:
+ raise NotImplementedError("get method is not implemented")
+
+ @overrides
+ def get_range(self, from_index: int, until_index: int) -> Iterator[T]:
+ raise NotImplementedError("get_range method is not implemented")
+
+ @overrides
+ def get_after(self, offset: int) -> Iterator[T]:
+ raise NotImplementedError("get_after method is not implemented")
+
+ @overrides
+ def get_count(self) -> int:
+ raise NotImplementedError("get_count method is not implemented")
+
+ def writer(self, writer_identifier: str) -> BufferedItemWriter[T]:
+ """
+ return a writer that buffers the items and performs the flush operation at
+ close time
+ :param writer_identifier: the id of the writer
+ :return: a buffered item writer
+ """
+ raise NotImplementedError("writer method is not implemented")
+
+ @abstractmethod
+ def clear(self) -> None:
+ """
+ Physically remove the current document.
+ """
+ pass
diff --git a/amber/src/main/python/core/storage/runnables/__init__.py b/amber/src/main/python/core/storage/runnables/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/storage/runnables/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/storage/runnables/input_port_materialization_reader_runnable.py b/amber/src/main/python/core/storage/runnables/input_port_materialization_reader_runnable.py
new file mode 100644
index 00000000000..e49c0316cc7
--- /dev/null
+++ b/amber/src/main/python/core/storage/runnables/input_port_materialization_reader_runnable.py
@@ -0,0 +1,214 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import typing
+from loguru import logger
+from pyarrow import Table
+from typing import Union
+
+from core.architecture.sendsemantics.broad_cast_partitioner import (
+ BroadcastPartitioner,
+)
+from core.architecture.sendsemantics.hash_based_shuffle_partitioner import (
+ HashBasedShufflePartitioner,
+)
+from core.architecture.sendsemantics.one_to_one_partitioner import OneToOnePartitioner
+from core.architecture.sendsemantics.partitioner import Partitioner
+from core.architecture.sendsemantics.range_based_shuffle_partitioner import (
+ RangeBasedShufflePartitioner,
+)
+from core.architecture.sendsemantics.round_robin_partitioner import (
+ RoundRobinPartitioner,
+)
+from core.models import Tuple, InternalQueue, DataFrame, DataPayload
+from core.models.internal_queue import DataElement, ECMElement
+from core.storage.document_factory import DocumentFactory
+from core.util import Stoppable, get_one_of
+from core.util.runnable.runnable import Runnable
+from core.util.virtual_identity import get_from_actor_id_for_input_port_storage
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ControlInvocation,
+ EmptyRequest,
+ EmbeddedControlMessageType,
+ EmbeddedControlMessage,
+ AsyncRpcContext,
+ ControlRequest,
+)
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ HashBasedShufflePartitioning,
+ OneToOnePartitioning,
+ Partitioning,
+ RoundRobinPartitioning,
+ RangeBasedShufflePartitioning,
+ BroadcastPartitioning,
+)
+
+
+class InputPortMaterializationReaderRunnable(Runnable, Stoppable):
+ def __init__(
+ self,
+ uri: str,
+ queue: InternalQueue,
+ worker_actor_id: ActorVirtualIdentity,
+ partitioning: Partitioning,
+ ):
+ """
+ Args:
+ uri (str): The URI of the materialized document.
+ queue: An instance of IQueue where messages are enqueued.
+ worker_actor_id (ActorVirtualIdentity): The target worker actor's identity.
+ partitioning: The partitioning information for this virtual reader worker
+ """
+ self.uri = uri
+ self.queue = queue
+ self.worker_actor_id = worker_actor_id
+ from_actor_id = get_from_actor_id_for_input_port_storage(
+ self.uri, self.worker_actor_id
+ )
+ self.channel_id = ChannelIdentity(
+ from_actor_id, self.worker_actor_id, is_control=False
+ )
+ self._stopped = False
+ self._finished = False
+ self.materialization = None
+ self.tuple_schema = None
+ self._partitioning_to_partitioner: dict[
+ type(Partitioning), type(Partitioner)
+ ] = {
+ OneToOnePartitioning: OneToOnePartitioner,
+ RoundRobinPartitioning: RoundRobinPartitioner,
+ HashBasedShufflePartitioning: HashBasedShufflePartitioner,
+ RangeBasedShufflePartitioning: RangeBasedShufflePartitioner,
+ BroadcastPartitioning: BroadcastPartitioner,
+ }
+ the_partitioning: Partitioning = get_one_of(partitioning)
+ partitioner = self._partitioning_to_partitioner[type(the_partitioning)]
+ self.partitioner: Partitioner = (
+ partitioner(the_partitioning)
+ if partitioner != OneToOnePartitioner
+ else partitioner(the_partitioning, self.worker_actor_id)
+ )
+
+ def finished(self) -> bool:
+ """
+ :return: Whether this reader thread has finished its logic.
+ """
+ return self._finished
+
+ def tuple_to_batch_with_filter(self, tuple_: Tuple) -> typing.Iterator[DataFrame]:
+ """
+ Let the partitioner produce batches to each (hypothetical) downstream
+ worker but only selects the worker that this thread is running on
+ as the input. This mimics the iterator logic of that in output
+ manager.
+ """
+ for receiver, tuples in self.partitioner.add_tuple_to_batch(tuple_):
+ if receiver == self.worker_actor_id:
+ yield self.tuples_to_data_frame(tuples)
+
+ def run(self) -> None:
+ """
+ Main execution logic that reads tuples from the materialized storage and
+ enqueues them in batches. It first emits a StartChannel ECM and, when finished,
+ emits an EndChannel ECM. Use the same partitioner implementation as that in
+ output manager, where a tuple is batched by the partitioner and only
+ selected as the input of this worker according to the partitioner.
+ """
+ try:
+ self.materialization, self.tuple_schema = DocumentFactory.open_document(
+ self.uri
+ )
+ self.emit_ecm("StartChannel", EmbeddedControlMessageType.NO_ALIGNMENT)
+ storage_iterator = self.materialization.get()
+
+ # Iterate and process tuples.
+ for tup in storage_iterator:
+ if self._stopped:
+ break
+ # Each tuple is sent to the partitioner and converted to
+ # a batch-based iterator.
+ tup.cast_to_schema(self.tuple_schema)
+ for data_frame in self.tuple_to_batch_with_filter(tup):
+ self.emit_payload(data_frame)
+ self.emit_ecm("EndChannel", EmbeddedControlMessageType.PORT_ALIGNMENT)
+ self._finished = True
+ except Exception as err:
+ logger.exception(err)
+
+ def stop(self):
+ """Sets the stop flag so the run loop may terminate."""
+ self._stopped = True
+
+ def emit_ecm(self, method_name: str, alignment: EmbeddedControlMessageType) -> None:
+ """
+ Emit an ECM (StartChannel or EndChannel), and
+ flush the remaining data batches if any. This mimics the
+ iterator logic of that in output manager.
+ """
+ ecm = EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity(method_name),
+ alignment,
+ [],
+ {
+ self.worker_actor_id.name: ControlInvocation(
+ method_name,
+ ControlRequest(empty_request=EmptyRequest()),
+ AsyncRpcContext(ActorVirtualIdentity(), ActorVirtualIdentity()),
+ -1,
+ )
+ },
+ )
+
+ for payload in self.partitioner.flush(self.worker_actor_id, ecm):
+ final_payload = (
+ payload
+ if isinstance(payload, EmbeddedControlMessage)
+ else self.tuples_to_data_frame(payload)
+ )
+ self.emit_payload(final_payload)
+
+ def emit_payload(self, payload: Union[DataPayload, EmbeddedControlMessage]) -> None:
+ """
+ Put the payload to the DP internal queue.
+ """
+ queue_element = (
+ ECMElement(tag=self.channel_id, payload=payload)
+ if isinstance(payload, EmbeddedControlMessage)
+ else DataElement(tag=self.channel_id, payload=payload)
+ )
+ self.queue.put(queue_element)
+
+ def tuples_to_data_frame(self, tuples: typing.List[Tuple]) -> DataFrame:
+ """
+ Converts a list of tuples to a DataFrame using pyarrow.Table.from_pydict
+ :param tuples:
+ :return:
+ """
+ return DataFrame(
+ frame=Table.from_pydict(
+ {
+ name: [t[name] for t in tuples]
+ for name in self.tuple_schema.get_attr_names()
+ },
+ schema=self.tuple_schema.as_arrow_schema(),
+ )
+ )
diff --git a/amber/src/main/python/core/storage/runnables/port_storage_writer.py b/amber/src/main/python/core/storage/runnables/port_storage_writer.py
new file mode 100644
index 00000000000..5a026162526
--- /dev/null
+++ b/amber/src/main/python/core/storage/runnables/port_storage_writer.py
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from dataclasses import dataclass
+from overrides import overrides
+
+from core.models import Tuple
+from core.storage.model.buffered_item_writer import BufferedItemWriter
+from core.util import StoppableQueueBlockingRunnable, IQueue
+from core.util.customized_queue.queue_base import QueueElement
+
+
+@dataclass
+class PortStorageWriterElement(QueueElement):
+ data_tuple: Tuple
+
+
+class PortStorageWriter(StoppableQueueBlockingRunnable):
+ def __init__(self, buffered_item_writer: BufferedItemWriter, queue: IQueue):
+ super().__init__(name=self.__class__.__name__, queue=queue)
+ self.buffered_item_writer: BufferedItemWriter = buffered_item_writer
+
+ @overrides
+ def receive(self, next_entry: QueueElement) -> None:
+ if isinstance(next_entry, PortStorageWriterElement):
+ self.buffered_item_writer.put_one(next_entry.data_tuple)
+ else:
+ raise TypeError(f"Unexpected entry {next_entry}")
+
+ @overrides
+ def pre_start(self) -> None:
+ self.buffered_item_writer.open()
+
+ @overrides
+ def post_stop(self) -> None:
+ self.buffered_item_writer.close()
diff --git a/amber/src/main/python/core/storage/storage_config.py b/amber/src/main/python/core/storage/storage_config.py
new file mode 100644
index 00000000000..0e47bdb71ae
--- /dev/null
+++ b/amber/src/main/python/core/storage/storage_config.py
@@ -0,0 +1,86 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+class StorageConfig:
+ """
+ A static class to keep the storage-related configs.
+ This class should be initialized with the configs passed from Java side and
+ is used by all storage-related classes.
+ """
+
+ _initialized = False
+
+ ICEBERG_CATALOG_TYPE = None
+ ICEBERG_POSTGRES_CATALOG_URI_WITHOUT_SCHEME = None
+ ICEBERG_POSTGRES_CATALOG_USERNAME = None
+ ICEBERG_POSTGRES_CATALOG_PASSWORD = None
+ ICEBERG_REST_CATALOG_URI = None
+ ICEBERG_REST_CATALOG_WAREHOUSE_NAME = None
+ ICEBERG_TABLE_RESULT_NAMESPACE = None
+ ICEBERG_FILE_STORAGE_DIRECTORY_PATH = None
+ ICEBERG_TABLE_COMMIT_BATCH_SIZE = None
+
+ # S3 configs
+ S3_ENDPOINT = None
+ S3_REGION = None
+ S3_AUTH_USERNAME = None
+ S3_AUTH_PASSWORD = None
+
+ @classmethod
+ def initialize(
+ cls,
+ catalog_type,
+ postgres_uri_without_scheme,
+ postgres_username,
+ postgres_password,
+ rest_catalog_uri,
+ rest_catalog_warehouse_name,
+ table_result_namespace,
+ directory_path,
+ commit_batch_size,
+ s3_endpoint,
+ s3_region,
+ s3_auth_username,
+ s3_auth_password,
+ ):
+ if cls._initialized:
+ raise RuntimeError(
+ "Storage config has already been initialized and cannot be modified."
+ )
+
+ cls.ICEBERG_CATALOG_TYPE = catalog_type
+ cls.ICEBERG_POSTGRES_CATALOG_URI_WITHOUT_SCHEME = postgres_uri_without_scheme
+ cls.ICEBERG_POSTGRES_CATALOG_USERNAME = postgres_username
+ cls.ICEBERG_POSTGRES_CATALOG_PASSWORD = postgres_password
+ cls.ICEBERG_REST_CATALOG_URI = rest_catalog_uri
+ cls.ICEBERG_REST_CATALOG_WAREHOUSE_NAME = rest_catalog_warehouse_name
+
+ cls.ICEBERG_TABLE_RESULT_NAMESPACE = table_result_namespace
+ cls.ICEBERG_FILE_STORAGE_DIRECTORY_PATH = directory_path
+ cls.ICEBERG_TABLE_COMMIT_BATCH_SIZE = int(commit_batch_size)
+
+ # S3 configs
+ cls.S3_ENDPOINT = s3_endpoint
+ cls.S3_REGION = s3_region
+ cls.S3_AUTH_USERNAME = s3_auth_username
+ cls.S3_AUTH_PASSWORD = s3_auth_password
+
+ cls._initialized = True
+
+ def __new__(cls, *args, **kwargs):
+ raise TypeError(f"{cls.__name__} is a static class and cannot be instantiated.")
diff --git a/amber/src/main/python/core/storage/vfs_uri_factory.py b/amber/src/main/python/core/storage/vfs_uri_factory.py
new file mode 100644
index 00000000000..de0c5db56ec
--- /dev/null
+++ b/amber/src/main/python/core/storage/vfs_uri_factory.py
@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from enum import Enum
+from typing import Optional
+from urllib.parse import urlparse
+
+from core.util.virtual_identity import (
+ serialize_global_port_identity,
+ deserialize_global_port_identity,
+)
+from proto.org.apache.texera.amber.core import (
+ WorkflowIdentity,
+ ExecutionIdentity,
+ GlobalPortIdentity,
+)
+
+
+class VFSResourceType(str, Enum):
+ RESULT = "result"
+ RUNTIME_STATISTICS = "runtimeStatistics"
+ CONSOLE_MESSAGES = "consoleMessages"
+
+
+class VFSURIFactory:
+ VFS_FILE_URI_SCHEME = "vfs"
+
+ @staticmethod
+ def decode_uri(
+ uri: str,
+ ) -> (
+ WorkflowIdentity,
+ ExecutionIdentity,
+ Optional[GlobalPortIdentity],
+ VFSResourceType,
+ ):
+ """
+ Parses a VFS URI and extracts its components.
+ """
+ parsed_uri = urlparse(uri)
+
+ if parsed_uri.scheme != VFSURIFactory.VFS_FILE_URI_SCHEME:
+ raise ValueError(f"Invalid URI scheme: {parsed_uri.scheme}")
+
+ segments = parsed_uri.path.lstrip("/").split("/")
+
+ def extract_value(key: str) -> str:
+ try:
+ index = segments.index(key)
+ return segments[index + 1]
+ except (ValueError, IndexError):
+ raise ValueError(f"Missing value for key: {key} in URI: {uri}")
+
+ workflow_id = WorkflowIdentity(int(extract_value("wid")))
+ execution_id = ExecutionIdentity(int(extract_value("eid")))
+
+ global_port_id = (
+ deserialize_global_port_identity(extract_value("globalportid"))
+ if "globalportid" in segments
+ else None
+ )
+
+ resource_type_str = segments[-1].lower()
+ try:
+ resource_type = VFSResourceType(resource_type_str)
+ except ValueError:
+ raise ValueError(f"Unknown resource type: {resource_type_str}")
+
+ return (
+ workflow_id,
+ execution_id,
+ global_port_id,
+ resource_type,
+ )
+
+ @staticmethod
+ def create_result_uri(workflow_id, execution_id, global_port_id) -> str:
+ """Creates a URI pointing to a result storage."""
+ base_uri = (
+ f"{VFSURIFactory.VFS_FILE_URI_SCHEME}:///wid/{workflow_id.id}"
+ f"/eid/{execution_id.id}/globalportid/"
+ f"{serialize_global_port_identity(global_port_id)}"
+ )
+
+ return f"{base_uri}/{VFSResourceType.RESULT.value}"
diff --git a/amber/src/main/python/core/util/__init__.py b/amber/src/main/python/core/util/__init__.py
new file mode 100644
index 00000000000..b747d28f4ba
--- /dev/null
+++ b/amber/src/main/python/core/util/__init__.py
@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .proto import get_one_of, set_one_of
+from .customized_queue import LinkedBlockingMultiQueue, IQueue
+from .stoppable import Stoppable, StoppableQueueBlockingRunnable
+
+__all__ = [
+ "get_one_of",
+ "set_one_of",
+ "LinkedBlockingMultiQueue",
+ "IQueue",
+ "StoppableQueueBlockingRunnable",
+ "Stoppable",
+]
diff --git a/amber/src/main/python/core/util/buffer/buffer_base.py b/amber/src/main/python/core/util/buffer/buffer_base.py
new file mode 100644
index 00000000000..244b7357407
--- /dev/null
+++ b/amber/src/main/python/core/util/buffer/buffer_base.py
@@ -0,0 +1,24 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import ABCMeta
+
+from core.util.protocol.base_protocols import FlushedGetable, Putable
+
+
+class IBuffer(FlushedGetable, Putable, metaclass=ABCMeta):
+ pass
diff --git a/amber/src/main/python/core/util/buffer/timed_buffer.py b/amber/src/main/python/core/util/buffer/timed_buffer.py
new file mode 100644
index 00000000000..abaacb15705
--- /dev/null
+++ b/amber/src/main/python/core/util/buffer/timed_buffer.py
@@ -0,0 +1,44 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from datetime import datetime
+from typing import List, Iterator
+
+from core.util.buffer.buffer_base import IBuffer
+from proto.org.apache.texera.amber.engine.architecture.rpc import ConsoleMessage
+
+
+class TimedBuffer(IBuffer):
+ def __init__(self, max_message_num=10, max_flush_interval_in_ms=500):
+ self._max_message_num = max_message_num
+ self._max_flush_interval_in_ms = max_flush_interval_in_ms
+ self._buffer: List[ConsoleMessage]() = list()
+ self._last_output_time = datetime.now()
+
+ def put(self, message: ConsoleMessage) -> None:
+ self._buffer.append(message)
+
+ def get(self, flush: bool = False) -> Iterator[ConsoleMessage]:
+ if (
+ flush
+ or len(self._buffer) >= self._max_message_num
+ or (datetime.now() - self._last_output_time).seconds
+ >= self._max_flush_interval_in_ms / 1000
+ ):
+ self._last_output_time = datetime.now()
+ yield from self._buffer
+ self._buffer.clear()
diff --git a/amber/src/main/python/core/util/console_message/replace_print.py b/amber/src/main/python/core/util/console_message/replace_print.py
new file mode 100644
index 00000000000..7feeeeb52d0
--- /dev/null
+++ b/amber/src/main/python/core/util/console_message/replace_print.py
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import builtins
+import inspect
+from contextlib import redirect_stdout
+from io import StringIO
+from typing import ContextManager
+
+from core.util.buffer.buffer_base import IBuffer
+from core.util.console_message.timestamp import current_time_in_local_timezone
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ ConsoleMessage,
+ ConsoleMessageType,
+)
+
+
+class replace_print(ContextManager):
+ """
+ A context manager to support replace builtin print function.
+
+ With in the context, we use a customized print function which does the following:
+ 1. writes to a given buffer instead of stdout
+ 2. writes as a complete string, which is made of joining of all stringify-ed
+ arguments and the end argument of the original print function. It calls the
+ buf.write once per print call, which is different from
+ contextlib.redirect_stdout who calls the buf.write for each argument in the
+ print function.
+ """
+
+ def __init__(self, worker_id: str, buf: IBuffer):
+ # save a reference to the original builtin.print before we replace it.
+ # it will always replace back when the context manager exits, with exception
+ # or not.
+ self.builtins_print = builtins.print
+ self.worker_id = worker_id
+ self.buf = buf # the provided buffer to write to
+
+ def __enter__(self) -> None:
+ """
+ Enters the context, replace builtin.print function with a wrapped function.
+ Now we hard code the wrapped_print to output complete print result to the
+ given buffer.
+ :return:
+ """
+
+ def wrapped_print(*args, **kwargs):
+ # use StringIO to obtain the written complete string from the original
+ # print function.
+ if "file" in kwargs:
+ self.builtins_print(*args, **kwargs)
+ return
+ with StringIO() as tmp_buf, redirect_stdout(tmp_buf):
+ self.builtins_print(*args, **kwargs)
+ complete_str = tmp_buf.getvalue()
+ console_message = ConsoleMessage(
+ worker_id=self.worker_id,
+ timestamp=current_time_in_local_timezone(),
+ msg_type=ConsoleMessageType.PRINT,
+ source=(
+ f"{inspect.currentframe().f_back.f_globals['__name__']}"
+ f":{inspect.currentframe().f_back.f_code.co_name}"
+ f":{inspect.currentframe().f_back.f_lineno}"
+ ),
+ title=complete_str,
+ message="",
+ )
+ self.buf.put(console_message)
+
+ builtins.print = wrapped_print
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+ """
+ Exits the context, revert the replacement to recover the original
+ builtin.print function.
+
+ It does not handle exception within the context, it simply raises it outside
+ the context.
+
+ :param exc_type: potential exception type.
+ :param exc_val: potential exception value.
+ :param exc_tb: potential exception traceback.
+ :return: bool, if no exception was raised, return True, otherwise,
+ return False.
+ """
+ builtins.print = self.builtins_print
+ return exc_val is None
diff --git a/amber/src/main/python/core/util/console_message/timed_buffer.py b/amber/src/main/python/core/util/console_message/timed_buffer.py
new file mode 100644
index 00000000000..59369f16b21
--- /dev/null
+++ b/amber/src/main/python/core/util/console_message/timed_buffer.py
@@ -0,0 +1,45 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from datetime import datetime
+from typing import Tuple, List, Iterator
+
+from proto.org.apache.texera.amber.engine.architecture.worker import (
+ PythonConsoleMessageV2,
+)
+
+
+class TimedBuffer:
+ def __init__(self, max_message_num=10, max_flush_interval_in_ms=500):
+ self._max_message_num = max_message_num
+ self._max_flush_interval_in_ms = max_flush_interval_in_ms
+ self._buffer: List[Tuple[datetime, str]]() = list()
+ self._last_output_time = datetime.now()
+
+ def add(self, console_message: PythonConsoleMessageV2) -> None:
+ self._buffer.append(console_message)
+
+ def get(self, flush=False) -> Iterator[PythonConsoleMessageV2]:
+ if (
+ flush
+ or len(self._buffer) >= self._max_message_num
+ or (datetime.now() - self._last_output_time).seconds
+ >= self._max_flush_interval_in_ms / 1000
+ ):
+ self._last_output_time = datetime.now()
+ yield from self._buffer
+ self._buffer.clear()
diff --git a/amber/src/main/python/core/util/console_message/timestamp.py b/amber/src/main/python/core/util/console_message/timestamp.py
new file mode 100644
index 00000000000..9ced3a712bc
--- /dev/null
+++ b/amber/src/main/python/core/util/console_message/timestamp.py
@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import tzlocal
+
+
+def current_time_in_local_timezone():
+ # Get the system's local timezone
+ local_timezone = tzlocal.get_localzone()
+
+ # Get the current time in the local timezone
+ local_time = datetime.datetime.now(local_timezone)
+
+ return local_time
diff --git a/amber/src/main/python/core/util/customized_queue/__init__.py b/amber/src/main/python/core/util/customized_queue/__init__.py
new file mode 100644
index 00000000000..0e4c20e8c60
--- /dev/null
+++ b/amber/src/main/python/core/util/customized_queue/__init__.py
@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .linked_blocking_multi_queue import LinkedBlockingMultiQueue
+from .queue_base import IQueue
+
+__all__ = ["LinkedBlockingMultiQueue", "IQueue"]
diff --git a/amber/src/main/python/core/util/customized_queue/double_blocking_queue.py b/amber/src/main/python/core/util/customized_queue/double_blocking_queue.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/util/customized_queue/double_blocking_queue.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/util/customized_queue/inner.py b/amber/src/main/python/core/util/customized_queue/inner.py
new file mode 100644
index 00000000000..b39762f4f83
--- /dev/null
+++ b/amber/src/main/python/core/util/customized_queue/inner.py
@@ -0,0 +1,154 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+This class is taken from https://github.com/sebkeim/inner-class at commit sha
+0856be1feee38710005a7ef27ae998af95dedbf8.
+"""
+
+from functools import update_wrapper
+
+
+# TODO: re-arrange to another module
+def raw_inner(x):
+ """do nothing decorator for future backward compatibility :
+ this will preserve current behavior for inner-class if a future version
+ of the language change the default semantic for inner classes"""
+ return x
+
+
+class static_inner:
+ """decorator for outer attribute"""
+
+ def __init__(self, cls):
+ self.icls = cls
+ self.__doc__ = cls.__doc__
+
+ def __set_name__(self, owner, name):
+ self.icls.owner = owner
+ # now that outer is set, replace decorator by the actual class
+ setattr(owner, name, self.icls)
+
+
+class class_inner(static_inner):
+ """decorator for outer attribute, inner derivation and carried inheritance"""
+
+ def __init__(self, cls):
+ for method in ("__get__", "__set__", "__del__"):
+ if hasattr(cls, method):
+ raise ValueError("descriptors can't be used as inner class")
+ static_inner.__init__(self, cls)
+
+ def _innerparents(self, outercls):
+ mro = self.icls.mro()
+ name = self.name
+ innerparents = []
+ for parent in outercls.__bases__:
+ try:
+ innerparent = getattr(parent, name)
+ except AttributeError:
+ pass
+ else:
+ if innerparent not in mro:
+ innerparents.append(innerparent)
+ return tuple(innerparents)
+
+ def __set_name__(self, owner, name):
+ # inner derivation
+ self.name = name
+
+ bases = self._innerparents(owner)
+ if bases:
+ selfbases = self.icls.__bases__
+ if selfbases != (object,):
+ bases = selfbases + bases
+ self.icls = type(self.icls)(
+ self.icls.__name__,
+ bases,
+ dict(self.icls.__dict__),
+ )
+ assert "outer" not in self.icls.__dict__
+ self.icls.owner = owner
+
+ def __get__(self, outerobj, outercls):
+ # carried ineritence
+ cls = self.icls
+ if cls.owner != outercls:
+ assert self.name not in outercls.__dict__
+
+ bases = (self.icls,) + self._innerparents(outercls)
+ cls = type(cls)(
+ self.name,
+ bases,
+ {
+ "owner": outercls,
+ "__qualname__": outercls.__name__ + "." + self.name,
+ "__module__": cls.__module__,
+ "__doc__": cls.__doc__,
+ # '__annotations__':cls.__annotations__
+ },
+ )
+
+ inner = type(self)(cls)
+ inner.name = self.name
+ setattr(outercls, self.name, inner)
+ return cls
+
+
+class inner(class_inner):
+ """decorator for outer object attribute, inner derivation, carried inheritance
+ and instance"""
+
+ is_property = False
+ is_cached = False
+
+ @classmethod
+ def property(cls, icls):
+ """replicate standard @property decorator"""
+ obj = cls(icls)
+ obj.is_property = True
+ return obj
+
+ @classmethod
+ def cached_property(cls, icls):
+ """replicate sdtlib @cached_property decorator"""
+ obj = cls(icls)
+ obj.is_property = True
+ obj.is_cached = True
+ return obj
+
+ def __get__(self, outerobj, outercls):
+ icls = class_inner.__get__(self, outerobj, outercls)
+ if outerobj is None:
+ return icls
+ # properties
+ if self.is_property:
+ innerobj = icls()
+ innerobj.owner = outerobj
+ if self.is_cached:
+ setattr(outerobj, self.name, innerobj)
+ return innerobj
+
+ # constructor
+ def ctor(*args, **kw):
+ innerobj = icls.__new__(icls, *args, **kw)
+ innerobj.owner = outerobj
+ innerobj.__init__(*args, **kw)
+ return innerobj
+
+ update_wrapper(ctor, icls.__init__)
+ return ctor
diff --git a/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py b/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py
new file mode 100644
index 00000000000..3b46e6db4d7
--- /dev/null
+++ b/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py
@@ -0,0 +1,473 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import sys
+from threading import RLock, Condition
+from typing import List, Optional, Generic, TypeVar, MutableMapping
+
+from core.util.customized_queue.inner import inner
+from core.util.customized_queue.queue_base import IKeyedQueue
+from core.util.thread.atomic import AtomicInteger
+
+K = TypeVar("K")
+T = TypeVar("T")
+
+
+class LinkedBlockingMultiQueue(IKeyedQueue):
+ @inner
+ class Node(Generic[T]):
+ def __init__(self, item: T):
+ self.item = item
+ self.next: Optional[LinkedBlockingMultiQueue.Node[T]] = None
+ self.in_mem_size = sys.getsizeof(item)
+
+ @inner
+ class SubQueue(Generic[T]):
+ def __init__(self, key: K):
+ self.key: K = key
+ self.priority_group: Optional[LinkedBlockingMultiQueue.PriorityGroup] = None
+ self.put_lock: RLock = RLock()
+ self.count: AtomicInteger = AtomicInteger()
+ self.enabled: bool = True
+ self.in_mem_size = AtomicInteger()
+ self.head: LinkedBlockingMultiQueue.Node = LinkedBlockingMultiQueue.Node(
+ None
+ )
+ self.last: Optional[LinkedBlockingMultiQueue.Node[T]] = self.head
+
+ def clear(self) -> None:
+ self.fully_lock()
+ try:
+ h: LinkedBlockingMultiQueue.Node[T] = self.head
+ p: LinkedBlockingMultiQueue.Node = h.next
+ while p is not None:
+ h.next = h
+ p.item = None
+ h = p
+ p = h.next
+ self.head = self.last
+ old_count = self.count.get_and_set(0)
+ if self.enabled:
+ self.owner.total_count.get_and_dec(old_count)
+ finally:
+ self.fully_unlock()
+
+ def disable(self) -> None:
+ self.fully_lock()
+ try:
+ if not self.enabled:
+ return
+ self.owner.total_count.dec(self.count.value)
+ self.enabled = False
+ finally:
+ self.fully_unlock()
+
+ def enable(self) -> None:
+ self.fully_lock()
+ try:
+ if self.enabled:
+ return
+ self.enabled = True
+
+ # potentially unlock waiting polls
+ c = self.count.value
+ if c > 0:
+ self.owner.total_count.inc(c)
+ self.owner.not_empty.notify()
+
+ finally:
+ self.fully_unlock()
+
+ def is_enabled(self) -> bool:
+ self.owner.take_lock.acquire()
+ try:
+ return self.enabled
+ finally:
+ self.owner.take_lock.release()
+
+ def enqueue(self, node: LinkedBlockingMultiQueue.Node[T]) -> None:
+ self.last.next = node
+ self.last = node
+ self.in_mem_size.inc(node.in_mem_size)
+
+ def dequeue(self) -> T:
+ assert self.size() > 0
+ h = self.head
+ first = h.next
+ h.next = h
+ self.head = first
+ x = first.item
+ self.in_mem_size.dec(first.in_mem_size)
+ first.item = None
+ return x
+
+ def __str__(self) -> str:
+ res = ""
+ h = self.head
+ while h.next is not None:
+ res += h.next.item
+ res += " -> "
+ h = h.next
+ return res
+
+ def size(self) -> int:
+ return self.count.value
+
+ def is_empty(self) -> bool:
+ return self.size() == 0
+
+ def put(self, obj: T) -> None:
+ if obj is None:
+ raise ValueError("Does not support NoneType.")
+ old_size = -1
+ node = LinkedBlockingMultiQueue.Node(obj)
+ self.put_lock.acquire()
+ try:
+ self.enqueue(node)
+ self.count.inc()
+ if self.enabled:
+ old_size = self.owner.total_count.get_and_inc()
+ finally:
+ self.put_lock.release()
+
+ if old_size == 0:
+ self.owner._signal_not_empty()
+
+ def remove(self, obj: T) -> bool:
+ if obj is None:
+ return False
+ self.fully_lock()
+ try:
+ trail = self.head
+ while trail.next is not None:
+ if trail.item == obj:
+ self.unlink(trail, trail.next)
+ return True
+ trail = trail.next
+ return False
+ finally:
+ self.fully_unlock()
+
+ def unlink(
+ self,
+ trail: LinkedBlockingMultiQueue.Node,
+ next_: LinkedBlockingMultiQueue.Node,
+ ) -> None:
+ trail.item = None
+ trail.next = next_.next
+ if self.last == next_:
+ self.last = trail
+ if self.enabled:
+ self.owner.total_count.get_and_dec()
+
+ def fully_lock(self) -> None:
+ self.put_lock.acquire()
+ self.owner.take_lock.acquire()
+
+ def fully_unlock(self) -> None:
+ self.put_lock.release()
+ self.owner.take_lock.release()
+
+ @inner
+ class PriorityGroup(Generic[T]):
+ def __init__(self, priority: int = 0):
+ # non-negative number, the smaller number means higher priority.
+ self.priority: int = priority
+ self.queues: List[LinkedBlockingMultiQueue.SubQueue[T]] = list()
+ self.next_idx: int = 0
+
+ def add_queue(self, to_add: LinkedBlockingMultiQueue.SubQueue[T]) -> None:
+ self.queues.append(to_add)
+ to_add.priority_group = self
+
+ def remove_queue(self, to_remove: LinkedBlockingMultiQueue.SubQueue[T]) -> None:
+ for queue in self.queues:
+ if queue.key == to_remove.key:
+ to_remove.put_lock.acquire()
+ try:
+ self.queues[:] = [q for q in self.queues if q != queue]
+ if self.next_idx == len(self.queues):
+ self.next_idx = 0
+ if queue.enabled:
+ self.owner.total_count.get_and_dec(to_remove.size())
+ finally:
+ to_remove.put_lock.release()
+
+ def get_next_sub_queue(self) -> Optional[LinkedBlockingMultiQueue.SubQueue[T]]:
+ start_idx = self.next_idx
+ queues = [q for q in self.queues]
+ while True:
+ child = queues[self.next_idx]
+ self.next_idx += 1
+ if self.next_idx == len(queues):
+ self.next_idx = 0
+ if child.enabled and child.size() > 0:
+ return child
+ if self.next_idx == start_idx:
+ break
+ return None
+
+ def peek(self) -> Optional[T]:
+ start_idx = self.next_idx
+ while True:
+ child = self.queues[self.next_idx]
+ if child.enabled and child.size() > 0:
+ return child.head.next.item
+ else:
+ self.next_idx += 1
+ if self.next_idx == len(self.queues):
+ self.next_idx = 0
+ if self.next_idx == start_idx:
+ break
+ return None
+
+ @inner
+ class DefaultSubQueueSelection(Generic[T]):
+ def __init__(
+ self, priority_groups: List[LinkedBlockingMultiQueue.PriorityGroup[T]]
+ ):
+ self.priority_groups: List[LinkedBlockingMultiQueue.PriorityGroup[T]] = (
+ priority_groups
+ )
+
+ def get_next(self) -> Optional[LinkedBlockingMultiQueue.SubQueue[T]]:
+ for pg in self.priority_groups:
+ sub_queue = pg.get_next_sub_queue()
+ if sub_queue is not None:
+ return sub_queue
+ return None
+
+ def peek(self) -> Optional[T]:
+ for pg in self.priority_groups:
+ deque = pg.peek()
+ if deque is not None:
+ return deque
+ return None
+
+ def set_priority_groups(
+ self, priority_groups: List[LinkedBlockingMultiQueue.PriorityGroup[T]]
+ ) -> None:
+ self.priority_groups = priority_groups
+
+ def __init__(self):
+ self.take_lock: RLock = RLock()
+ self.not_empty: Condition = Condition(self.take_lock)
+
+ # thread-safe in CPython
+ self.sub_queues: MutableMapping[K, LinkedBlockingMultiQueue.SubQueue] = dict()
+
+ # the count of the queue, describing how many element are getable;
+ # disabled subqueues will not be included in this count
+ self.total_count = AtomicInteger()
+
+ # thread-safe in CPython
+ self.priority_groups: List[LinkedBlockingMultiQueue.PriorityGroup] = list()
+ self.sub_queue_selection = LinkedBlockingMultiQueue.DefaultSubQueueSelection(
+ self.priority_groups
+ )
+
+ def in_mem_size(self, key: K) -> int:
+ return self.sub_queues[key].in_mem_size.value
+
+ def put(self, key: K, item: T) -> None:
+ """
+ Put one item into the SubQueue specified by the key.
+
+ :param key: the identifier of a SubQueue.
+ :param item: Any instance.
+ :raises KeyError for non-existing keys.
+ :return: None
+ """
+ self.get_sub_queue(key).put(item)
+
+ def get(self) -> T:
+ """
+ Blocking get the next available item from the queue.
+ - Disabled SubQueues are considered empty and will not be fetched.
+ - When multiple SubQueues are enabled and have items, it selects the SubQueue
+ by the order specified by the self.sub_queue_selection strategy.
+
+ :return: T, Any item that is available to the fetched.
+ """
+ self.take_lock.acquire()
+ try:
+ while self.total_count.value == 0:
+ self.not_empty.wait()
+
+ # at this point we know there is an element
+ sub_queue = self.sub_queue_selection.get_next()
+ item = sub_queue.dequeue()
+ sub_queue.count.dec()
+ if self.total_count.get_and_dec() > 1:
+ # sub queue still has element
+ self.not_empty.notify()
+ finally:
+ self.take_lock.release()
+
+ return item
+
+ def peek(self) -> Optional[T]:
+ """
+ Peek the next available item from the queue.
+ - When no item is available, it returns None.
+ - Otherwise, it acts the same as LinkedBlockingMultiQueue.get() but
+ without actually taking the item out from the queue.
+
+ :return: Optional[T], could be the available item or None.
+ """
+ self.take_lock.acquire()
+ try:
+ if self.total_count.value == 0:
+ return None
+ else:
+ return self.sub_queue_selection.peek()
+ finally:
+ self.take_lock.release()
+
+ def enable(self, key: K) -> None:
+ """
+ Enables a SubQueue, specified by key. This action acquires all locks.
+
+ :param key: the identifier of the SubQueue.
+ :raises KeyError for non-existing keys.
+ :return: None
+ """
+ self.get_sub_queue(key).enable()
+
+ def disable(self, key: K) -> None:
+ """
+ Disables a SubQueue, specified by key. This action acquires all locks.
+
+ :param key: the identifier of the SubQueue.
+ :raises KeyError for non-existing keys.
+ :return: None
+ """
+ self.get_sub_queue(key).disable()
+
+ def size(self, key: Optional[K] = None) -> int:
+ """
+ Get the total number of elements of all the SubQueues, or of a specific
+ SubQueue if a key is provided. This action acquires NO locks.
+
+ :param key: an optional identifier of a SubQueue.
+ If provided, give the size of the SubQueue.
+ Otherwise, return the total size of all SubQueues.
+ :raises KeyError for non-existing keys.
+ :return: Integer for size.
+ """
+ if key is not None:
+ return self.get_sub_queue(key).size()
+ else:
+ return self.total_count.value
+
+ def __len__(self) -> int:
+ return self.size()
+
+ def is_empty(self, key: Optional[K] = None) -> bool:
+ """
+ Check if the queue is empty, or check a specific SubQueue if
+ key is provided. This action acquires NO locks.
+
+ :param key: optional identifier of a SubQueue.
+ :raises KeyError for non-existing keys.
+ :return: Boolean representing empty or not.
+ """
+ return self.size(key) == 0
+
+ def is_enabled(self, key: K) -> bool:
+ return self.get_sub_queue(key).is_enabled()
+
+ def add_sub_queue(self, key: K, priority: int) -> Optional[SubQueue]:
+ """
+ Create a new SubQueue if absent, with the key and priority.
+
+ :param key: SubQueue identifier for future reference.
+ :param priority: int value of priority, the lower number means the higher
+ priority.
+ :return: returns None if the key is new, or returns the previous SubQueue
+ mapped by the key if key is repeated.
+ """
+ sub_queue = self.SubQueue(key)
+ self.take_lock.acquire()
+
+ try:
+ old_queue = self.sub_queues.get(key)
+ self.sub_queues[key] = sub_queue
+ if old_queue is None:
+ i = 0
+ added = False
+ for pg in self.priority_groups:
+ if pg.priority == priority:
+ pg.add_queue(sub_queue)
+ added = True
+ break
+ elif pg.priority > priority:
+ new_pg = LinkedBlockingMultiQueue.PriorityGroup(priority)
+ new_pg.add_queue(sub_queue)
+ self.priority_groups.append(new_pg)
+ added = True
+ break
+
+ i += 1
+ if not added:
+ new_pg = LinkedBlockingMultiQueue.PriorityGroup(priority)
+ new_pg.add_queue(sub_queue)
+ self.priority_groups.append(new_pg)
+
+ return old_queue
+ finally:
+ self.take_lock.release()
+
+ def remove_sub_queue(self, key: K) -> SubQueue:
+ self.take_lock.acquire()
+ try:
+ removed: Optional[LinkedBlockingMultiQueue.SubQueue] = self.sub_queues.get(
+ key
+ )
+ if removed is not None:
+ del self.sub_queues[key]
+ removed.priority_group.remove_queue(removed)
+ if len(removed.priority_group.queues) == 0:
+ self.priority_groups.remove(removed.priority_group)
+ return removed
+ finally:
+ self.take_lock.release()
+
+ def get_sub_queue(self, key: K) -> SubQueue:
+ """
+ Get the SubQueue specified by the key.
+
+ :param key: the identifier of a SubQueue.
+ :raises KeyError for non-existing keys.
+ :return: the SubQueue.
+ """
+ return self.sub_queues[key]
+
+ def _signal_not_empty(self) -> None:
+ """
+ Notifies a (the next) consumer that the queue is not empty.
+ Should only be invoked by the producer.
+
+ :return: None
+ """
+ self.take_lock.acquire()
+ try:
+ self.not_empty.notify()
+ finally:
+ self.take_lock.release()
diff --git a/amber/src/main/python/core/util/customized_queue/queue_base.py b/amber/src/main/python/core/util/customized_queue/queue_base.py
new file mode 100644
index 00000000000..47b8aac94e4
--- /dev/null
+++ b/amber/src/main/python/core/util/customized_queue/queue_base.py
@@ -0,0 +1,45 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import ABCMeta
+from dataclasses import dataclass
+
+from core.util.protocol.base_protocols import (
+ Putable,
+ Getable,
+ EmtpyCheckable,
+ KeyedPutable,
+ KeyedEmtpyCheckable,
+)
+
+
+@dataclass
+class QueueElement:
+ pass
+
+
+@dataclass
+class QueueControl(QueueElement):
+ msg: str
+
+
+class IQueue(Putable, Getable, EmtpyCheckable, metaclass=ABCMeta):
+ pass
+
+
+class IKeyedQueue(KeyedPutable, Getable, KeyedEmtpyCheckable, metaclass=ABCMeta):
+ pass
diff --git a/amber/src/main/python/core/util/customized_queue/test_linked_blocking_multi_queue.py b/amber/src/main/python/core/util/customized_queue/test_linked_blocking_multi_queue.py
new file mode 100644
index 00000000000..1df9423b69b
--- /dev/null
+++ b/amber/src/main/python/core/util/customized_queue/test_linked_blocking_multi_queue.py
@@ -0,0 +1,237 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+import random
+import time
+from threading import Thread
+
+from core.util.customized_queue.linked_blocking_multi_queue import (
+ LinkedBlockingMultiQueue,
+)
+
+
+class TestLinkedBlockingMultiQueue:
+ @pytest.fixture
+ def queue(self):
+ lbmq = LinkedBlockingMultiQueue()
+ lbmq.add_sub_queue("control", 0)
+ lbmq.add_sub_queue("data", 1)
+ return lbmq
+
+ def test_sub_can_emit(self, queue):
+ assert queue.is_empty()
+ queue.put("data", 1)
+ assert not queue.is_empty()
+ assert queue.is_empty("control")
+ assert queue.get() == 1
+ assert queue.is_empty()
+ assert queue.is_empty("control")
+
+ def test_main_can_emit(self, queue):
+ assert queue.is_empty()
+ queue.put("control", "s")
+ assert not queue.is_empty()
+ assert queue.get() == "s"
+ assert queue.is_empty()
+
+ def test_main_can_emit_before_sub(self, queue):
+ assert queue.is_empty()
+ queue.put("data", 1)
+ queue.put("control", "s")
+ assert not queue.is_empty()
+ assert queue.get() == "s"
+ assert queue.is_empty("control")
+ assert not queue.is_empty()
+ assert queue.get() == 1
+ assert queue.is_empty()
+
+ def test_can_maintain_order_respectively(self, queue):
+ queue.put("data", 1)
+ queue.put("control", "s1")
+ queue.put("data", 99)
+ queue.put("control", "s2")
+ queue.put("control", "s3")
+ queue.put("data", 3)
+ queue.put("control", "s4")
+ res = list()
+ while not queue.is_empty():
+ res.append(queue.get())
+
+ assert res == ["s1", "s2", "s3", "s4", 1, 99, 3]
+
+ def test_can_disable_sub(self, queue):
+ queue.disable("data")
+ queue.put("data", 1)
+ queue.put("control", "s1")
+ queue.put("data", 99)
+ queue.put("control", "s2")
+ queue.put("control", "s3")
+ queue.put("data", 3)
+ queue.put("control", "s4")
+ res = list()
+ while not queue.is_empty():
+ res.append(queue.get())
+
+ assert res == ["s1", "s2", "s3", "s4"]
+ assert queue.is_empty()
+ queue.enable("data")
+ assert not queue.is_empty()
+ res = list()
+ while not queue.is_empty():
+ res.append(queue.get())
+
+ assert res == [1, 99, 3]
+ assert queue.is_empty()
+
+ @pytest.mark.timeout(2)
+ def test_producer_first_insert_sub(self, queue, reraise):
+ def producer():
+ with reraise:
+ time.sleep(0.2)
+ queue.put("data", 1)
+
+ producer_thread = Thread(target=producer)
+ producer_thread.start()
+ producer_thread.join()
+ assert queue.get() == 1
+ reraise()
+
+ @pytest.mark.timeout(2)
+ def test_consumer_first_insert_sub(self, queue, reraise):
+ def consumer():
+ with reraise:
+ assert queue.get() == 1
+ assert queue.is_empty()
+
+ consumer_thread = Thread(target=consumer)
+ consumer_thread.start()
+ time.sleep(0.2)
+ queue.put("data", 1)
+ consumer_thread.join()
+ reraise()
+
+ @pytest.mark.timeout(2)
+ def test_producer_first_insert_main(self, queue, reraise):
+ def producer():
+ with reraise:
+ time.sleep(0.2)
+ queue.put("control", "s")
+
+ producer_thread = Thread(target=producer)
+ producer_thread.start()
+ producer_thread.join()
+ assert queue.get() == "s"
+ reraise()
+
+ @pytest.mark.timeout(2)
+ def test_consumer_first_insert_main(self, queue, reraise):
+ def consumer():
+ with reraise:
+ assert queue.get() == "s"
+ assert queue.is_empty()
+
+ consumer_thread = Thread(target=consumer)
+ consumer_thread.start()
+ time.sleep(0.2)
+ queue.put("control", "s")
+ consumer_thread.join()
+ reraise()
+
+ @pytest.mark.timeout(10)
+ def test_multiple_producer_race(self, queue, reraise):
+ def producer(k):
+ with reraise:
+ if isinstance(k, int):
+ for i in range(k):
+ queue.put("data", i)
+ else:
+ queue.put("control", k)
+
+ threads = []
+ target = set()
+ for i in range(1000):
+ if random.random() > 0.5:
+ i = chr(i)
+ target.add(i)
+ producer_thread = Thread(target=producer, args=(i,))
+ producer_thread.start()
+ threads.append(producer_thread)
+ res = set()
+
+ def consumer():
+ with reraise:
+ queue.disable("data")
+ while len(res) < len(target):
+ res.add(queue.get())
+
+ consumer_thread = Thread(target=consumer)
+ consumer_thread.start()
+ for thread in threads:
+ thread.join()
+
+ consumer_thread.join()
+ assert res == target
+
+ reraise()
+
+ def test_multi_types(
+ self,
+ queue,
+ ):
+ queue.put("data", 1)
+ queue.put("data", 1.1)
+ queue.put("control", "s")
+ queue.disable("data")
+ assert queue.get() == "s"
+ assert queue.is_empty()
+
+ @pytest.mark.timeout(2)
+ def test_common_single_producer_single_consumer(self, queue, reraise):
+ def producer():
+ with reraise:
+ for i in range(11):
+ if i % 3 == 0:
+ queue.put("control", "s")
+ else:
+ queue.put("data", i)
+
+ producer_thread = Thread(target=producer)
+ producer_thread.start()
+ producer_thread.join()
+
+ total: int = 0
+ while True:
+ queue.enable("data")
+ queue.enable("data")
+ t = queue.get()
+ queue.is_empty("control")
+
+ if isinstance(t, int):
+ total += t
+ else:
+ assert t == "s"
+ queue.is_empty()
+ queue.disable("data")
+ queue.disable("data")
+ queue.is_empty()
+ queue.disable("data")
+ if t == 10:
+ break
+ assert total == sum(filter(lambda x: x % 3 != 0, range(11)))
+
+ reraise()
diff --git a/amber/src/main/python/core/util/expression_evaluator/__init__.py b/amber/src/main/python/core/util/expression_evaluator/__init__.py
new file mode 100644
index 00000000000..b13e6a05191
--- /dev/null
+++ b/amber/src/main/python/core/util/expression_evaluator/__init__.py
@@ -0,0 +1,211 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import inspect
+import re
+from collections.abc import Iterator, Mapping
+from typing import Any, Dict, List, Optional, Pattern, Tuple
+
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EvaluatedValue,
+ TypedValue,
+)
+
+
+class ExpressionEvaluator:
+ """
+ Provides a series of static evaluation methods of a given expression, with an
+ optional context.
+ """
+
+ @staticmethod
+ def evaluate(
+ expression: str, runtime_context: Optional[Dict[str, Any]] = None
+ ) -> EvaluatedValue:
+ """
+ Evaluates the given expression and return a EvaluatedValue.
+
+ Right now, there is no validation performed on the input expression. User
+ takes full
+ responsibility of using this method.
+ :param expression: a python statement string
+ :param runtime_context: a Mapping of expressions to values, to be used for
+ evaluation
+ :return: EvaluatedValue which contains the current value and its children's
+ value, all in the format of TypedValue.
+
+ A TypedValue contains:
+ - expression: str, to match the request expression being evaluated;
+ - value_ref: str, the reference of this value, can be used to
+ construct the next expression which expands the current value
+ further;
+ - value_str: str, the value in string format, to be displayed;
+ - value_type: str, the type of this value, in string format,
+ to be displayed;
+ - expandable: bool, whether this value can be expanded or not.
+
+ The TypedValue could be expanded. For now it supports the following types:
+ - Primitives (expandable = False);
+ - Collections
+ - Array/Tuple like (expandable = True);
+ - Dict/Mapping like (expandable = True);
+ - Set like (expandable = True, but its elements' expandable = False);
+ - Iterables (expandable = True);
+ - Iterators (expandable = False);
+ - Generators (expandable = True).
+
+ See test cases for more usage details.
+ """
+
+ value = eval(expression, runtime_context)
+ value_str = repr(value)
+ type_str = type(value).__name__
+
+ to_be_expanded = list()
+
+ if ExpressionEvaluator._has_attributes(value):
+ to_be_expanded += ExpressionEvaluator._extract_attributes(value)
+
+ if ExpressionEvaluator._is_iterable(value):
+ if ExpressionEvaluator._is_generator(value):
+ to_be_expanded += ExpressionEvaluator._extract_generator_locals(value)
+ elif ExpressionEvaluator._is_iterator(value):
+ pass
+ else:
+ to_be_expanded += ExpressionEvaluator._extract_container_items(value)
+
+ return EvaluatedValue(
+ value=TypedValue(
+ expression=expression,
+ value_ref=expression,
+ value_str=value_str,
+ value_type=type_str,
+ expandable=ExpressionEvaluator._is_expandable(value),
+ ),
+ attributes=to_be_expanded,
+ )
+
+ @staticmethod
+ def _has_attributes(value: Any) -> bool:
+ return hasattr(value, "__dict__")
+
+ @staticmethod
+ def _is_expandable(obj, parent=None) -> bool:
+ # for set and set-like subclasses, the internal values cannot be expanded
+ # easily, disable for now
+ return (
+ not isinstance(parent, set)
+ and not (
+ ExpressionEvaluator._is_iterator(obj)
+ and not ExpressionEvaluator._is_generator(obj)
+ )
+ and (
+ ExpressionEvaluator._contains_attributes(obj)
+ or (
+ ExpressionEvaluator._is_iterable(obj)
+ and not ExpressionEvaluator._is_empty_container(obj)
+ )
+ )
+ )
+
+ @staticmethod
+ def _is_mapping(obj) -> bool:
+ return isinstance(obj, Mapping)
+
+ @staticmethod
+ def _is_generator(obj) -> bool:
+ return inspect.isgenerator(obj)
+
+ @staticmethod
+ def _is_iterator(obj) -> bool:
+ return isinstance(obj, Iterator)
+
+ @staticmethod
+ def _is_iterable(obj) -> bool:
+ """
+ According to
+ https://www.pythonlikeyoumeanit.com/Module2_EssentialsOfPython/Iterables.html#Iterables,
+ an iterable is any Python object with an __iter__() method or with a
+ __getitem__() method that implements Sequence semantics.
+ """
+ return hasattr(obj, "__iter__") or hasattr(obj, "__getitem__")
+
+ @staticmethod
+ def _contains_attributes(obj) -> bool:
+ return hasattr(obj, "__dict__") and len(obj.__dict__) > 0
+
+ @staticmethod
+ def _is_empty_container(obj) -> bool:
+ return hasattr(obj, "__len__") and len(obj) == 0
+
+ @staticmethod
+ def _contextualize_expression(
+ expression: str, context_replacements: Dict[Pattern[str], str]
+ ) -> str:
+ contextualized_expression = expression
+ for pattern, contextualized_pattern in context_replacements.items():
+ contextualized_expression = re.sub(
+ pattern, contextualized_pattern, contextualized_expression
+ )
+ return contextualized_expression
+
+ @staticmethod
+ def _extract_container_items(value: Any) -> List[TypedValue]:
+ return ExpressionEvaluator._to_typed_values(
+ (
+ value.items()
+ if ExpressionEvaluator._is_mapping(value)
+ else enumerate(value)
+ ),
+ parent=value,
+ to_getitem=True,
+ ref_as_repr=True,
+ )
+
+ @staticmethod
+ def _extract_attributes(value: Any) -> List[TypedValue]:
+ return ExpressionEvaluator._to_typed_values(vars(value).items())
+
+ @staticmethod
+ def _extract_generator_locals(value: Any) -> List[TypedValue]:
+ return ExpressionEvaluator._to_typed_values(
+ filter(lambda t: t[0] != ".0", inspect.getgeneratorlocals(value).items()),
+ check_expandable=False,
+ )
+
+ @staticmethod
+ def _to_typed_values(
+ kv_iter: List[Tuple[str, Any]],
+ parent=None,
+ to_getitem=False,
+ ref_as_repr=False,
+ check_expandable=True,
+ ):
+ return [
+ TypedValue(
+ expression=f"__getitem__({repr(k)})" if to_getitem else k,
+ value_ref=repr(k) if ref_as_repr else k,
+ value_str=repr(v),
+ value_type=type(v).__name__,
+ expandable=(
+ ExpressionEvaluator._is_expandable(v, parent=parent)
+ if check_expandable
+ else False
+ ),
+ )
+ for k, v in kv_iter
+ ]
diff --git a/amber/src/main/python/core/util/expression_evaluator/test_expression_evaluator.py b/amber/src/main/python/core/util/expression_evaluator/test_expression_evaluator.py
new file mode 100644
index 00000000000..19340bdccda
--- /dev/null
+++ b/amber/src/main/python/core/util/expression_evaluator/test_expression_evaluator.py
@@ -0,0 +1,480 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.util.expression_evaluator import ExpressionEvaluator
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EvaluatedValue,
+ TypedValue,
+)
+
+
+class TestExpressionEvaluator:
+ def test_evaluate_basic_expressions(self):
+ i = 10
+ assert ExpressionEvaluator.evaluate(
+ "i", runtime_context={"i": i}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="i",
+ value_ref="i",
+ value_str="10",
+ value_type="int",
+ expandable=False,
+ ),
+ attributes=[],
+ )
+
+ f = 1.1
+ assert ExpressionEvaluator.evaluate(
+ "f", runtime_context={"f": f}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="f",
+ value_ref="f",
+ value_str="1.1",
+ value_type="float",
+ expandable=False,
+ ),
+ attributes=[],
+ )
+
+ def test_evaluate_str_expression(self):
+ s = "hello world"
+ assert ExpressionEvaluator.evaluate(
+ "s", runtime_context={"s": s}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="s",
+ value_ref="s",
+ value_str="'hello world'",
+ value_type="str",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="__getitem__(0)",
+ value_ref="0",
+ value_str="'h'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(1)",
+ value_ref="1",
+ value_str="'e'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(2)",
+ value_ref="2",
+ value_str="'l'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(3)",
+ value_ref="3",
+ value_str="'l'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(4)",
+ value_ref="4",
+ value_str="'o'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(5)",
+ value_ref="5",
+ value_str="' '",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(6)",
+ value_ref="6",
+ value_str="'w'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(7)",
+ value_ref="7",
+ value_str="'o'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(8)",
+ value_ref="8",
+ value_str="'r'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(9)",
+ value_ref="9",
+ value_str="'l'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__(10)",
+ value_ref="10",
+ value_str="'d'",
+ value_type="str",
+ expandable=True,
+ ),
+ ],
+ )
+ assert ExpressionEvaluator.evaluate(
+ "s[4]", runtime_context={"s": s}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="s[4]",
+ value_ref="s[4]",
+ value_str="'o'",
+ value_type="str",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="__getitem__(0)",
+ value_ref="0",
+ value_str="'o'",
+ value_type="str",
+ expandable=True,
+ )
+ ],
+ )
+
+ assert ExpressionEvaluator.evaluate(
+ "s.__getitem__(2)", runtime_context={"s": s}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="s.__getitem__(2)",
+ value_ref="s.__getitem__(2)",
+ value_str="'l'",
+ value_type="str",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="__getitem__(0)",
+ value_ref="0",
+ value_str="'l'",
+ value_type="str",
+ expandable=True,
+ )
+ ],
+ )
+
+ def test_evaluate_object_expression(self):
+ class A:
+ def __init__(self):
+ self.i = 10
+ self.j = 1.1
+
+ a = A()
+
+ assert ExpressionEvaluator.evaluate(
+ "a", runtime_context={"a": a}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="a",
+ value_ref="a",
+ value_str=(
+ ".A"
+ f" object at {hex(id(a))}>"
+ ),
+ value_type="A",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="i",
+ value_ref="i",
+ value_str="10",
+ value_type="int",
+ expandable=False,
+ ),
+ TypedValue(
+ expression="j",
+ value_ref="j",
+ value_str="1.1",
+ value_type="float",
+ expandable=False,
+ ),
+ ],
+ )
+
+ def test_evaluate_container_expressions(self):
+ i = 10
+ f = 1.1
+
+ a_list = [i, f, (i, f)]
+ assert ExpressionEvaluator.evaluate(
+ "a_list", runtime_context={"a_list": a_list}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="a_list",
+ value_ref="a_list",
+ value_str="[10, 1.1, (10, 1.1)]",
+ value_type="list",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="__getitem__(0)",
+ value_ref="0",
+ value_str="10",
+ value_type="int",
+ expandable=False,
+ ),
+ TypedValue(
+ expression="__getitem__(1)",
+ value_ref="1",
+ value_str="1.1",
+ value_type="float",
+ expandable=False,
+ ),
+ TypedValue(
+ expression="__getitem__(2)",
+ value_ref="2",
+ value_str="(10, 1.1)",
+ value_type="tuple",
+ expandable=True,
+ ),
+ ],
+ )
+ t = (i, f, {i, f})
+ assert ExpressionEvaluator.evaluate(
+ "t", runtime_context={"t": t}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="t",
+ value_ref="t",
+ value_str="(10, 1.1, {1.1, 10})",
+ value_type="tuple",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="__getitem__(0)",
+ value_ref="0",
+ value_str="10",
+ value_type="int",
+ expandable=False,
+ ),
+ TypedValue(
+ expression="__getitem__(1)",
+ value_ref="1",
+ value_str="1.1",
+ value_type="float",
+ expandable=False,
+ ),
+ TypedValue(
+ expression="__getitem__(2)",
+ value_ref="2",
+ value_str="{1.1, 10}",
+ value_type="set",
+ expandable=True,
+ ),
+ ],
+ )
+ s = {i, f, (i, f)}
+ assert ExpressionEvaluator.evaluate(
+ "s", runtime_context={"s": s}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="s",
+ value_ref="s",
+ value_str="{1.1, 10, (10, 1.1)}",
+ value_type="set",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="__getitem__(0)",
+ value_ref="0",
+ value_str="1.1",
+ value_type="float",
+ expandable=False,
+ ),
+ TypedValue(
+ expression="__getitem__(1)",
+ value_ref="1",
+ value_str="10",
+ value_type="int",
+ expandable=False,
+ ),
+ TypedValue(
+ expression="__getitem__(2)",
+ value_ref="2",
+ value_str="(10, 1.1)",
+ value_type="tuple",
+ expandable=False,
+ ),
+ ],
+ )
+
+ d = {1: "a", "b": [{i, f}], (i,): f}
+ assert ExpressionEvaluator.evaluate(
+ "d", runtime_context={"d": d}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="d",
+ value_ref="d",
+ value_str="{1: 'a', 'b': [{1.1, 10}], (10,): 1.1}",
+ value_type="dict",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="__getitem__(1)",
+ value_ref="1",
+ value_str="'a'",
+ value_type="str",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__('b')",
+ value_ref="'b'",
+ value_str="[{1.1, 10}]",
+ value_type="list",
+ expandable=True,
+ ),
+ TypedValue(
+ expression="__getitem__((10,))",
+ value_ref="(10,)",
+ value_str="1.1",
+ value_type="float",
+ expandable=False,
+ ),
+ ],
+ )
+
+ g = (i for i in range(10))
+ assert ExpressionEvaluator.evaluate(
+ "g", runtime_context={"g": g}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="g",
+ value_ref="g",
+ value_str=(
+ ". at {hex(id(g))}>"
+ ),
+ value_type="generator",
+ expandable=True,
+ ),
+ attributes=[],
+ )
+
+ def gen():
+ for i in range(10):
+ yield i
+
+ g = gen()
+ next(g)
+ assert ExpressionEvaluator.evaluate(
+ "g", runtime_context={"g": g}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="g",
+ value_ref="g",
+ value_str=(
+ ".gen at {hex(id(g))}>"
+ ),
+ value_type="generator",
+ expandable=True,
+ ),
+ attributes=[
+ TypedValue(
+ expression="i",
+ value_ref="i",
+ value_str="0",
+ value_type="int",
+ expandable=False,
+ )
+ ],
+ )
+
+ it = iter([1, 2, 3])
+ assert ExpressionEvaluator.evaluate(
+ "it", runtime_context={"it": it}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="it",
+ value_ref="it",
+ value_str=f"",
+ value_type="list_iterator",
+ expandable=False,
+ ),
+ attributes=[],
+ )
+
+ it = iter([1, 2, 3])
+ next(it)
+ assert ExpressionEvaluator.evaluate(
+ "it", runtime_context={"it": it}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="it",
+ value_ref="it",
+ value_str=f"",
+ value_type="list_iterator",
+ expandable=False,
+ ),
+ attributes=[],
+ )
+
+ def test_evaluate_in_another_context(self):
+ i = 10
+ j = 20
+ assert ExpressionEvaluator.evaluate(
+ "j", runtime_context={"j": i, "i": j}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="j",
+ value_ref="j",
+ value_str="10",
+ value_type="int",
+ expandable=False,
+ ),
+ attributes=[],
+ )
+
+ assert ExpressionEvaluator.evaluate(
+ "i", runtime_context={"j": i, "i": j}
+ ) == EvaluatedValue(
+ value=TypedValue(
+ expression="i",
+ value_ref="i",
+ value_str="20",
+ value_type="int",
+ expandable=False,
+ ),
+ attributes=[],
+ )
diff --git a/amber/src/main/python/core/util/operator/__init__.py b/amber/src/main/python/core/util/operator/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/util/operator/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/util/proto/__init__.py b/amber/src/main/python/core/util/proto/__init__.py
new file mode 100644
index 00000000000..7ad3af5af4f
--- /dev/null
+++ b/amber/src/main/python/core/util/proto/__init__.py
@@ -0,0 +1,42 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import re
+from typing import T
+
+from betterproto import Message, which_one_of
+
+camel_case_pattern = re.compile(r"(? T:
+ _, value = which_one_of(base, ("sealed_" if sealed else "") + "value")
+ return value
+
+
+def set_one_of(base: T, value: Message) -> T:
+ name = value.__class__.__name__
+ name = name.strip("V2")
+ snake_case_name = re.sub(camel_case_pattern, "_", name).lower()
+ ret = base()
+ ret.__setattr__(snake_case_name, value)
+ return ret
+
+
+# implicitly used when being imported, this is to make betterproto
+# Messages hashable.
+Message.__hash__ = lambda x: hash(x.__repr__())
diff --git a/amber/src/main/python/core/util/protocol/base_protocols.py b/amber/src/main/python/core/util/protocol/base_protocols.py
new file mode 100644
index 00000000000..2cd42f3f37a
--- /dev/null
+++ b/amber/src/main/python/core/util/protocol/base_protocols.py
@@ -0,0 +1,59 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import abstractmethod
+from typing import TypeVar, Sized, Optional
+from typing_extensions import Protocol
+
+T = TypeVar("T")
+K = TypeVar("K")
+
+
+class Putable(Protocol):
+ @abstractmethod
+ def put(self, item: T) -> None:
+ pass
+
+
+class KeyedPutable(Protocol):
+ @abstractmethod
+ def put(self, key: K, item: T) -> None:
+ pass
+
+
+class Getable(Protocol):
+ @abstractmethod
+ def get(self) -> T:
+ pass
+
+
+class FlushedGetable(Protocol):
+ @abstractmethod
+ def get(self, flush: bool) -> T:
+ pass
+
+
+class EmtpyCheckable(Sized):
+ @abstractmethod
+ def is_empty(self) -> bool:
+ pass
+
+
+class KeyedEmtpyCheckable(Sized):
+ @abstractmethod
+ def is_empty(self, key: Optional[K] = None) -> bool:
+ pass
diff --git a/amber/src/main/python/core/util/runnable/runnable.py b/amber/src/main/python/core/util/runnable/runnable.py
new file mode 100644
index 00000000000..433c86fc8bb
--- /dev/null
+++ b/amber/src/main/python/core/util/runnable/runnable.py
@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import abstractmethod
+
+from typing_extensions import Protocol
+
+
+class Runnable(Protocol):
+ @abstractmethod
+ def run(self) -> None:
+ """run some logic"""
diff --git a/amber/src/main/python/core/util/stoppable/__init__.py b/amber/src/main/python/core/util/stoppable/__init__.py
new file mode 100644
index 00000000000..ad4690fe7d1
--- /dev/null
+++ b/amber/src/main/python/core/util/stoppable/__init__.py
@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .stoppable import Stoppable
+from .stoppable_queue_blocking_thread import StoppableQueueBlockingRunnable
+
+__all__ = ["Stoppable", "StoppableQueueBlockingRunnable"]
diff --git a/amber/src/main/python/core/util/stoppable/stoppable.py b/amber/src/main/python/core/util/stoppable/stoppable.py
new file mode 100644
index 00000000000..2c79e69d21d
--- /dev/null
+++ b/amber/src/main/python/core/util/stoppable/stoppable.py
@@ -0,0 +1,25 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import abstractmethod
+from typing_extensions import Protocol
+
+
+class Stoppable(Protocol):
+ @abstractmethod
+ def stop(self):
+ """stop self"""
diff --git a/amber/src/main/python/core/util/stoppable/stoppable_queue_blocking_thread.py b/amber/src/main/python/core/util/stoppable/stoppable_queue_blocking_thread.py
new file mode 100644
index 00000000000..d20073631b3
--- /dev/null
+++ b/amber/src/main/python/core/util/stoppable/stoppable_queue_blocking_thread.py
@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from loguru import logger
+from overrides import overrides
+
+from core.util.customized_queue.queue_base import IQueue, QueueControl, QueueElement
+from core.util.runnable.runnable import Runnable
+from core.util.stoppable.stoppable import Stoppable
+
+
+class StoppableQueueBlockingRunnable(Runnable, Stoppable):
+ """
+ An implementation of Stoppable, assuming the Runnable.run() would be blocked
+ by a blocking Queue.get(block=True, timeout=None).
+
+ For example:
+ ```
+ def run(self) -> None:
+ while True:
+ entry = queue.get() # here is a blocking Queue.get()
+ # do something with the entry
+ ```
+
+ According to https://docs.python.org/3/library/queue.html#queue.Queue.get, which
+ quoted as: "Prior to 3.0 on POSIX systems, and for all versions on Windows, if
+ block is true and timeout is None, this operation goes into an uninterruptible
+ wait on an underlying lock."
+
+ Currently, there is no other workaround for interrupting a waiting stoppable,
+ safely.
+
+ This implementation adds a special marker called
+ `StoppableQueueBlockingRunnable.RUNNABLE_STOP` into the queue, and when the
+ marker is consumed, it should break the Runnable.run().
+
+ """
+
+ RUNNABLE_STOP = QueueControl(msg="__RUNNABLE__STOP__MARKER__")
+
+ def __init__(self, name: str, queue: IQueue):
+ self._internal_queue = queue
+ self.name = name
+
+ @logger.catch(reraise=True)
+ @overrides
+ def run(self):
+ self.pre_start()
+ try:
+ while True:
+ self.receive(self.interruptible_get())
+ except StoppableQueueBlockingRunnable.InterruptRunnable:
+ # surpassed the expected interruption
+ logger.debug(f"{self.name}-interrupting")
+ finally:
+ self.post_stop()
+
+ @logger.catch(reraise=True)
+ def receive(self, next_entry: QueueElement):
+ pass
+
+ @logger.catch(reraise=True)
+ def pre_start(self) -> None:
+ pass
+
+ @logger.catch(reraise=True)
+ def post_stop(self) -> None:
+ pass
+
+ @logger.catch(reraise=True)
+ @overrides
+ def stop(self):
+ self._internal_queue.put(StoppableQueueBlockingRunnable.RUNNABLE_STOP)
+
+ def interruptible_get(self):
+ next_entry = self._internal_queue.get()
+ if next_entry == StoppableQueueBlockingRunnable.RUNNABLE_STOP:
+ raise StoppableQueueBlockingRunnable.InterruptRunnable
+ return next_entry
+
+ class InterruptRunnable(Exception):
+ """
+ Used to interrupt a runnable.
+ """
diff --git a/amber/src/main/python/core/util/thread/__init__.py b/amber/src/main/python/core/util/thread/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/core/util/thread/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/core/util/thread/atomic.py b/amber/src/main/python/core/util/thread/atomic.py
new file mode 100644
index 00000000000..a73619489fd
--- /dev/null
+++ b/amber/src/main/python/core/util/thread/atomic.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import threading
+
+
+class AtomicInteger:
+ def __init__(self, value=0):
+ self._value = int(value)
+ self._lock = threading.Lock()
+
+ def inc(self, d=1):
+ with self._lock:
+ self._value += int(d)
+ return self._value
+
+ def dec(self, d=1):
+ return self.inc(-d)
+
+ def get_and_inc(self, d=1):
+ with self._lock:
+ old_value = self._value
+ self._value += int(d)
+ return old_value
+
+ def get_and_dec(self, d=1):
+ return self.get_and_inc(-d)
+
+ @property
+ def value(self):
+ with self._lock:
+ return self._value
+
+ @value.setter
+ def value(self, v):
+ with self._lock:
+ self._value = int(v)
+ return self._value
+
+ def get_and_set(self, v):
+ with self._lock:
+ old_value = self.value
+ self._value = int(v)
+ return old_value
diff --git a/amber/src/main/python/core/util/virtual_identity/__init__.py b/amber/src/main/python/core/util/virtual_identity/__init__.py
new file mode 100644
index 00000000000..93a887d7c68
--- /dev/null
+++ b/amber/src/main/python/core/util/virtual_identity/__init__.py
@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import re
+from proto.org.apache.texera.amber.core import (
+ GlobalPortIdentity,
+ PhysicalOpIdentity,
+ OperatorIdentity,
+ PortIdentity,
+ ActorVirtualIdentity,
+)
+
+worker_name_pattern = re.compile(r"Worker:WF\d+-.+-(\w+)-(\d+)")
+
+MATERIALIZATION_READER_ACTOR_PREFIX = "MATERIALIZATION_READER_"
+
+
+def get_worker_index(worker_id: str) -> int:
+ match = worker_name_pattern.match(worker_id)
+ if match:
+ return int(match.group(2))
+ raise ValueError("Invalid worker ID format")
+
+
+def serialize_global_port_identity(obj: GlobalPortIdentity) -> str:
+ """
+ Serialize GlobalPortIdentity into a custom human-readable string.
+ Expected format:
+ ``(logicalOpId=,layerName=,
+ portId=,isInternal=,isInput= )``
+ """
+ logical_op_id = obj.op_id.logical_op_id.id
+ layer_name = obj.op_id.layer_name
+ port_id = obj.port_id.id
+ is_internal = obj.port_id.internal
+ is_input_port = obj.input
+ return (
+ f"(logicalOpId={logical_op_id},layerName={layer_name},portId={port_id},"
+ f"isInternal={str(is_internal).lower()},isInput={str(is_input_port).lower()})"
+ )
+
+
+def deserialize_global_port_identity(encoded_str: str) -> GlobalPortIdentity:
+ """
+ Deserialize a custom string from the format
+ ``(logicalOpId=,layerName=,
+ portId=,isInternal=,isInput= )``
+ back into a GlobalPortIdentity object.
+ """
+ pattern = (
+ r"\(logicalOpId=([^,]+),layerName=([^,]+),"
+ r"portId=([^,]+),isInternal=([^,]+),isInput=([^)]+)\)"
+ )
+ match = re.fullmatch(pattern, encoded_str)
+ if not match:
+ raise ValueError(f"Invalid GlobalPortIdentity format: {encoded_str}")
+ logical_op_id, layer_name, port_id_str, is_internal_str, is_input_str = (
+ match.groups()
+ )
+ port_id = int(port_id_str)
+ is_internal = is_internal_str.lower() == "true"
+ is_input_port = is_input_str.lower() == "true"
+ op_id = PhysicalOpIdentity(
+ logical_op_id=OperatorIdentity(id=logical_op_id), layer_name=layer_name
+ )
+ port = PortIdentity(id=port_id, internal=is_internal)
+ return GlobalPortIdentity(op_id=op_id, port_id=port, input=is_input_port)
+
+
+def get_from_actor_id_for_input_port_storage(
+ storage_uri_str: str, to_worker_actor_id: ActorVirtualIdentity
+) -> ActorVirtualIdentity:
+ """
+ Constructs an ActorVirtualIdentity for input port storage.
+
+ Args:
+ storage_uri_str (str): The string representation of the storage URI.
+
+ Returns:
+ ActorVirtualIdentity: A new virtual identity created by
+ prefixing the storage URI.
+ """
+ return ActorVirtualIdentity(
+ MATERIALIZATION_READER_ACTOR_PREFIX + storage_uri_str + to_worker_actor_id.name
+ )
diff --git a/core/sandbox/src/main/java/edu/uci/ics/texera/sandbox/AlchemyAPIexample/thirdparty/__init__.py b/amber/src/main/python/proto/__init__.py
similarity index 100%
rename from core/sandbox/src/main/java/edu/uci/ics/texera/sandbox/AlchemyAPIexample/thirdparty/__init__.py
rename to amber/src/main/python/proto/__init__.py
diff --git a/core/perftest/index/standard/promed/write.lock b/amber/src/main/python/proto/org/__init__.py
similarity index 100%
rename from core/perftest/index/standard/promed/write.lock
rename to amber/src/main/python/proto/org/__init__.py
diff --git a/core/sandbox/src/main/java/edu/uci/ics/texera/sandbox/AlchemyAPIexample/api_key.txt b/amber/src/main/python/proto/org/apache/__init__.py
similarity index 100%
rename from core/sandbox/src/main/java/edu/uci/ics/texera/sandbox/AlchemyAPIexample/api_key.txt
rename to amber/src/main/python/proto/org/apache/__init__.py
diff --git a/amber/src/main/python/proto/org/apache/texera/__init__.py b/amber/src/main/python/proto/org/apache/texera/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/amber/src/main/python/proto/org/apache/texera/amber/__init__.py b/amber/src/main/python/proto/org/apache/texera/amber/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/amber/src/main/python/proto/org/apache/texera/amber/core/__init__.py b/amber/src/main/python/proto/org/apache/texera/amber/core/__init__.py
new file mode 100644
index 00000000000..2d21638c263
--- /dev/null
+++ b/amber/src/main/python/proto/org/apache/texera/amber/core/__init__.py
@@ -0,0 +1,146 @@
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# sources: org/apache/texera/amber/core/executor.proto, org/apache/texera/amber/core/virtualidentity.proto, org/apache/texera/amber/core/workflow.proto, org/apache/texera/amber/core/workflowruntimestate.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import (
+ List,
+)
+
+import betterproto
+
+
+class OutputPortOutputMode(betterproto.Enum):
+ SET_SNAPSHOT = 0
+ """outputs complete result set snapshot for each update"""
+
+ SET_DELTA = 1
+ """outputs incremental result set delta for each update"""
+
+ SINGLE_SNAPSHOT = 2
+ """
+ outputs a single snapshot for the entire execution,
+ used explicitly to support visualization operators that may exceed the memory limit
+ TODO: remove this mode after we have a better solution for output size limit
+ """
+
+
+class FatalErrorType(betterproto.Enum):
+ COMPILATION_ERROR = 0
+ EXECUTION_FAILURE = 1
+
+
+@dataclass(eq=False, repr=False)
+class WorkflowIdentity(betterproto.Message):
+ id: int = betterproto.int64_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionIdentity(betterproto.Message):
+ id: int = betterproto.int64_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class ActorVirtualIdentity(betterproto.Message):
+ name: str = betterproto.string_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class ChannelIdentity(betterproto.Message):
+ from_worker_id: "ActorVirtualIdentity" = betterproto.message_field(1)
+ to_worker_id: "ActorVirtualIdentity" = betterproto.message_field(2)
+ is_control: bool = betterproto.bool_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class OperatorIdentity(betterproto.Message):
+ id: str = betterproto.string_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class PhysicalOpIdentity(betterproto.Message):
+ logical_op_id: "OperatorIdentity" = betterproto.message_field(1)
+ layer_name: str = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class EmbeddedControlMessageIdentity(betterproto.Message):
+ id: str = betterproto.string_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class PortIdentity(betterproto.Message):
+ id: int = betterproto.int32_field(1)
+ internal: bool = betterproto.bool_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class GlobalPortIdentity(betterproto.Message):
+ op_id: "PhysicalOpIdentity" = betterproto.message_field(1)
+ port_id: "PortIdentity" = betterproto.message_field(2)
+ input: bool = betterproto.bool_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class InputPort(betterproto.Message):
+ id: "PortIdentity" = betterproto.message_field(1)
+ display_name: str = betterproto.string_field(2)
+ disallow_multi_links: bool = betterproto.bool_field(3)
+ dependencies: List["PortIdentity"] = betterproto.message_field(4)
+
+
+@dataclass(eq=False, repr=False)
+class OutputPort(betterproto.Message):
+ id: "PortIdentity" = betterproto.message_field(1)
+ display_name: str = betterproto.string_field(2)
+ blocking: bool = betterproto.bool_field(3)
+ mode: "OutputPortOutputMode" = betterproto.enum_field(4)
+
+
+@dataclass(eq=False, repr=False)
+class PhysicalLink(betterproto.Message):
+ from_op_id: "PhysicalOpIdentity" = betterproto.message_field(1)
+ from_port_id: "PortIdentity" = betterproto.message_field(2)
+ to_op_id: "PhysicalOpIdentity" = betterproto.message_field(3)
+ to_port_id: "PortIdentity" = betterproto.message_field(4)
+
+
+@dataclass(eq=False, repr=False)
+class OpExecWithCode(betterproto.Message):
+ code: str = betterproto.string_field(1)
+ language: str = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class OpExecWithClassName(betterproto.Message):
+ class_name: str = betterproto.string_field(1)
+ desc_string: str = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class OpExecSource(betterproto.Message):
+ storage_key: str = betterproto.string_field(1)
+ workflow_identity: "WorkflowIdentity" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class OpExecInitInfo(betterproto.Message):
+ op_exec_with_class_name: "OpExecWithClassName" = betterproto.message_field(
+ 1, group="sealed_value"
+ )
+ op_exec_with_code: "OpExecWithCode" = betterproto.message_field(
+ 2, group="sealed_value"
+ )
+ op_exec_source: "OpExecSource" = betterproto.message_field(3, group="sealed_value")
+
+
+@dataclass(eq=False, repr=False)
+class WorkflowFatalError(betterproto.Message):
+ type: "FatalErrorType" = betterproto.enum_field(1)
+ timestamp: datetime = betterproto.message_field(2)
+ message: str = betterproto.string_field(3)
+ details: str = betterproto.string_field(4)
+ operator_id: str = betterproto.string_field(5)
+ worker_id: str = betterproto.string_field(6)
diff --git a/amber/src/main/python/proto/org/apache/texera/amber/engine/__init__.py b/amber/src/main/python/proto/org/apache/texera/amber/engine/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/__init__.py b/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/rpc/__init__.py b/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/rpc/__init__.py
new file mode 100644
index 00000000000..77d51933af6
--- /dev/null
+++ b/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/rpc/__init__.py
@@ -0,0 +1,2161 @@
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# sources: org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto, org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto, org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto, org/apache/texera/amber/engine/architecture/rpc/testerservice.proto, org/apache/texera/amber/engine/architecture/rpc/workerservice.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import (
+ TYPE_CHECKING,
+ Dict,
+ List,
+ Optional,
+)
+
+import betterproto
+import grpclib
+from betterproto.grpc.grpclib_server import ServiceBase
+
+from .... import core as ___core__
+from .. import (
+ sendsemantics as _sendsemantics__,
+ worker as _worker__,
+)
+
+
+if TYPE_CHECKING:
+ import grpclib.server
+ from betterproto.grpc.grpclib_client import MetadataLike
+ from grpclib.metadata import Deadline
+
+
+class EmbeddedControlMessageType(betterproto.Enum):
+ ALL_ALIGNMENT = 0
+ NO_ALIGNMENT = 1
+ PORT_ALIGNMENT = 2
+
+
+class ConsoleMessageType(betterproto.Enum):
+ PRINT = 0
+ ERROR = 1
+ COMMAND = 2
+ DEBUGGER = 3
+
+
+class StatisticsUpdateTarget(betterproto.Enum):
+ BOTH_UI_AND_PERSISTENCE = 0
+ UI_ONLY = 1
+ PERSISTENCE_ONLY = 2
+
+
+class ErrorLanguage(betterproto.Enum):
+ PYTHON = 0
+ SCALA = 1
+
+
+class WorkflowAggregatedState(betterproto.Enum):
+ UNINITIALIZED = 0
+ READY = 1
+ RUNNING = 2
+ PAUSING = 3
+ PAUSED = 4
+ RESUMING = 5
+ COMPLETED = 6
+ FAILED = 7
+ UNKNOWN = 8
+ KILLED = 9
+ TERMINATED = 10
+
+
+@dataclass(eq=False, repr=False)
+class ControlRequest(betterproto.Message):
+ propagate_embedded_control_message_request: (
+ "PropagateEmbeddedControlMessageRequest"
+ ) = betterproto.message_field(1, group="sealed_value")
+ """request for controller"""
+
+ take_global_checkpoint_request: "TakeGlobalCheckpointRequest" = (
+ betterproto.message_field(2, group="sealed_value")
+ )
+ debug_command_request: "DebugCommandRequest" = betterproto.message_field(
+ 3, group="sealed_value"
+ )
+ evaluate_python_expression_request: "EvaluatePythonExpressionRequest" = (
+ betterproto.message_field(4, group="sealed_value")
+ )
+ retry_workflow_request: "RetryWorkflowRequest" = betterproto.message_field(
+ 5, group="sealed_value"
+ )
+ console_message_triggered_request: "ConsoleMessageTriggeredRequest" = (
+ betterproto.message_field(6, group="sealed_value")
+ )
+ port_completed_request: "PortCompletedRequest" = betterproto.message_field(
+ 7, group="sealed_value"
+ )
+ worker_state_updated_request: "WorkerStateUpdatedRequest" = (
+ betterproto.message_field(8, group="sealed_value")
+ )
+ link_workers_request: "LinkWorkersRequest" = betterproto.message_field(
+ 9, group="sealed_value"
+ )
+ workflow_reconfigure_request: "WorkflowReconfigureRequest" = (
+ betterproto.message_field(10, group="sealed_value")
+ )
+ add_input_channel_request: "AddInputChannelRequest" = betterproto.message_field(
+ 50, group="sealed_value"
+ )
+ """request for worker"""
+
+ add_partitioning_request: "AddPartitioningRequest" = betterproto.message_field(
+ 51, group="sealed_value"
+ )
+ assign_port_request: "AssignPortRequest" = betterproto.message_field(
+ 52, group="sealed_value"
+ )
+ finalize_checkpoint_request: "FinalizeCheckpointRequest" = (
+ betterproto.message_field(53, group="sealed_value")
+ )
+ initialize_executor_request: "InitializeExecutorRequest" = (
+ betterproto.message_field(54, group="sealed_value")
+ )
+ update_executor_request: "UpdateExecutorRequest" = betterproto.message_field(
+ 55, group="sealed_value"
+ )
+ empty_request: "EmptyRequest" = betterproto.message_field(56, group="sealed_value")
+ prepare_checkpoint_request: "PrepareCheckpointRequest" = betterproto.message_field(
+ 57, group="sealed_value"
+ )
+ query_statistics_request: "QueryStatisticsRequest" = betterproto.message_field(
+ 58, group="sealed_value"
+ )
+ ping: "Ping" = betterproto.message_field(100, group="sealed_value")
+ """request for testing"""
+
+ pong: "Pong" = betterproto.message_field(101, group="sealed_value")
+ nested: "Nested" = betterproto.message_field(102, group="sealed_value")
+ pass_: "Pass" = betterproto.message_field(103, group="sealed_value")
+ error_command: "ErrorCommand" = betterproto.message_field(104, group="sealed_value")
+ recursion: "Recursion" = betterproto.message_field(105, group="sealed_value")
+ collect: "Collect" = betterproto.message_field(106, group="sealed_value")
+ generate_number: "GenerateNumber" = betterproto.message_field(
+ 107, group="sealed_value"
+ )
+ multi_call: "MultiCall" = betterproto.message_field(108, group="sealed_value")
+ chain: "Chain" = betterproto.message_field(109, group="sealed_value")
+
+
+@dataclass(eq=False, repr=False)
+class EmptyRequest(betterproto.Message):
+ pass
+
+
+@dataclass(eq=False, repr=False)
+class AsyncRpcContext(betterproto.Message):
+ sender: "___core__.ActorVirtualIdentity" = betterproto.message_field(1)
+ receiver: "___core__.ActorVirtualIdentity" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class ControlInvocation(betterproto.Message):
+ method_name: str = betterproto.string_field(1)
+ command: "ControlRequest" = betterproto.message_field(2)
+ context: "AsyncRpcContext" = betterproto.message_field(3)
+ command_id: int = betterproto.int64_field(4)
+
+
+@dataclass(eq=False, repr=False)
+class EmbeddedControlMessage(betterproto.Message):
+ id: "___core__.EmbeddedControlMessageIdentity" = betterproto.message_field(1)
+ ecm_type: "EmbeddedControlMessageType" = betterproto.enum_field(2)
+ scope: List["___core__.ChannelIdentity"] = betterproto.message_field(3)
+ command_mapping: Dict[str, "ControlInvocation"] = betterproto.map_field(
+ 4, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+
+
+@dataclass(eq=False, repr=False)
+class PropagateEmbeddedControlMessageRequest(betterproto.Message):
+ source_op_to_start_prop: List["___core__.PhysicalOpIdentity"] = (
+ betterproto.message_field(1)
+ )
+ id: "___core__.EmbeddedControlMessageIdentity" = betterproto.message_field(2)
+ ecm_type: "EmbeddedControlMessageType" = betterproto.enum_field(3)
+ scope: List["___core__.PhysicalOpIdentity"] = betterproto.message_field(4)
+ target_ops: List["___core__.PhysicalOpIdentity"] = betterproto.message_field(5)
+ command: "ControlRequest" = betterproto.message_field(6)
+ method_name: str = betterproto.string_field(7)
+
+
+@dataclass(eq=False, repr=False)
+class TakeGlobalCheckpointRequest(betterproto.Message):
+ estimation_only: bool = betterproto.bool_field(1)
+ checkpoint_id: "___core__.EmbeddedControlMessageIdentity" = (
+ betterproto.message_field(2)
+ )
+ destination: str = betterproto.string_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class WorkflowReconfigureRequest(betterproto.Message):
+ reconfiguration: List["UpdateExecutorRequest"] = betterproto.message_field(1)
+ reconfiguration_id: str = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class DebugCommandRequest(betterproto.Message):
+ worker_id: str = betterproto.string_field(1)
+ cmd: str = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class EvaluatePythonExpressionRequest(betterproto.Message):
+ expression: str = betterproto.string_field(1)
+ operator_id: str = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class RetryWorkflowRequest(betterproto.Message):
+ workers: List["___core__.ActorVirtualIdentity"] = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class ConsoleMessage(betterproto.Message):
+ worker_id: str = betterproto.string_field(1)
+ timestamp: datetime = betterproto.message_field(2)
+ msg_type: "ConsoleMessageType" = betterproto.enum_field(3)
+ source: str = betterproto.string_field(4)
+ title: str = betterproto.string_field(5)
+ message: str = betterproto.string_field(6)
+
+
+@dataclass(eq=False, repr=False)
+class ConsoleMessageTriggeredRequest(betterproto.Message):
+ console_message: "ConsoleMessage" = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class PortCompletedRequest(betterproto.Message):
+ port_id: "___core__.PortIdentity" = betterproto.message_field(1)
+ input: bool = betterproto.bool_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class WorkerStateUpdatedRequest(betterproto.Message):
+ state: "_worker__.WorkerState" = betterproto.enum_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class LinkWorkersRequest(betterproto.Message):
+ link: "___core__.PhysicalLink" = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class Ping(betterproto.Message):
+ """Ping message"""
+
+ i: int = betterproto.int32_field(1)
+ end: int = betterproto.int32_field(2)
+ to: "___core__.ActorVirtualIdentity" = betterproto.message_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class Pong(betterproto.Message):
+ """Pong message"""
+
+ i: int = betterproto.int32_field(1)
+ end: int = betterproto.int32_field(2)
+ to: "___core__.ActorVirtualIdentity" = betterproto.message_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class Pass(betterproto.Message):
+ """Pass message"""
+
+ value: str = betterproto.string_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class Nested(betterproto.Message):
+ """Nested message"""
+
+ k: int = betterproto.int32_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class MultiCall(betterproto.Message):
+ """MultiCall message"""
+
+ seq: List["___core__.ActorVirtualIdentity"] = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class ErrorCommand(betterproto.Message):
+ """ErrorCommand message"""
+
+ pass
+
+
+@dataclass(eq=False, repr=False)
+class Collect(betterproto.Message):
+ """Collect message"""
+
+ workers: List["___core__.ActorVirtualIdentity"] = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class GenerateNumber(betterproto.Message):
+ """GenerateNumber message"""
+
+ pass
+
+
+@dataclass(eq=False, repr=False)
+class Chain(betterproto.Message):
+ """Chain message"""
+
+ nexts: List["___core__.ActorVirtualIdentity"] = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class Recursion(betterproto.Message):
+ """Recursion message"""
+
+ i: int = betterproto.int32_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class AddInputChannelRequest(betterproto.Message):
+ """Messages for the commands"""
+
+ channel_id: "___core__.ChannelIdentity" = betterproto.message_field(1)
+ port_id: "___core__.PortIdentity" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class AddPartitioningRequest(betterproto.Message):
+ tag: "___core__.PhysicalLink" = betterproto.message_field(1)
+ partitioning: "_sendsemantics__.Partitioning" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class AssignPortRequest(betterproto.Message):
+ port_id: "___core__.PortIdentity" = betterproto.message_field(1)
+ input: bool = betterproto.bool_field(2)
+ schema: Dict[str, str] = betterproto.map_field(
+ 3, betterproto.TYPE_STRING, betterproto.TYPE_STRING
+ )
+ storage_uris: List[str] = betterproto.string_field(4)
+ partitionings: List["_sendsemantics__.Partitioning"] = betterproto.message_field(5)
+
+
+@dataclass(eq=False, repr=False)
+class FinalizeCheckpointRequest(betterproto.Message):
+ checkpoint_id: "___core__.EmbeddedControlMessageIdentity" = (
+ betterproto.message_field(1)
+ )
+ write_to: str = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class InitializeExecutorRequest(betterproto.Message):
+ total_worker_count: int = betterproto.int32_field(1)
+ op_exec_init_info: "___core__.OpExecInitInfo" = betterproto.message_field(2)
+ is_source: bool = betterproto.bool_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class UpdateExecutorRequest(betterproto.Message):
+ target_op_id: "___core__.PhysicalOpIdentity" = betterproto.message_field(1)
+ new_exec_init_info: "___core__.OpExecInitInfo" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class PrepareCheckpointRequest(betterproto.Message):
+ checkpoint_id: "___core__.EmbeddedControlMessageIdentity" = (
+ betterproto.message_field(1)
+ )
+ estimation_only: bool = betterproto.bool_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class QueryStatisticsRequest(betterproto.Message):
+ filter_by_workers: List["___core__.ActorVirtualIdentity"] = (
+ betterproto.message_field(1)
+ )
+ update_target: "StatisticsUpdateTarget" = betterproto.enum_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class ControlReturn(betterproto.Message):
+ """The generic return message"""
+
+ retrieve_workflow_state_response: "RetrieveWorkflowStateResponse" = (
+ betterproto.message_field(1, group="sealed_value")
+ )
+ """controller responses"""
+
+ propagate_embedded_control_message_response: (
+ "PropagateEmbeddedControlMessageResponse"
+ ) = betterproto.message_field(2, group="sealed_value")
+ take_global_checkpoint_response: "TakeGlobalCheckpointResponse" = (
+ betterproto.message_field(3, group="sealed_value")
+ )
+ evaluate_python_expression_response: "EvaluatePythonExpressionResponse" = (
+ betterproto.message_field(4, group="sealed_value")
+ )
+ start_workflow_response: "StartWorkflowResponse" = betterproto.message_field(
+ 5, group="sealed_value"
+ )
+ worker_state_response: "WorkerStateResponse" = betterproto.message_field(
+ 50, group="sealed_value"
+ )
+ """worker responses"""
+
+ worker_metrics_response: "WorkerMetricsResponse" = betterproto.message_field(
+ 51, group="sealed_value"
+ )
+ finalize_checkpoint_response: "FinalizeCheckpointResponse" = (
+ betterproto.message_field(52, group="sealed_value")
+ )
+ control_error: "ControlError" = betterproto.message_field(101, group="sealed_value")
+ """common responses"""
+
+ empty_return: "EmptyReturn" = betterproto.message_field(102, group="sealed_value")
+ string_response: "StringResponse" = betterproto.message_field(
+ 103, group="sealed_value"
+ )
+ int_response: "IntResponse" = betterproto.message_field(104, group="sealed_value")
+
+
+@dataclass(eq=False, repr=False)
+class EmptyReturn(betterproto.Message):
+ pass
+
+
+@dataclass(eq=False, repr=False)
+class ControlError(betterproto.Message):
+ error_message: str = betterproto.string_field(1)
+ error_details: str = betterproto.string_field(2)
+ stack_trace: str = betterproto.string_field(3)
+ language: "ErrorLanguage" = betterproto.enum_field(4)
+
+
+@dataclass(eq=False, repr=False)
+class ReturnInvocation(betterproto.Message):
+ command_id: int = betterproto.int64_field(1)
+ return_value: "ControlReturn" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class StringResponse(betterproto.Message):
+ value: str = betterproto.string_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class IntResponse(betterproto.Message):
+ value: int = betterproto.int32_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class RetrieveWorkflowStateResponse(betterproto.Message):
+ state: Dict[str, str] = betterproto.map_field(
+ 1, betterproto.TYPE_STRING, betterproto.TYPE_STRING
+ )
+
+
+@dataclass(eq=False, repr=False)
+class FinalizeCheckpointResponse(betterproto.Message):
+ size: int = betterproto.int64_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class PropagateEmbeddedControlMessageResponse(betterproto.Message):
+ returns: Dict[str, "ControlReturn"] = betterproto.map_field(
+ 1, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+
+
+@dataclass(eq=False, repr=False)
+class TakeGlobalCheckpointResponse(betterproto.Message):
+ total_size: int = betterproto.int64_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class TypedValue(betterproto.Message):
+ expression: str = betterproto.string_field(1)
+ value_ref: str = betterproto.string_field(2)
+ value_str: str = betterproto.string_field(3)
+ value_type: str = betterproto.string_field(4)
+ expandable: bool = betterproto.bool_field(5)
+
+
+@dataclass(eq=False, repr=False)
+class EvaluatedValue(betterproto.Message):
+ value: "TypedValue" = betterproto.message_field(1)
+ attributes: List["TypedValue"] = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class EvaluatePythonExpressionResponse(betterproto.Message):
+ values: List["EvaluatedValue"] = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class StartWorkflowResponse(betterproto.Message):
+ workflow_state: "WorkflowAggregatedState" = betterproto.enum_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class WorkerStateResponse(betterproto.Message):
+ state: "_worker__.WorkerState" = betterproto.enum_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class WorkerMetricsResponse(betterproto.Message):
+ metrics: "_worker__.WorkerMetrics" = betterproto.message_field(1)
+
+
+class RpcTesterStub(betterproto.ServiceStub):
+ async def send_ping(
+ self,
+ ping: "Ping",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "IntResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendPing",
+ ping,
+ IntResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_pong(
+ self,
+ pong: "Pong",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "IntResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendPong",
+ pong,
+ IntResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_nested(
+ self,
+ nested: "Nested",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "StringResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendNested",
+ nested,
+ StringResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_pass(
+ self,
+ pass_: "Pass",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "StringResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendPass",
+ pass_,
+ StringResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_error_command(
+ self,
+ error_command: "ErrorCommand",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "StringResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendErrorCommand",
+ error_command,
+ StringResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_recursion(
+ self,
+ recursion: "Recursion",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "StringResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendRecursion",
+ recursion,
+ StringResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_collect(
+ self,
+ collect: "Collect",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "StringResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendCollect",
+ collect,
+ StringResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_generate_number(
+ self,
+ generate_number: "GenerateNumber",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "IntResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendGenerateNumber",
+ generate_number,
+ IntResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_multi_call(
+ self,
+ multi_call: "MultiCall",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "StringResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendMultiCall",
+ multi_call,
+ StringResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def send_chain(
+ self,
+ chain: "Chain",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "StringResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendChain",
+ chain,
+ StringResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+
+class WorkerServiceStub(betterproto.ServiceStub):
+ async def add_input_channel(
+ self,
+ add_input_channel_request: "AddInputChannelRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/AddInputChannel",
+ add_input_channel_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def add_partitioning(
+ self,
+ add_partitioning_request: "AddPartitioningRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/AddPartitioning",
+ add_partitioning_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def assign_port(
+ self,
+ assign_port_request: "AssignPortRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/AssignPort",
+ assign_port_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def finalize_checkpoint(
+ self,
+ finalize_checkpoint_request: "FinalizeCheckpointRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "FinalizeCheckpointResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/FinalizeCheckpoint",
+ finalize_checkpoint_request,
+ FinalizeCheckpointResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def flush_network_buffer(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/FlushNetworkBuffer",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def initialize_executor(
+ self,
+ initialize_executor_request: "InitializeExecutorRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/InitializeExecutor",
+ initialize_executor_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def open_executor(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/OpenExecutor",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def pause_worker(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "WorkerStateResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/PauseWorker",
+ empty_request,
+ WorkerStateResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def prepare_checkpoint(
+ self,
+ prepare_checkpoint_request: "PrepareCheckpointRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/PrepareCheckpoint",
+ prepare_checkpoint_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def query_statistics(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "WorkerMetricsResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/QueryStatistics",
+ empty_request,
+ WorkerMetricsResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def resume_worker(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "WorkerStateResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/ResumeWorker",
+ empty_request,
+ WorkerStateResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def retrieve_state(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/RetrieveState",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def retry_current_tuple(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/RetryCurrentTuple",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def start_worker(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "WorkerStateResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/StartWorker",
+ empty_request,
+ WorkerStateResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def end_worker(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/EndWorker",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def start_channel(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/StartChannel",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def end_channel(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/EndChannel",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def debug_command(
+ self,
+ debug_command_request: "DebugCommandRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/DebugCommand",
+ debug_command_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def evaluate_python_expression(
+ self,
+ evaluate_python_expression_request: "EvaluatePythonExpressionRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EvaluatedValue":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/EvaluatePythonExpression",
+ evaluate_python_expression_request,
+ EvaluatedValue,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def no_operation(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/NoOperation",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def update_executor(
+ self,
+ update_executor_request: "UpdateExecutorRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/UpdateExecutor",
+ update_executor_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+
+class ControllerServiceStub(betterproto.ServiceStub):
+ async def retrieve_workflow_state(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "RetrieveWorkflowStateResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/RetrieveWorkflowState",
+ empty_request,
+ RetrieveWorkflowStateResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def propagate_embedded_control_message(
+ self,
+ propagate_embedded_control_message_request: "PropagateEmbeddedControlMessageRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "PropagateEmbeddedControlMessageResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/PropagateEmbeddedControlMessage",
+ propagate_embedded_control_message_request,
+ PropagateEmbeddedControlMessageResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def take_global_checkpoint(
+ self,
+ take_global_checkpoint_request: "TakeGlobalCheckpointRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "TakeGlobalCheckpointResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/TakeGlobalCheckpoint",
+ take_global_checkpoint_request,
+ TakeGlobalCheckpointResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def debug_command(
+ self,
+ debug_command_request: "DebugCommandRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/DebugCommand",
+ debug_command_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def evaluate_python_expression(
+ self,
+ evaluate_python_expression_request: "EvaluatePythonExpressionRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EvaluatePythonExpressionResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/EvaluatePythonExpression",
+ evaluate_python_expression_request,
+ EvaluatePythonExpressionResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def console_message_triggered(
+ self,
+ console_message_triggered_request: "ConsoleMessageTriggeredRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/ConsoleMessageTriggered",
+ console_message_triggered_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def port_completed(
+ self,
+ port_completed_request: "PortCompletedRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/PortCompleted",
+ port_completed_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def start_workflow(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "StartWorkflowResponse":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/StartWorkflow",
+ empty_request,
+ StartWorkflowResponse,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def resume_workflow(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/ResumeWorkflow",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def pause_workflow(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/PauseWorkflow",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def worker_state_updated(
+ self,
+ worker_state_updated_request: "WorkerStateUpdatedRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/WorkerStateUpdated",
+ worker_state_updated_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def worker_execution_completed(
+ self,
+ empty_request: "EmptyRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/WorkerExecutionCompleted",
+ empty_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def link_workers(
+ self,
+ link_workers_request: "LinkWorkersRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/LinkWorkers",
+ link_workers_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def controller_initiate_query_statistics(
+ self,
+ query_statistics_request: "QueryStatisticsRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/ControllerInitiateQueryStatistics",
+ query_statistics_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def retry_workflow(
+ self,
+ retry_workflow_request: "RetryWorkflowRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/RetryWorkflow",
+ retry_workflow_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+ async def reconfigure_workflow(
+ self,
+ workflow_reconfigure_request: "WorkflowReconfigureRequest",
+ *,
+ timeout: Optional[float] = None,
+ deadline: Optional["Deadline"] = None,
+ metadata: Optional["MetadataLike"] = None
+ ) -> "EmptyReturn":
+ return await self._unary_unary(
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/ReconfigureWorkflow",
+ workflow_reconfigure_request,
+ EmptyReturn,
+ timeout=timeout,
+ deadline=deadline,
+ metadata=metadata,
+ )
+
+
+class RpcTesterBase(ServiceBase):
+
+ async def send_ping(self, ping: "Ping") -> "IntResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_pong(self, pong: "Pong") -> "IntResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_nested(self, nested: "Nested") -> "StringResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_pass(self, pass_: "Pass") -> "StringResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_error_command(
+ self, error_command: "ErrorCommand"
+ ) -> "StringResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_recursion(self, recursion: "Recursion") -> "StringResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_collect(self, collect: "Collect") -> "StringResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_generate_number(
+ self, generate_number: "GenerateNumber"
+ ) -> "IntResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_multi_call(self, multi_call: "MultiCall") -> "StringResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def send_chain(self, chain: "Chain") -> "StringResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def __rpc_send_ping(
+ self, stream: "grpclib.server.Stream[Ping, IntResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_ping(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_pong(
+ self, stream: "grpclib.server.Stream[Pong, IntResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_pong(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_nested(
+ self, stream: "grpclib.server.Stream[Nested, StringResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_nested(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_pass(
+ self, stream: "grpclib.server.Stream[Pass, StringResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_pass(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_error_command(
+ self, stream: "grpclib.server.Stream[ErrorCommand, StringResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_error_command(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_recursion(
+ self, stream: "grpclib.server.Stream[Recursion, StringResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_recursion(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_collect(
+ self, stream: "grpclib.server.Stream[Collect, StringResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_collect(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_generate_number(
+ self, stream: "grpclib.server.Stream[GenerateNumber, IntResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_generate_number(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_multi_call(
+ self, stream: "grpclib.server.Stream[MultiCall, StringResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_multi_call(request)
+ await stream.send_message(response)
+
+ async def __rpc_send_chain(
+ self, stream: "grpclib.server.Stream[Chain, StringResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.send_chain(request)
+ await stream.send_message(response)
+
+ def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
+ return {
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendPing": grpclib.const.Handler(
+ self.__rpc_send_ping,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ Ping,
+ IntResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendPong": grpclib.const.Handler(
+ self.__rpc_send_pong,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ Pong,
+ IntResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendNested": grpclib.const.Handler(
+ self.__rpc_send_nested,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ Nested,
+ StringResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendPass": grpclib.const.Handler(
+ self.__rpc_send_pass,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ Pass,
+ StringResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendErrorCommand": grpclib.const.Handler(
+ self.__rpc_send_error_command,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ ErrorCommand,
+ StringResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendRecursion": grpclib.const.Handler(
+ self.__rpc_send_recursion,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ Recursion,
+ StringResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendCollect": grpclib.const.Handler(
+ self.__rpc_send_collect,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ Collect,
+ StringResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendGenerateNumber": grpclib.const.Handler(
+ self.__rpc_send_generate_number,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ GenerateNumber,
+ IntResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendMultiCall": grpclib.const.Handler(
+ self.__rpc_send_multi_call,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ MultiCall,
+ StringResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.RPCTester/SendChain": grpclib.const.Handler(
+ self.__rpc_send_chain,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ Chain,
+ StringResponse,
+ ),
+ }
+
+
+class WorkerServiceBase(ServiceBase):
+
+ async def add_input_channel(
+ self, add_input_channel_request: "AddInputChannelRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def add_partitioning(
+ self, add_partitioning_request: "AddPartitioningRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def assign_port(
+ self, assign_port_request: "AssignPortRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def finalize_checkpoint(
+ self, finalize_checkpoint_request: "FinalizeCheckpointRequest"
+ ) -> "FinalizeCheckpointResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def flush_network_buffer(
+ self, empty_request: "EmptyRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def initialize_executor(
+ self, initialize_executor_request: "InitializeExecutorRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def open_executor(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def pause_worker(
+ self, empty_request: "EmptyRequest"
+ ) -> "WorkerStateResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def prepare_checkpoint(
+ self, prepare_checkpoint_request: "PrepareCheckpointRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def query_statistics(
+ self, empty_request: "EmptyRequest"
+ ) -> "WorkerMetricsResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def resume_worker(
+ self, empty_request: "EmptyRequest"
+ ) -> "WorkerStateResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def retrieve_state(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def retry_current_tuple(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def start_worker(
+ self, empty_request: "EmptyRequest"
+ ) -> "WorkerStateResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def end_worker(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def start_channel(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def end_channel(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def debug_command(
+ self, debug_command_request: "DebugCommandRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def evaluate_python_expression(
+ self, evaluate_python_expression_request: "EvaluatePythonExpressionRequest"
+ ) -> "EvaluatedValue":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def no_operation(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def update_executor(
+ self, update_executor_request: "UpdateExecutorRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def __rpc_add_input_channel(
+ self, stream: "grpclib.server.Stream[AddInputChannelRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.add_input_channel(request)
+ await stream.send_message(response)
+
+ async def __rpc_add_partitioning(
+ self, stream: "grpclib.server.Stream[AddPartitioningRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.add_partitioning(request)
+ await stream.send_message(response)
+
+ async def __rpc_assign_port(
+ self, stream: "grpclib.server.Stream[AssignPortRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.assign_port(request)
+ await stream.send_message(response)
+
+ async def __rpc_finalize_checkpoint(
+ self,
+ stream: "grpclib.server.Stream[FinalizeCheckpointRequest, FinalizeCheckpointResponse]",
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.finalize_checkpoint(request)
+ await stream.send_message(response)
+
+ async def __rpc_flush_network_buffer(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.flush_network_buffer(request)
+ await stream.send_message(response)
+
+ async def __rpc_initialize_executor(
+ self, stream: "grpclib.server.Stream[InitializeExecutorRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.initialize_executor(request)
+ await stream.send_message(response)
+
+ async def __rpc_open_executor(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.open_executor(request)
+ await stream.send_message(response)
+
+ async def __rpc_pause_worker(
+ self, stream: "grpclib.server.Stream[EmptyRequest, WorkerStateResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.pause_worker(request)
+ await stream.send_message(response)
+
+ async def __rpc_prepare_checkpoint(
+ self, stream: "grpclib.server.Stream[PrepareCheckpointRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.prepare_checkpoint(request)
+ await stream.send_message(response)
+
+ async def __rpc_query_statistics(
+ self, stream: "grpclib.server.Stream[EmptyRequest, WorkerMetricsResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.query_statistics(request)
+ await stream.send_message(response)
+
+ async def __rpc_resume_worker(
+ self, stream: "grpclib.server.Stream[EmptyRequest, WorkerStateResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.resume_worker(request)
+ await stream.send_message(response)
+
+ async def __rpc_retrieve_state(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.retrieve_state(request)
+ await stream.send_message(response)
+
+ async def __rpc_retry_current_tuple(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.retry_current_tuple(request)
+ await stream.send_message(response)
+
+ async def __rpc_start_worker(
+ self, stream: "grpclib.server.Stream[EmptyRequest, WorkerStateResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.start_worker(request)
+ await stream.send_message(response)
+
+ async def __rpc_end_worker(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.end_worker(request)
+ await stream.send_message(response)
+
+ async def __rpc_start_channel(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.start_channel(request)
+ await stream.send_message(response)
+
+ async def __rpc_end_channel(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.end_channel(request)
+ await stream.send_message(response)
+
+ async def __rpc_debug_command(
+ self, stream: "grpclib.server.Stream[DebugCommandRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.debug_command(request)
+ await stream.send_message(response)
+
+ async def __rpc_evaluate_python_expression(
+ self,
+ stream: "grpclib.server.Stream[EvaluatePythonExpressionRequest, EvaluatedValue]",
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.evaluate_python_expression(request)
+ await stream.send_message(response)
+
+ async def __rpc_no_operation(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.no_operation(request)
+ await stream.send_message(response)
+
+ async def __rpc_update_executor(
+ self, stream: "grpclib.server.Stream[UpdateExecutorRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.update_executor(request)
+ await stream.send_message(response)
+
+ def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
+ return {
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/AddInputChannel": grpclib.const.Handler(
+ self.__rpc_add_input_channel,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ AddInputChannelRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/AddPartitioning": grpclib.const.Handler(
+ self.__rpc_add_partitioning,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ AddPartitioningRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/AssignPort": grpclib.const.Handler(
+ self.__rpc_assign_port,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ AssignPortRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/FinalizeCheckpoint": grpclib.const.Handler(
+ self.__rpc_finalize_checkpoint,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ FinalizeCheckpointRequest,
+ FinalizeCheckpointResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/FlushNetworkBuffer": grpclib.const.Handler(
+ self.__rpc_flush_network_buffer,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/InitializeExecutor": grpclib.const.Handler(
+ self.__rpc_initialize_executor,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ InitializeExecutorRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/OpenExecutor": grpclib.const.Handler(
+ self.__rpc_open_executor,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/PauseWorker": grpclib.const.Handler(
+ self.__rpc_pause_worker,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ WorkerStateResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/PrepareCheckpoint": grpclib.const.Handler(
+ self.__rpc_prepare_checkpoint,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ PrepareCheckpointRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/QueryStatistics": grpclib.const.Handler(
+ self.__rpc_query_statistics,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ WorkerMetricsResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/ResumeWorker": grpclib.const.Handler(
+ self.__rpc_resume_worker,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ WorkerStateResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/RetrieveState": grpclib.const.Handler(
+ self.__rpc_retrieve_state,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/RetryCurrentTuple": grpclib.const.Handler(
+ self.__rpc_retry_current_tuple,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/StartWorker": grpclib.const.Handler(
+ self.__rpc_start_worker,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ WorkerStateResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/EndWorker": grpclib.const.Handler(
+ self.__rpc_end_worker,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/StartChannel": grpclib.const.Handler(
+ self.__rpc_start_channel,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/EndChannel": grpclib.const.Handler(
+ self.__rpc_end_channel,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/DebugCommand": grpclib.const.Handler(
+ self.__rpc_debug_command,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ DebugCommandRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/EvaluatePythonExpression": grpclib.const.Handler(
+ self.__rpc_evaluate_python_expression,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EvaluatePythonExpressionRequest,
+ EvaluatedValue,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/NoOperation": grpclib.const.Handler(
+ self.__rpc_no_operation,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.WorkerService/UpdateExecutor": grpclib.const.Handler(
+ self.__rpc_update_executor,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ UpdateExecutorRequest,
+ EmptyReturn,
+ ),
+ }
+
+
+class ControllerServiceBase(ServiceBase):
+
+ async def retrieve_workflow_state(
+ self, empty_request: "EmptyRequest"
+ ) -> "RetrieveWorkflowStateResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def propagate_embedded_control_message(
+ self,
+ propagate_embedded_control_message_request: "PropagateEmbeddedControlMessageRequest",
+ ) -> "PropagateEmbeddedControlMessageResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def take_global_checkpoint(
+ self, take_global_checkpoint_request: "TakeGlobalCheckpointRequest"
+ ) -> "TakeGlobalCheckpointResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def debug_command(
+ self, debug_command_request: "DebugCommandRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def evaluate_python_expression(
+ self, evaluate_python_expression_request: "EvaluatePythonExpressionRequest"
+ ) -> "EvaluatePythonExpressionResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def console_message_triggered(
+ self, console_message_triggered_request: "ConsoleMessageTriggeredRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def port_completed(
+ self, port_completed_request: "PortCompletedRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def start_workflow(
+ self, empty_request: "EmptyRequest"
+ ) -> "StartWorkflowResponse":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def resume_workflow(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def pause_workflow(self, empty_request: "EmptyRequest") -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def worker_state_updated(
+ self, worker_state_updated_request: "WorkerStateUpdatedRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def worker_execution_completed(
+ self, empty_request: "EmptyRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def link_workers(
+ self, link_workers_request: "LinkWorkersRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def controller_initiate_query_statistics(
+ self, query_statistics_request: "QueryStatisticsRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def retry_workflow(
+ self, retry_workflow_request: "RetryWorkflowRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def reconfigure_workflow(
+ self, workflow_reconfigure_request: "WorkflowReconfigureRequest"
+ ) -> "EmptyReturn":
+ raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+ async def __rpc_retrieve_workflow_state(
+ self,
+ stream: "grpclib.server.Stream[EmptyRequest, RetrieveWorkflowStateResponse]",
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.retrieve_workflow_state(request)
+ await stream.send_message(response)
+
+ async def __rpc_propagate_embedded_control_message(
+ self,
+ stream: "grpclib.server.Stream[PropagateEmbeddedControlMessageRequest, PropagateEmbeddedControlMessageResponse]",
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.propagate_embedded_control_message(request)
+ await stream.send_message(response)
+
+ async def __rpc_take_global_checkpoint(
+ self,
+ stream: "grpclib.server.Stream[TakeGlobalCheckpointRequest, TakeGlobalCheckpointResponse]",
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.take_global_checkpoint(request)
+ await stream.send_message(response)
+
+ async def __rpc_debug_command(
+ self, stream: "grpclib.server.Stream[DebugCommandRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.debug_command(request)
+ await stream.send_message(response)
+
+ async def __rpc_evaluate_python_expression(
+ self,
+ stream: "grpclib.server.Stream[EvaluatePythonExpressionRequest, EvaluatePythonExpressionResponse]",
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.evaluate_python_expression(request)
+ await stream.send_message(response)
+
+ async def __rpc_console_message_triggered(
+ self,
+ stream: "grpclib.server.Stream[ConsoleMessageTriggeredRequest, EmptyReturn]",
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.console_message_triggered(request)
+ await stream.send_message(response)
+
+ async def __rpc_port_completed(
+ self, stream: "grpclib.server.Stream[PortCompletedRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.port_completed(request)
+ await stream.send_message(response)
+
+ async def __rpc_start_workflow(
+ self, stream: "grpclib.server.Stream[EmptyRequest, StartWorkflowResponse]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.start_workflow(request)
+ await stream.send_message(response)
+
+ async def __rpc_resume_workflow(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.resume_workflow(request)
+ await stream.send_message(response)
+
+ async def __rpc_pause_workflow(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.pause_workflow(request)
+ await stream.send_message(response)
+
+ async def __rpc_worker_state_updated(
+ self, stream: "grpclib.server.Stream[WorkerStateUpdatedRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.worker_state_updated(request)
+ await stream.send_message(response)
+
+ async def __rpc_worker_execution_completed(
+ self, stream: "grpclib.server.Stream[EmptyRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.worker_execution_completed(request)
+ await stream.send_message(response)
+
+ async def __rpc_link_workers(
+ self, stream: "grpclib.server.Stream[LinkWorkersRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.link_workers(request)
+ await stream.send_message(response)
+
+ async def __rpc_controller_initiate_query_statistics(
+ self, stream: "grpclib.server.Stream[QueryStatisticsRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.controller_initiate_query_statistics(request)
+ await stream.send_message(response)
+
+ async def __rpc_retry_workflow(
+ self, stream: "grpclib.server.Stream[RetryWorkflowRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.retry_workflow(request)
+ await stream.send_message(response)
+
+ async def __rpc_reconfigure_workflow(
+ self, stream: "grpclib.server.Stream[WorkflowReconfigureRequest, EmptyReturn]"
+ ) -> None:
+ request = await stream.recv_message()
+ response = await self.reconfigure_workflow(request)
+ await stream.send_message(response)
+
+ def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
+ return {
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/RetrieveWorkflowState": grpclib.const.Handler(
+ self.__rpc_retrieve_workflow_state,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ RetrieveWorkflowStateResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/PropagateEmbeddedControlMessage": grpclib.const.Handler(
+ self.__rpc_propagate_embedded_control_message,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ PropagateEmbeddedControlMessageRequest,
+ PropagateEmbeddedControlMessageResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/TakeGlobalCheckpoint": grpclib.const.Handler(
+ self.__rpc_take_global_checkpoint,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ TakeGlobalCheckpointRequest,
+ TakeGlobalCheckpointResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/DebugCommand": grpclib.const.Handler(
+ self.__rpc_debug_command,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ DebugCommandRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/EvaluatePythonExpression": grpclib.const.Handler(
+ self.__rpc_evaluate_python_expression,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EvaluatePythonExpressionRequest,
+ EvaluatePythonExpressionResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/ConsoleMessageTriggered": grpclib.const.Handler(
+ self.__rpc_console_message_triggered,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ ConsoleMessageTriggeredRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/PortCompleted": grpclib.const.Handler(
+ self.__rpc_port_completed,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ PortCompletedRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/StartWorkflow": grpclib.const.Handler(
+ self.__rpc_start_workflow,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ StartWorkflowResponse,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/ResumeWorkflow": grpclib.const.Handler(
+ self.__rpc_resume_workflow,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/PauseWorkflow": grpclib.const.Handler(
+ self.__rpc_pause_workflow,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/WorkerStateUpdated": grpclib.const.Handler(
+ self.__rpc_worker_state_updated,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ WorkerStateUpdatedRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/WorkerExecutionCompleted": grpclib.const.Handler(
+ self.__rpc_worker_execution_completed,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ EmptyRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/LinkWorkers": grpclib.const.Handler(
+ self.__rpc_link_workers,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ LinkWorkersRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/ControllerInitiateQueryStatistics": grpclib.const.Handler(
+ self.__rpc_controller_initiate_query_statistics,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ QueryStatisticsRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/RetryWorkflow": grpclib.const.Handler(
+ self.__rpc_retry_workflow,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ RetryWorkflowRequest,
+ EmptyReturn,
+ ),
+ "/org.apache.texera.amber.engine.architecture.rpc.ControllerService/ReconfigureWorkflow": grpclib.const.Handler(
+ self.__rpc_reconfigure_workflow,
+ grpclib.const.Cardinality.UNARY_UNARY,
+ WorkflowReconfigureRequest,
+ EmptyReturn,
+ ),
+ }
diff --git a/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/sendsemantics/__init__.py b/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/sendsemantics/__init__.py
new file mode 100644
index 00000000000..bc241806b5c
--- /dev/null
+++ b/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/sendsemantics/__init__.py
@@ -0,0 +1,66 @@
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# sources: org/apache/texera/amber/engine/architecture/sendsemantics/partitionings.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+from typing import (
+ List,
+)
+
+import betterproto
+
+from .... import core as ___core__
+
+
+@dataclass(eq=False, repr=False)
+class Partitioning(betterproto.Message):
+ one_to_one_partitioning: "OneToOnePartitioning" = betterproto.message_field(
+ 1, group="sealed_value"
+ )
+ round_robin_partitioning: "RoundRobinPartitioning" = betterproto.message_field(
+ 2, group="sealed_value"
+ )
+ hash_based_shuffle_partitioning: "HashBasedShufflePartitioning" = (
+ betterproto.message_field(3, group="sealed_value")
+ )
+ range_based_shuffle_partitioning: "RangeBasedShufflePartitioning" = (
+ betterproto.message_field(4, group="sealed_value")
+ )
+ broadcast_partitioning: "BroadcastPartitioning" = betterproto.message_field(
+ 5, group="sealed_value"
+ )
+
+
+@dataclass(eq=False, repr=False)
+class OneToOnePartitioning(betterproto.Message):
+ batch_size: int = betterproto.int32_field(1)
+ channels: List["___core__.ChannelIdentity"] = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class RoundRobinPartitioning(betterproto.Message):
+ batch_size: int = betterproto.int32_field(1)
+ channels: List["___core__.ChannelIdentity"] = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class HashBasedShufflePartitioning(betterproto.Message):
+ batch_size: int = betterproto.int32_field(1)
+ channels: List["___core__.ChannelIdentity"] = betterproto.message_field(2)
+ hash_attribute_names: List[str] = betterproto.string_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class RangeBasedShufflePartitioning(betterproto.Message):
+ batch_size: int = betterproto.int32_field(1)
+ channels: List["___core__.ChannelIdentity"] = betterproto.message_field(2)
+ range_attribute_names: List[str] = betterproto.string_field(3)
+ range_min: int = betterproto.int64_field(4)
+ range_max: int = betterproto.int64_field(5)
+
+
+@dataclass(eq=False, repr=False)
+class BroadcastPartitioning(betterproto.Message):
+ batch_size: int = betterproto.int32_field(1)
+ channels: List["___core__.ChannelIdentity"] = betterproto.message_field(2)
diff --git a/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/worker/__init__.py b/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/worker/__init__.py
new file mode 100644
index 00000000000..6a7b210e185
--- /dev/null
+++ b/amber/src/main/python/proto/org/apache/texera/amber/engine/architecture/worker/__init__.py
@@ -0,0 +1,49 @@
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# sources: org/apache/texera/amber/engine/architecture/worker/statistics.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+from typing import (
+ List,
+)
+
+import betterproto
+
+from .... import core as ___core__
+
+
+class WorkerState(betterproto.Enum):
+ UNINITIALIZED = 0
+ READY = 1
+ RUNNING = 2
+ PAUSED = 3
+ COMPLETED = 4
+ TERMINATED = 5
+
+
+@dataclass(eq=False, repr=False)
+class PortTupleMetricsMapping(betterproto.Message):
+ port_id: "___core__.PortIdentity" = betterproto.message_field(1)
+ tuple_metrics: "TupleMetrics" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class TupleMetrics(betterproto.Message):
+ count: int = betterproto.int64_field(1)
+ size: int = betterproto.int64_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class WorkerStatistics(betterproto.Message):
+ input_tuple_metrics: List["PortTupleMetricsMapping"] = betterproto.message_field(1)
+ output_tuple_metrics: List["PortTupleMetricsMapping"] = betterproto.message_field(2)
+ data_processing_time: int = betterproto.int64_field(3)
+ control_processing_time: int = betterproto.int64_field(4)
+ idle_time: int = betterproto.int64_field(5)
+
+
+@dataclass(eq=False, repr=False)
+class WorkerMetrics(betterproto.Message):
+ worker_state: "WorkerState" = betterproto.enum_field(1)
+ worker_statistics: "WorkerStatistics" = betterproto.message_field(2)
diff --git a/amber/src/main/python/proto/org/apache/texera/amber/engine/common/__init__.py b/amber/src/main/python/proto/org/apache/texera/amber/engine/common/__init__.py
new file mode 100644
index 00000000000..55c789aa395
--- /dev/null
+++ b/amber/src/main/python/proto/org/apache/texera/amber/engine/common/__init__.py
@@ -0,0 +1,156 @@
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# sources: org/apache/texera/amber/engine/common/actormessage.proto, org/apache/texera/amber/engine/common/ambermessage.proto, org/apache/texera/amber/engine/common/executionruntimestate.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+from typing import (
+ Dict,
+ List,
+)
+
+import betterproto
+
+from ... import core as __core__
+from ..architecture import (
+ rpc as _architecture_rpc__,
+ worker as _architecture_worker__,
+)
+
+
+@dataclass(eq=False, repr=False)
+class DirectControlMessagePayloadV2(betterproto.Message):
+ control_invocation: "_architecture_rpc__.ControlInvocation" = (
+ betterproto.message_field(1, group="value")
+ )
+ return_invocation: "_architecture_rpc__.ReturnInvocation" = (
+ betterproto.message_field(2, group="value")
+ )
+
+
+@dataclass(eq=False, repr=False)
+class PythonDataHeader(betterproto.Message):
+ tag: "__core__.ChannelIdentity" = betterproto.message_field(1)
+ payload_type: str = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class PythonControlMessage(betterproto.Message):
+ tag: "__core__.ChannelIdentity" = betterproto.message_field(1)
+ payload: "DirectControlMessagePayloadV2" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class BreakpointFault(betterproto.Message):
+ worker_name: str = betterproto.string_field(1)
+ faulted_tuple: "BreakpointFaultBreakpointTuple" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class BreakpointFaultBreakpointTuple(betterproto.Message):
+ id: int = betterproto.int64_field(1)
+ is_input: bool = betterproto.bool_field(2)
+ tuple: List[str] = betterproto.string_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class OperatorBreakpoints(betterproto.Message):
+ unresolved_breakpoints: List["BreakpointFault"] = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionBreakpointStore(betterproto.Message):
+ operator_info: Dict[str, "OperatorBreakpoints"] = betterproto.map_field(
+ 1, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+
+
+@dataclass(eq=False, repr=False)
+class EvaluatedValueList(betterproto.Message):
+ values: List["_architecture_rpc__.EvaluatedValue"] = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class OperatorConsole(betterproto.Message):
+ console_messages: List["_architecture_rpc__.ConsoleMessage"] = (
+ betterproto.message_field(1)
+ )
+ evaluate_expr_results: Dict[str, "EvaluatedValueList"] = betterproto.map_field(
+ 2, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionConsoleStore(betterproto.Message):
+ operator_console: Dict[str, "OperatorConsole"] = betterproto.map_field(
+ 1, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+
+
+@dataclass(eq=False, repr=False)
+class OperatorWorkerMapping(betterproto.Message):
+ operator_id: str = betterproto.string_field(1)
+ worker_ids: List[str] = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class OperatorStatistics(betterproto.Message):
+ input_metrics: List["_architecture_worker__.PortTupleMetricsMapping"] = (
+ betterproto.message_field(1)
+ )
+ output_metrics: List["_architecture_worker__.PortTupleMetricsMapping"] = (
+ betterproto.message_field(2)
+ )
+ num_workers: int = betterproto.int32_field(3)
+ data_processing_time: int = betterproto.int64_field(4)
+ control_processing_time: int = betterproto.int64_field(5)
+ idle_time: int = betterproto.int64_field(6)
+
+
+@dataclass(eq=False, repr=False)
+class OperatorMetrics(betterproto.Message):
+ operator_state: "_architecture_rpc__.WorkflowAggregatedState" = (
+ betterproto.enum_field(1)
+ )
+ operator_statistics: "OperatorStatistics" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionStatsStore(betterproto.Message):
+ start_time_stamp: int = betterproto.int64_field(1)
+ end_time_stamp: int = betterproto.int64_field(2)
+ operator_info: Dict[str, "OperatorMetrics"] = betterproto.map_field(
+ 3, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+ operator_worker_mapping: List["OperatorWorkerMapping"] = betterproto.message_field(
+ 4
+ )
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionMetadataStore(betterproto.Message):
+ state: "_architecture_rpc__.WorkflowAggregatedState" = betterproto.enum_field(1)
+ fatal_errors: List["__core__.WorkflowFatalError"] = betterproto.message_field(2)
+ execution_id: "__core__.ExecutionIdentity" = betterproto.message_field(3)
+ is_recovering: bool = betterproto.bool_field(4)
+
+
+@dataclass(eq=False, repr=False)
+class Backpressure(betterproto.Message):
+ enable_backpressure: bool = betterproto.bool_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class CreditUpdate(betterproto.Message):
+ pass
+
+
+@dataclass(eq=False, repr=False)
+class ActorCommand(betterproto.Message):
+ backpressure: "Backpressure" = betterproto.message_field(1, group="sealed_value")
+ credit_update: "CreditUpdate" = betterproto.message_field(2, group="sealed_value")
+
+
+@dataclass(eq=False, repr=False)
+class PythonActorMessage(betterproto.Message):
+ payload: "ActorCommand" = betterproto.message_field(1)
diff --git a/amber/src/main/python/proto/org/apache/texera/web/__init__.py b/amber/src/main/python/proto/org/apache/texera/web/__init__.py
new file mode 100644
index 00000000000..adb5848bb0c
--- /dev/null
+++ b/amber/src/main/python/proto/org/apache/texera/web/__init__.py
@@ -0,0 +1,158 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# sources: org/apache/texera/workflowruntimestate.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import (
+ Dict,
+ List,
+)
+
+import betterproto
+
+from ...amber.engine import common as __amber_engine_common__
+from ...amber.engine.architecture import worker as __amber_engine_architecture_worker__
+
+
+class FatalErrorType(betterproto.Enum):
+ COMPILATION_ERROR = 0
+ EXECUTION_FAILURE = 1
+
+
+class WorkflowAggregatedState(betterproto.Enum):
+ UNINITIALIZED = 0
+ READY = 1
+ RUNNING = 2
+ PAUSING = 3
+ PAUSED = 4
+ RESUMING = 5
+ COMPLETED = 6
+ FAILED = 7
+ UNKNOWN = 8
+ KILLED = 9
+
+
+@dataclass(eq=False, repr=False)
+class BreakpointFault(betterproto.Message):
+ worker_name: str = betterproto.string_field(1)
+ faulted_tuple: "BreakpointFaultBreakpointTuple" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class BreakpointFaultBreakpointTuple(betterproto.Message):
+ id: int = betterproto.int64_field(1)
+ is_input: bool = betterproto.bool_field(2)
+ tuple: List[str] = betterproto.string_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class OperatorBreakpoints(betterproto.Message):
+ unresolved_breakpoints: List["BreakpointFault"] = betterproto.message_field(1)
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionBreakpointStore(betterproto.Message):
+ operator_info: Dict[str, "OperatorBreakpoints"] = betterproto.map_field(
+ 1, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+
+
+@dataclass(eq=False, repr=False)
+class EvaluatedValueList(betterproto.Message):
+ values: List["__amber_engine_architecture_worker__.EvaluatedValue"] = (
+ betterproto.message_field(1)
+ )
+
+
+@dataclass(eq=False, repr=False)
+class OperatorConsole(betterproto.Message):
+ console_messages: List["__amber_engine_architecture_worker__.ConsoleMessage"] = (
+ betterproto.message_field(1)
+ )
+ evaluate_expr_results: Dict[str, "EvaluatedValueList"] = betterproto.map_field(
+ 2, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionConsoleStore(betterproto.Message):
+ operator_console: Dict[str, "OperatorConsole"] = betterproto.map_field(
+ 1, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+
+
+@dataclass(eq=False, repr=False)
+class OperatorWorkerMapping(betterproto.Message):
+ operator_id: str = betterproto.string_field(1)
+ worker_ids: List[str] = betterproto.string_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class OperatorStatistics(betterproto.Message):
+ input_count: List["__amber_engine_architecture_worker__.PortTupleCountMapping"] = (
+ betterproto.message_field(1)
+ )
+ output_count: List["__amber_engine_architecture_worker__.PortTupleCountMapping"] = (
+ betterproto.message_field(2)
+ )
+ num_workers: int = betterproto.int32_field(3)
+ data_processing_time: int = betterproto.int64_field(4)
+ control_processing_time: int = betterproto.int64_field(5)
+ idle_time: int = betterproto.int64_field(6)
+
+
+@dataclass(eq=False, repr=False)
+class OperatorMetrics(betterproto.Message):
+ operator_state: "WorkflowAggregatedState" = betterproto.enum_field(1)
+ operator_statistics: "OperatorStatistics" = betterproto.message_field(2)
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionStatsStore(betterproto.Message):
+ start_time_stamp: int = betterproto.int64_field(1)
+ end_time_stamp: int = betterproto.int64_field(2)
+ operator_info: Dict[str, "OperatorMetrics"] = betterproto.map_field(
+ 3, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
+ operator_worker_mapping: List["OperatorWorkerMapping"] = betterproto.message_field(
+ 4
+ )
+
+
+@dataclass(eq=False, repr=False)
+class WorkflowFatalError(betterproto.Message):
+ type: "FatalErrorType" = betterproto.enum_field(1)
+ timestamp: datetime = betterproto.message_field(2)
+ message: str = betterproto.string_field(3)
+ details: str = betterproto.string_field(4)
+ operator_id: str = betterproto.string_field(5)
+ worker_id: str = betterproto.string_field(6)
+
+
+@dataclass(eq=False, repr=False)
+class ExecutionMetadataStore(betterproto.Message):
+ state: "WorkflowAggregatedState" = betterproto.enum_field(1)
+ fatal_errors: List["WorkflowFatalError"] = betterproto.message_field(2)
+ execution_id: "__amber_engine_common__.ExecutionIdentity" = (
+ betterproto.message_field(3)
+ )
+ is_recovering: bool = betterproto.bool_field(4)
diff --git a/amber/src/main/python/proto/scalapb/__init__.py b/amber/src/main/python/proto/scalapb/__init__.py
new file mode 100644
index 00000000000..49c713815a5
--- /dev/null
+++ b/amber/src/main/python/proto/scalapb/__init__.py
@@ -0,0 +1,421 @@
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# sources: scalapb/scalapb.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+from typing import (
+ Dict,
+ List,
+)
+
+import betterproto
+import betterproto.lib.google.protobuf as betterproto_lib_google_protobuf
+
+
+class MatchType(betterproto.Enum):
+ CONTAINS = 0
+ EXACT = 1
+ PRESENCE = 2
+
+
+class ScalaPbOptionsOptionsScope(betterproto.Enum):
+ """
+ Whether to apply the options only to this file, or for the entire package (and its subpackages)
+ """
+
+ FILE = 0
+ """Apply the options for this file only (default)"""
+
+ PACKAGE = 1
+ """Apply the options for the entire package and its subpackages."""
+
+
+class ScalaPbOptionsEnumValueNaming(betterproto.Enum):
+ """Naming convention for generated enum values"""
+
+ AS_IN_PROTO = 0
+ CAMEL_CASE = 1
+
+
+@dataclass(eq=False, repr=False)
+class ScalaPbOptions(betterproto.Message):
+ package_name: str = betterproto.string_field(1)
+ """If set then it overrides the java_package and package."""
+
+ flat_package: bool = betterproto.bool_field(2)
+ """
+ If true, the compiler does not append the proto base file name
+ into the generated package name. If false (the default), the
+ generated scala package name is the package_name.basename where
+ basename is the proto file name without the .proto extension.
+ """
+
+ import_: List[str] = betterproto.string_field(3)
+ """
+ Adds the following imports at the top of the file (this is meant
+ to provide implicit TypeMappers)
+ """
+
+ preamble: List[str] = betterproto.string_field(4)
+ """
+ Text to add to the generated scala file. This can be used only
+ when single_file is true.
+ """
+
+ single_file: bool = betterproto.bool_field(5)
+ """
+ If true, all messages and enums (but not services) will be written
+ to a single Scala file.
+ """
+
+ no_primitive_wrappers: bool = betterproto.bool_field(7)
+ """
+ By default, wrappers defined at
+ https://github.com/google/protobuf/blob/master/src/google/protobuf/wrappers.proto,
+ are mapped to an Option[T] where T is a primitive type. When this field
+ is set to true, we do not perform this transformation.
+ """
+
+ primitive_wrappers: bool = betterproto.bool_field(6)
+ """
+ DEPRECATED. In ScalaPB <= 0.5.47, it was necessary to explicitly enable
+ primitive_wrappers. This field remains here for backwards compatibility,
+ but it has no effect on generated code. It is an error to set both
+ `primitive_wrappers` and `no_primitive_wrappers`.
+ """
+
+ collection_type: str = betterproto.string_field(8)
+ """
+ Scala type to be used for repeated fields. If unspecified,
+ `scala.collection.Seq` will be used.
+ """
+
+ preserve_unknown_fields: bool = betterproto.bool_field(9)
+ """
+ If set to true, all generated messages in this file will preserve unknown
+ fields.
+ """
+
+ object_name: str = betterproto.string_field(10)
+ """
+ If defined, sets the name of the file-level object that would be generated. This
+ object extends `GeneratedFileObject` and contains descriptors, and list of message
+ and enum companions.
+ """
+
+ scope: "ScalaPbOptionsOptionsScope" = betterproto.enum_field(11)
+ """Experimental: scope to apply the given options."""
+
+ lenses: bool = betterproto.bool_field(12)
+ """If true, lenses will be generated."""
+
+ retain_source_code_info: bool = betterproto.bool_field(13)
+ """
+ If true, then source-code info information will be included in the
+ generated code - normally the source code info is cleared out to reduce
+ code size. The source code info is useful for extracting source code
+ location from the descriptors as well as comments.
+ """
+
+ map_type: str = betterproto.string_field(14)
+ """
+ Scala type to be used for maps. If unspecified,
+ `scala.collection.immutable.Map` will be used.
+ """
+
+ no_default_values_in_constructor: bool = betterproto.bool_field(15)
+ """
+ If true, no default values will be generated in message constructors.
+ """
+
+ enum_value_naming: "ScalaPbOptionsEnumValueNaming" = betterproto.enum_field(16)
+ enum_strip_prefix: bool = betterproto.bool_field(17)
+ """
+ Indicate if prefix (enum name + optional underscore) should be removed in scala code
+ Strip is applied before enum value naming changes.
+ """
+
+ bytes_type: str = betterproto.string_field(21)
+ """Scala type to use for bytes fields."""
+
+ java_conversions: bool = betterproto.bool_field(23)
+ """Enable java conversions for this file."""
+
+ aux_message_options: List["ScalaPbOptionsAuxMessageOptions"] = (
+ betterproto.message_field(18)
+ )
+ """List of message options to apply to some messages."""
+
+ aux_field_options: List["ScalaPbOptionsAuxFieldOptions"] = (
+ betterproto.message_field(19)
+ )
+ """List of message options to apply to some fields."""
+
+ aux_enum_options: List["ScalaPbOptionsAuxEnumOptions"] = betterproto.message_field(
+ 20
+ )
+ """List of message options to apply to some enums."""
+
+ aux_enum_value_options: List["ScalaPbOptionsAuxEnumValueOptions"] = (
+ betterproto.message_field(22)
+ )
+ """List of enum value options to apply to some enum values."""
+
+ preprocessors: List[str] = betterproto.string_field(24)
+ """List of preprocessors to apply."""
+
+ field_transformations: List["FieldTransformation"] = betterproto.message_field(25)
+ ignore_all_transformations: bool = betterproto.bool_field(26)
+ """
+ Ignores all transformations for this file. This is meant to allow specific files to
+ opt out from transformations inherited through package-scoped options.
+ """
+
+ getters: bool = betterproto.bool_field(27)
+ """If true, getters will be generated."""
+
+ test_only_no_java_conversions: bool = betterproto.bool_field(999)
+ """
+ For use in tests only. Inhibit Java conversions even when when generator parameters
+ request for it.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class ScalaPbOptionsAuxMessageOptions(betterproto.Message):
+ """
+ AuxMessageOptions enables you to set message-level options through package-scoped options.
+ This is useful when you can't add a dependency on scalapb.proto from the proto file that
+ defines the message.
+ """
+
+ target: str = betterproto.string_field(1)
+ """The fully-qualified name of the message in the proto name space."""
+
+ options: "MessageOptions" = betterproto.message_field(2)
+ """
+ Options to apply to the message. If there are any options defined on the target message
+ they take precedence over the options.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class ScalaPbOptionsAuxFieldOptions(betterproto.Message):
+ """
+ AuxFieldOptions enables you to set field-level options through package-scoped options.
+ This is useful when you can't add a dependency on scalapb.proto from the proto file that
+ defines the field.
+ """
+
+ target: str = betterproto.string_field(1)
+ """The fully-qualified name of the field in the proto name space."""
+
+ options: "FieldOptions" = betterproto.message_field(2)
+ """
+ Options to apply to the field. If there are any options defined on the target message
+ they take precedence over the options.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class ScalaPbOptionsAuxEnumOptions(betterproto.Message):
+ """
+ AuxEnumOptions enables you to set enum-level options through package-scoped options.
+ This is useful when you can't add a dependency on scalapb.proto from the proto file that
+ defines the enum.
+ """
+
+ target: str = betterproto.string_field(1)
+ """The fully-qualified name of the enum in the proto name space."""
+
+ options: "EnumOptions" = betterproto.message_field(2)
+ """
+ Options to apply to the enum. If there are any options defined on the target enum
+ they take precedence over the options.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class ScalaPbOptionsAuxEnumValueOptions(betterproto.Message):
+ """
+ AuxEnumValueOptions enables you to set enum value level options through package-scoped
+ options. This is useful when you can't add a dependency on scalapb.proto from the proto
+ file that defines the enum.
+ """
+
+ target: str = betterproto.string_field(1)
+ """The fully-qualified name of the enum value in the proto name space."""
+
+ options: "EnumValueOptions" = betterproto.message_field(2)
+ """
+ Options to apply to the enum value. If there are any options defined on
+ the target enum value they take precedence over the options.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class MessageOptions(betterproto.Message):
+ extends: List[str] = betterproto.string_field(1)
+ """Additional classes and traits to mix in to the case class."""
+
+ companion_extends: List[str] = betterproto.string_field(2)
+ """Additional classes and traits to mix in to the companion object."""
+
+ annotations: List[str] = betterproto.string_field(3)
+ """Custom annotations to add to the generated case class."""
+
+ type: str = betterproto.string_field(4)
+ """
+ All instances of this message will be converted to this type. An implicit TypeMapper
+ must be present.
+ """
+
+ companion_annotations: List[str] = betterproto.string_field(5)
+ """
+ Custom annotations to add to the companion object of the generated class.
+ """
+
+ sealed_oneof_extends: List[str] = betterproto.string_field(6)
+ """
+ Additional classes and traits to mix in to generated sealed_oneof base trait.
+ """
+
+ no_box: bool = betterproto.bool_field(7)
+ """
+ If true, when this message is used as an optional field, do not wrap it in an `Option`.
+ This is equivalent of setting `(field).no_box` to true on each field with the message type.
+ """
+
+ unknown_fields_annotations: List[str] = betterproto.string_field(8)
+ """
+ Custom annotations to add to the generated `unknownFields` case class field.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class Collection(betterproto.Message):
+ """
+ Represents a custom Collection type in Scala. This allows ScalaPB to integrate with
+ collection types that are different enough from the ones in the standard library.
+ """
+
+ type: str = betterproto.string_field(1)
+ """Type of the collection"""
+
+ non_empty: bool = betterproto.bool_field(2)
+ """
+ Set to true if this collection type is not allowed to be empty, for example
+ cats.data.NonEmptyList. When true, ScalaPB will not generate `clearX` for the repeated
+ field and not provide a default argument in the constructor.
+ """
+
+ adapter: str = betterproto.string_field(3)
+ """
+ An Adapter is a Scala object available at runtime that provides certain static methods
+ that can operate on this collection type.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class FieldOptions(betterproto.Message):
+ type: str = betterproto.string_field(1)
+ scala_name: str = betterproto.string_field(2)
+ collection_type: str = betterproto.string_field(3)
+ """
+ Can be specified only if this field is repeated. If unspecified,
+ it falls back to the file option named `collection_type`, which defaults
+ to `scala.collection.Seq`.
+ """
+
+ collection: "Collection" = betterproto.message_field(8)
+ key_type: str = betterproto.string_field(4)
+ """
+ If the field is a map, you can specify custom Scala types for the key
+ or value.
+ """
+
+ value_type: str = betterproto.string_field(5)
+ annotations: List[str] = betterproto.string_field(6)
+ """Custom annotations to add to the field."""
+
+ map_type: str = betterproto.string_field(7)
+ """
+ Can be specified only if this field is a map. If unspecified,
+ it falls back to the file option named `map_type` which defaults to
+ `scala.collection.immutable.Map`
+ """
+
+ no_box: bool = betterproto.bool_field(30)
+ """
+ Do not box this value in Option[T]. If set, this overrides MessageOptions.no_box
+ """
+
+ required: bool = betterproto.bool_field(31)
+ """
+ Like no_box it does not box a value in Option[T], but also fails parsing when a value
+ is not provided. This enables to emulate required fields in proto3.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class EnumOptions(betterproto.Message):
+ extends: List[str] = betterproto.string_field(1)
+ """Additional classes and traits to mix in to the base trait"""
+
+ companion_extends: List[str] = betterproto.string_field(2)
+ """Additional classes and traits to mix in to the companion object."""
+
+ type: str = betterproto.string_field(3)
+ """
+ All instances of this enum will be converted to this type. An implicit TypeMapper
+ must be present.
+ """
+
+ base_annotations: List[str] = betterproto.string_field(4)
+ """Custom annotations to add to the generated enum's base class."""
+
+ recognized_annotations: List[str] = betterproto.string_field(5)
+ """Custom annotations to add to the generated trait."""
+
+ unrecognized_annotations: List[str] = betterproto.string_field(6)
+ """Custom annotations to add to the generated Unrecognized case class."""
+
+
+@dataclass(eq=False, repr=False)
+class EnumValueOptions(betterproto.Message):
+ extends: List[str] = betterproto.string_field(1)
+ """Additional classes and traits to mix in to an individual enum value."""
+
+ scala_name: str = betterproto.string_field(2)
+ """Name in Scala to use for this enum value."""
+
+ annotations: List[str] = betterproto.string_field(3)
+ """
+ Custom annotations to add to the generated case object for this enum value.
+ """
+
+
+@dataclass(eq=False, repr=False)
+class OneofOptions(betterproto.Message):
+ extends: List[str] = betterproto.string_field(1)
+ """Additional traits to mix in to a oneof."""
+
+ scala_name: str = betterproto.string_field(2)
+ """Name in Scala to use for this oneof field."""
+
+
+@dataclass(eq=False, repr=False)
+class FieldTransformation(betterproto.Message):
+ when: "betterproto_lib_google_protobuf.FieldDescriptorProto" = (
+ betterproto.message_field(1)
+ )
+ match_type: "MatchType" = betterproto.enum_field(2)
+ set: "betterproto_lib_google_protobuf.FieldOptions" = betterproto.message_field(3)
+
+
+@dataclass(eq=False, repr=False)
+class PreprocessorOutput(betterproto.Message):
+ options_by_file: Dict[str, "ScalaPbOptions"] = betterproto.map_field(
+ 1, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
+ )
diff --git a/amber/src/main/python/pyamber/__init__.py b/amber/src/main/python/pyamber/__init__.py
new file mode 100644
index 00000000000..01ee5e08279
--- /dev/null
+++ b/amber/src/main/python/pyamber/__init__.py
@@ -0,0 +1,44 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from core.models import (
+ Tuple,
+ TupleLike,
+ Table,
+ TableLike,
+ Batch,
+ BatchLike,
+ TableOperator,
+ BatchOperator,
+ SourceOperator,
+ TupleOperatorV2,
+ State,
+)
+
+__all__ = [
+ "Tuple",
+ "TupleLike",
+ "Table",
+ "TableLike",
+ "Batch",
+ "BatchLike",
+ "TableOperator",
+ "BatchOperator",
+ "TupleOperatorV2",
+ "SourceOperator",
+ "State",
+]
diff --git a/amber/src/main/python/pyproject.toml b/amber/src/main/python/pyproject.toml
new file mode 100644
index 00000000000..72bfeb57245
--- /dev/null
+++ b/amber/src/main/python/pyproject.toml
@@ -0,0 +1,27 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+extend-exclude = ["proto"]
+
+[tool.ruff.lint]
+ignore = ["F403", "F405", "E203"]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 10
\ No newline at end of file
diff --git a/amber/src/main/python/pytexera/__init__.py b/amber/src/main/python/pytexera/__init__.py
new file mode 100644
index 00000000000..e40d1a43fe0
--- /dev/null
+++ b/amber/src/main/python/pytexera/__init__.py
@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from loguru import logger
+from overrides import overrides
+from typing import Iterator, Optional, Union
+
+from pyamber import *
+from .storage.dataset_file_document import DatasetFileDocument
+from .storage.large_binary_input_stream import LargeBinaryInputStream
+from .storage.large_binary_output_stream import LargeBinaryOutputStream
+from .udf.udf_operator import (
+ UDFOperatorV2,
+ UDFTableOperator,
+ UDFBatchOperator,
+ UDFSourceOperator,
+)
+from core.models.type.large_binary import largebinary
+
+__all__ = [
+ "State",
+ "Tuple",
+ "TupleLike",
+ "UDFOperatorV2",
+ "Table",
+ "TableLike",
+ "Batch",
+ "BatchLike",
+ "UDFTableOperator",
+ "UDFBatchOperator",
+ "UDFSourceOperator",
+ "DatasetFileDocument",
+ "largebinary",
+ "LargeBinaryInputStream",
+ "LargeBinaryOutputStream",
+ # export external tools to be used
+ "overrides",
+ "logger",
+ "Iterator",
+ "Optional",
+ "Union",
+]
diff --git a/amber/src/main/python/pytexera/storage/__init__.py b/amber/src/main/python/pytexera/storage/__init__.py
new file mode 100644
index 00000000000..da1b6392b23
--- /dev/null
+++ b/amber/src/main/python/pytexera/storage/__init__.py
@@ -0,0 +1,20 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .dataset_file_document import DatasetFileDocument
+
+__all__ = ["DatasetFileDocument"]
diff --git a/amber/src/main/python/pytexera/storage/dataset_file_document.py b/amber/src/main/python/pytexera/storage/dataset_file_document.py
new file mode 100644
index 00000000000..3d077735833
--- /dev/null
+++ b/amber/src/main/python/pytexera/storage/dataset_file_document.py
@@ -0,0 +1,97 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import io
+import os
+import requests
+import urllib.parse
+
+
+class DatasetFileDocument:
+ def __init__(self, file_path: str):
+ """
+ Parses the file path into dataset metadata.
+
+ :param file_path:
+ Expected format - "/ownerEmail/datasetName/versionName/fileRelativePath"
+ Example: "/bob@texera.com/twitterDataset/v1/california/irvine/tw1.csv"
+ """
+ parts = file_path.strip("/").split("/")
+ if len(parts) < 4:
+ raise ValueError(
+ "Invalid file path format. "
+ "Expected: /ownerEmail/datasetName/versionName/fileRelativePath"
+ )
+
+ self.owner_email = parts[0]
+ self.dataset_name = parts[1]
+ self.version_name = parts[2]
+ self.file_relative_path = "/".join(parts[3:])
+
+ self.jwt_token = os.getenv("USER_JWT_TOKEN")
+ self.presign_endpoint = os.getenv("FILE_SERVICE_GET_PRESIGNED_URL_ENDPOINT")
+
+ if not self.jwt_token:
+ raise ValueError(
+ "JWT token is required but not set in environment variables."
+ )
+ if not self.presign_endpoint:
+ self.presign_endpoint = "http://localhost:9092/api/dataset/presign-download"
+
+ def get_presigned_url(self) -> str:
+ """
+ Requests a presigned URL from the API.
+
+ :return: The presigned URL as a string.
+ :raises: RuntimeError if the request fails.
+ """
+ headers = {"Authorization": f"Bearer {self.jwt_token}"}
+ encoded_file_path = urllib.parse.quote(
+ f"/{self.owner_email}"
+ f"/{self.dataset_name}"
+ f"/{self.version_name}"
+ f"/{self.file_relative_path}"
+ )
+
+ params = {"filePath": encoded_file_path}
+
+ response = requests.get(self.presign_endpoint, headers=headers, params=params)
+
+ if response.status_code != 200:
+ raise RuntimeError(
+ f"Failed to get presigned URL: {response.status_code} {response.text}"
+ )
+
+ return response.json().get("presignedUrl")
+
+ def read_file(self) -> io.BytesIO:
+ """
+ Reads the file content from the presigned URL.
+
+ :return: A file-like object.
+ :raises: RuntimeError if the retrieval fails.
+ """
+ presigned_url = self.get_presigned_url()
+ response = requests.get(presigned_url)
+
+ if response.status_code != 200:
+ raise RuntimeError(
+ f"Failed to retrieve file content: "
+ f"{response.status_code} {response.text}"
+ )
+
+ return io.BytesIO(response.content)
diff --git a/amber/src/main/python/pytexera/storage/large_binary_input_stream.py b/amber/src/main/python/pytexera/storage/large_binary_input_stream.py
new file mode 100644
index 00000000000..8e7d8640403
--- /dev/null
+++ b/amber/src/main/python/pytexera/storage/large_binary_input_stream.py
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+LargeBinaryInputStream for reading largebinary data from S3.
+
+Usage:
+ with LargeBinaryInputStream(large_binary) as stream:
+ content = stream.read()
+"""
+
+from typing import BinaryIO, Optional
+from functools import wraps
+from io import IOBase
+from core.models.type.large_binary import largebinary
+
+
+def _require_open(func):
+ """Decorator to ensure stream is open before reading operations."""
+
+ @wraps(func)
+ def wrapper(self, *args, **kwargs):
+ if self._closed:
+ raise ValueError("I/O operation on closed stream")
+ if self._underlying is None:
+ self._lazy_init()
+ return func(self, *args, **kwargs)
+
+ return wrapper
+
+
+class LargeBinaryInputStream(IOBase):
+ """
+ InputStream for reading largebinary data from S3.
+
+ Lazily downloads from S3 on first read. Supports context manager and iteration.
+ """
+
+ def __init__(self, large_binary: largebinary):
+ """Initialize stream for reading the given largebinary."""
+ super().__init__()
+ if large_binary is None:
+ raise ValueError("largebinary cannot be None")
+ self._large_binary = large_binary
+ self._underlying: Optional[BinaryIO] = None
+ self._closed = False
+
+ def _lazy_init(self):
+ """Download from S3 on first read operation."""
+ from pytexera.storage import large_binary_manager
+
+ s3 = large_binary_manager._get_s3_client()
+ response = s3.get_object(
+ Bucket=self._large_binary.get_bucket_name(),
+ Key=self._large_binary.get_object_key(),
+ )
+ self._underlying = response["Body"]
+
+ @_require_open
+ def read(self, n: int = -1) -> bytes:
+ """Read and return up to n bytes (-1 reads all)."""
+ return self._underlying.read(n)
+
+ @_require_open
+ def readline(self, size: int = -1) -> bytes:
+ """Read and return one line from the stream."""
+ return self._underlying.readline(size)
+
+ @_require_open
+ def readlines(self, hint: int = -1) -> list[bytes]:
+ """Read and return a list of lines from the stream."""
+ return self._underlying.readlines(hint)
+
+ def readable(self) -> bool:
+ """Return True if the stream can be read from."""
+ return not self._closed
+
+ def seekable(self) -> bool:
+ """Return False - this stream does not support seeking."""
+ return False
+
+ @property
+ def closed(self) -> bool:
+ """Return True if the stream is closed."""
+ return self._closed
+
+ def close(self) -> None:
+ """Close the stream and release resources."""
+ if not self._closed:
+ self._closed = True
+ if self._underlying is not None:
+ self._underlying.close()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.close()
+
+ def __iter__(self):
+ return self
+
+ def __next__(self) -> bytes:
+ line = self.readline()
+ if not line:
+ raise StopIteration
+ return line
diff --git a/amber/src/main/python/pytexera/storage/large_binary_manager.py b/amber/src/main/python/pytexera/storage/large_binary_manager.py
new file mode 100644
index 00000000000..e061eac6228
--- /dev/null
+++ b/amber/src/main/python/pytexera/storage/large_binary_manager.py
@@ -0,0 +1,78 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Internal largebinary manager for S3 operations.
+
+Users should not interact with this module directly. Use largebinary() constructor
+and LargeBinaryInputStream/LargeBinaryOutputStream instead.
+"""
+
+import time
+import uuid
+from loguru import logger
+from core.storage.storage_config import StorageConfig
+
+# Module-level state
+_s3_client = None
+DEFAULT_BUCKET = "texera-large-binaries"
+
+
+def _get_s3_client():
+ """Get or initialize S3 client (lazy initialization, cached)."""
+ global _s3_client
+ if _s3_client is None:
+ try:
+ import boto3
+ from botocore.config import Config
+ except ImportError as e:
+ raise RuntimeError("boto3 required. Install with: pip install boto3") from e
+
+ _s3_client = boto3.client(
+ "s3",
+ endpoint_url=StorageConfig.S3_ENDPOINT,
+ aws_access_key_id=StorageConfig.S3_AUTH_USERNAME,
+ aws_secret_access_key=StorageConfig.S3_AUTH_PASSWORD,
+ region_name=StorageConfig.S3_REGION,
+ config=Config(signature_version="s3v4", s3={"addressing_style": "path"}),
+ )
+ return _s3_client
+
+
+def _ensure_bucket_exists(bucket: str):
+ """Ensure S3 bucket exists, creating it if necessary."""
+ s3 = _get_s3_client()
+ try:
+ s3.head_bucket(Bucket=bucket)
+ except s3.exceptions.NoSuchBucket:
+ logger.debug(f"Bucket {bucket} not found, creating it")
+ s3.create_bucket(Bucket=bucket)
+ logger.info(f"Created bucket: {bucket}")
+
+
+def create() -> str:
+ """
+ Creates a new largebinary reference with a unique S3 URI.
+
+ Returns:
+ S3 URI string (format: s3://bucket/key)
+ """
+ _ensure_bucket_exists(DEFAULT_BUCKET)
+ timestamp_ms = int(time.time() * 1000)
+ unique_id = uuid.uuid4()
+ object_key = f"objects/{timestamp_ms}/{unique_id}"
+ return f"s3://{DEFAULT_BUCKET}/{object_key}"
diff --git a/amber/src/main/python/pytexera/storage/large_binary_output_stream.py b/amber/src/main/python/pytexera/storage/large_binary_output_stream.py
new file mode 100644
index 00000000000..af4f1a275c2
--- /dev/null
+++ b/amber/src/main/python/pytexera/storage/large_binary_output_stream.py
@@ -0,0 +1,244 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+LargeBinaryOutputStream for streaming largebinary data to S3.
+
+Usage:
+ from pytexera import largebinary, LargeBinaryOutputStream
+
+ large_binary = largebinary()
+ with LargeBinaryOutputStream(large_binary) as out:
+ out.write(b"data")
+"""
+
+from typing import Optional, Union
+from io import IOBase
+from core.models.type.large_binary import largebinary
+from pytexera.storage import large_binary_manager
+import threading
+import queue
+
+# Constants
+_CHUNK_SIZE = 64 * 1024 # 64KB
+_QUEUE_TIMEOUT = 0.1
+
+
+class _QueueReader:
+ """File-like object that reads from a queue."""
+
+ def __init__(self, q: queue.Queue):
+ self._queue = q
+ self._buffer = b""
+ self._eof = False
+
+ def read(self, size=-1):
+ """Read bytes from the queue."""
+ if self._eof and not self._buffer:
+ return b""
+
+ # Collect chunks until we have enough data or reach EOF
+ chunks = [self._buffer] if self._buffer else []
+ total_size = len(self._buffer)
+ self._buffer = b""
+ needed = size if size != -1 else None
+
+ while not self._eof and (needed is None or total_size < needed):
+ try:
+ chunk = self._queue.get(timeout=_QUEUE_TIMEOUT)
+ if chunk is None: # EOF marker
+ self._eof = True
+ break
+ chunks.append(chunk)
+ total_size += len(chunk)
+ except queue.Empty:
+ continue
+
+ result = b"".join(chunks)
+
+ # If size was specified, split and buffer remainder
+ if needed is not None and len(result) > needed:
+ self._buffer = result[needed:]
+ result = result[:needed]
+
+ return result
+
+
+class LargeBinaryOutputStream(IOBase):
+ """
+ OutputStream for streaming largebinary data to S3.
+
+ Data is uploaded in the background using multipart upload as you write.
+ Call close() to complete the upload and ensure all data is persisted.
+
+ This class follows Python's standard I/O interface (io.IOBase).
+
+ Usage:
+ from pytexera import largebinary, LargeBinaryOutputStream
+
+ # Create a new largebinary and write to it
+ large_binary = largebinary()
+ with LargeBinaryOutputStream(large_binary) as out:
+ out.write(b"Hello, World!")
+ out.write(b"More data")
+ # large_binary is now ready to be added to tuples
+
+ Note: Not thread-safe. Do not access from multiple threads concurrently.
+ """
+
+ def __init__(self, large_binary: largebinary):
+ """
+ Initialize a LargeBinaryOutputStream.
+
+ Args:
+ large_binary: The largebinary reference to write to
+
+ Raises:
+ ValueError: If large_binary is None
+ """
+ super().__init__()
+ if large_binary is None:
+ raise ValueError("largebinary cannot be None")
+
+ self._large_binary = large_binary
+ self._bucket_name = large_binary.get_bucket_name()
+ self._object_key = large_binary.get_object_key()
+ self._closed = False
+
+ # Background upload thread state
+ self._queue: queue.Queue = queue.Queue(maxsize=_CHUNK_SIZE)
+ self._upload_exception: Optional[Exception] = None
+ self._upload_complete = threading.Event()
+ self._upload_thread: Optional[threading.Thread] = None
+ self._lock = threading.Lock()
+
+ def write(self, b: Union[bytes, bytearray]) -> int:
+ """
+ Write bytes to the stream.
+
+ Args:
+ b: Bytes to write
+
+ Returns:
+ Number of bytes written
+
+ Raises:
+ ValueError: If stream is closed
+ IOError: If previous upload failed
+ """
+ if self._closed:
+ raise ValueError("I/O operation on closed stream")
+
+ # Check if upload has failed
+ with self._lock:
+ if self._upload_exception is not None:
+ raise IOError(
+ f"Background upload failed: {self._upload_exception}"
+ ) from self._upload_exception
+
+ # Start upload thread on first write
+ if self._upload_thread is None:
+
+ def upload_worker():
+ try:
+ large_binary_manager._ensure_bucket_exists(self._bucket_name)
+ s3 = large_binary_manager._get_s3_client()
+ reader = _QueueReader(self._queue)
+ s3.upload_fileobj(reader, self._bucket_name, self._object_key)
+ except Exception as e:
+ with self._lock:
+ self._upload_exception = e
+ finally:
+ self._upload_complete.set()
+
+ self._upload_thread = threading.Thread(target=upload_worker, daemon=True)
+ self._upload_thread.start()
+
+ # Write data in chunks
+ data = bytes(b)
+ for offset in range(0, len(data), _CHUNK_SIZE):
+ self._queue.put(data[offset : offset + _CHUNK_SIZE], block=True)
+
+ return len(data)
+
+ def writable(self) -> bool:
+ """Return True if the stream can be written to."""
+ return not self._closed
+
+ def seekable(self) -> bool:
+ """Return False - this stream does not support seeking."""
+ return False
+
+ @property
+ def closed(self) -> bool:
+ """Return True if the stream is closed."""
+ return self._closed
+
+ def flush(self) -> None:
+ """
+ Flush the write buffer.
+
+ Note: This doesn't guarantee data is uploaded to S3 yet.
+ Call close() to ensure upload completion.
+ """
+ # No-op: data is already being consumed by the upload thread
+ pass
+
+ def close(self) -> None:
+ """
+ Close the stream and complete the S3 upload.
+ Blocks until upload is complete. Raises IOError if upload failed.
+
+ Raises:
+ IOError: If upload failed
+ """
+ if self._closed:
+ return
+
+ self._closed = True
+
+ # Signal EOF to upload thread and wait for completion
+ if self._upload_thread is not None:
+ self._queue.put(None, block=True) # EOF marker
+ self._upload_thread.join()
+ self._upload_complete.wait()
+
+ # Check for errors and cleanup if needed
+ with self._lock:
+ exception = self._upload_exception
+
+ if exception is not None:
+ self._cleanup_failed_upload()
+ raise IOError(f"Failed to complete upload: {exception}") from exception
+
+ def _cleanup_failed_upload(self):
+ """Clean up a failed upload by deleting the S3 object."""
+ try:
+ s3 = large_binary_manager._get_s3_client()
+ s3.delete_object(Bucket=self._bucket_name, Key=self._object_key)
+ except Exception:
+ # Ignore cleanup errors - we're already handling an upload failure
+ pass
+
+ def __enter__(self):
+ """Context manager entry."""
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ """Context manager exit - automatically cleanup."""
+ self.close()
+ return False
diff --git a/amber/src/main/python/pytexera/storage/test_large_binary_input_stream.py b/amber/src/main/python/pytexera/storage/test_large_binary_input_stream.py
new file mode 100644
index 00000000000..85bdbd13fa1
--- /dev/null
+++ b/amber/src/main/python/pytexera/storage/test_large_binary_input_stream.py
@@ -0,0 +1,222 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+from unittest.mock import patch, MagicMock
+from io import BytesIO
+from core.models.type.large_binary import largebinary
+from pytexera.storage.large_binary_input_stream import LargeBinaryInputStream
+from pytexera.storage import large_binary_manager
+
+
+class TestLargeBinaryInputStream:
+ @pytest.fixture
+ def large_binary(self):
+ """Create a test largebinary."""
+ return largebinary("s3://test-bucket/path/to/object")
+
+ @pytest.fixture
+ def mock_s3_response(self):
+ """Create a mock S3 response with a BytesIO body."""
+ return {"Body": BytesIO(b"test data content")}
+
+ def test_init_with_valid_large_binary(self, large_binary):
+ """Test initialization with a valid largebinary."""
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ assert stream._large_binary == large_binary
+ assert stream._underlying is None
+ assert not stream._closed
+ finally:
+ stream.close()
+
+ def test_init_with_none_raises_error(self):
+ """Test that initializing with None raises ValueError."""
+ with pytest.raises(ValueError, match="largebinary cannot be None"):
+ LargeBinaryInputStream(None)
+
+ def test_lazy_init_downloads_from_s3(self, large_binary, mock_s3_response):
+ """Test that _lazy_init downloads from S3 on first read."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = mock_s3_response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ assert stream._underlying is None # Not initialized yet
+
+ # Trigger lazy init by reading
+ data = stream.read()
+ assert data == b"test data content"
+ assert stream._underlying is not None
+
+ # Verify S3 was called correctly
+ mock_s3_client.get_object.assert_called_once_with(
+ Bucket="test-bucket", Key="path/to/object"
+ )
+ finally:
+ stream.close()
+
+ def test_read_all(self, large_binary, mock_s3_response):
+ """Test reading all data."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = mock_s3_response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ data = stream.read()
+ assert data == b"test data content"
+ finally:
+ stream.close()
+
+ def test_read_partial(self, large_binary, mock_s3_response):
+ """Test reading partial data."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = mock_s3_response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ data = stream.read(4)
+ assert data == b"test"
+ finally:
+ stream.close()
+
+ def test_readline(self, large_binary):
+ """Test reading a line."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ response = {"Body": BytesIO(b"line1\nline2\nline3")}
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ line = stream.readline()
+ assert line == b"line1\n"
+ finally:
+ stream.close()
+
+ def test_readlines(self, large_binary):
+ """Test reading all lines."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ response = {"Body": BytesIO(b"line1\nline2\nline3")}
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ lines = stream.readlines()
+ assert lines == [b"line1\n", b"line2\n", b"line3"]
+ finally:
+ stream.close()
+
+ def test_readable(self, large_binary):
+ """Test readable() method."""
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ assert stream.readable() is True
+
+ stream.close()
+ assert stream.readable() is False
+ finally:
+ if not stream._closed:
+ stream.close()
+
+ def test_seekable(self, large_binary):
+ """Test seekable() method (should always return False)."""
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ assert stream.seekable() is False
+ finally:
+ stream.close()
+
+ def test_closed_property(self, large_binary):
+ """Test closed property."""
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ assert stream.closed is False
+
+ stream.close()
+ assert stream.closed is True
+ finally:
+ if not stream._closed:
+ stream.close()
+
+ def test_close(self, large_binary, mock_s3_response):
+ """Test closing the stream."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = mock_s3_response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ stream = LargeBinaryInputStream(large_binary)
+ stream.read(1) # Trigger lazy init
+ assert stream._underlying is not None
+
+ stream.close()
+ assert stream._closed is True
+ assert stream._underlying.closed
+
+ def test_context_manager(self, large_binary, mock_s3_response):
+ """Test using as context manager."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = mock_s3_response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ with LargeBinaryInputStream(large_binary) as stream:
+ data = stream.read()
+ assert data == b"test data content"
+ assert not stream._closed
+
+ # Stream should be closed after context exit
+ assert stream._closed
+
+ def test_iteration(self, large_binary):
+ """Test iteration over lines."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ response = {"Body": BytesIO(b"line1\nline2\nline3")}
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ stream = LargeBinaryInputStream(large_binary)
+ try:
+ lines = list(stream)
+ assert lines == [b"line1\n", b"line2\n", b"line3"]
+ finally:
+ stream.close()
+
+ def test_read_after_close_raises_error(self, large_binary, mock_s3_response):
+ """Test that reading after close raises ValueError."""
+ with patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client:
+ mock_s3_client = MagicMock()
+ mock_s3_client.get_object.return_value = mock_s3_response
+ mock_get_s3_client.return_value = mock_s3_client
+
+ stream = LargeBinaryInputStream(large_binary)
+ stream.close()
+
+ with pytest.raises(ValueError, match="I/O operation on closed stream"):
+ stream.read()
+ # Stream is already closed, no need to close again
diff --git a/amber/src/main/python/pytexera/storage/test_large_binary_manager.py b/amber/src/main/python/pytexera/storage/test_large_binary_manager.py
new file mode 100644
index 00000000000..64c7080e520
--- /dev/null
+++ b/amber/src/main/python/pytexera/storage/test_large_binary_manager.py
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+from unittest.mock import patch, MagicMock
+from pytexera.storage import large_binary_manager
+from core.storage.storage_config import StorageConfig
+
+
+class TestLargeBinaryManager:
+ @pytest.fixture(autouse=True)
+ def setup_storage_config(self):
+ """Initialize StorageConfig for tests."""
+ if not StorageConfig._initialized:
+ StorageConfig.initialize(
+ catalog_type="postgres",
+ postgres_uri_without_scheme="localhost:5432/test",
+ postgres_username="test",
+ postgres_password="test",
+ rest_catalog_uri="http://localhost:8181/catalog/",
+ rest_catalog_warehouse_name="texera",
+ table_result_namespace="test",
+ directory_path="/tmp/test",
+ commit_batch_size=1000,
+ s3_endpoint="http://localhost:9000",
+ s3_region="us-east-1",
+ s3_auth_username="minioadmin",
+ s3_auth_password="minioadmin",
+ )
+
+ def test_get_s3_client_initializes_once(self):
+ """Test that S3 client is initialized and cached."""
+ # Reset the client
+ large_binary_manager._s3_client = None
+
+ with patch("boto3.client") as mock_boto3_client:
+ mock_client = MagicMock()
+ mock_boto3_client.return_value = mock_client
+
+ # First call should create client
+ client1 = large_binary_manager._get_s3_client()
+ assert client1 == mock_client
+ assert mock_boto3_client.call_count == 1
+
+ # Second call should return cached client
+ client2 = large_binary_manager._get_s3_client()
+ assert client2 == mock_client
+ assert mock_boto3_client.call_count == 1 # Still 1, not 2
+
+ def test_get_s3_client_without_boto3_raises_error(self):
+ """Test that missing boto3 raises RuntimeError."""
+ large_binary_manager._s3_client = None
+
+ import sys
+
+ # Temporarily remove boto3 from sys.modules to simulate it not being installed
+ boto3_backup = sys.modules.pop("boto3", None)
+ try:
+ # Mock the import to raise ImportError
+ original_import = __import__
+
+ def mock_import(name, *args, **kwargs):
+ if name == "boto3":
+ raise ImportError("No module named boto3")
+ return original_import(name, *args, **kwargs)
+
+ with patch("builtins.__import__", side_effect=mock_import):
+ with pytest.raises(RuntimeError, match="boto3 required"):
+ large_binary_manager._get_s3_client()
+ finally:
+ # Restore boto3 if it was there
+ if boto3_backup is not None:
+ sys.modules["boto3"] = boto3_backup
+
+ def test_ensure_bucket_exists_when_bucket_exists(self):
+ """Test that existing bucket doesn't trigger creation."""
+ large_binary_manager._s3_client = None
+
+ with patch("boto3.client") as mock_boto3_client:
+ mock_client = MagicMock()
+ mock_boto3_client.return_value = mock_client
+ # head_bucket doesn't raise exception (bucket exists)
+ mock_client.head_bucket.return_value = None
+ mock_client.exceptions.NoSuchBucket = type("NoSuchBucket", (Exception,), {})
+
+ large_binary_manager._ensure_bucket_exists("test-bucket")
+ mock_client.head_bucket.assert_called_once_with(Bucket="test-bucket")
+ mock_client.create_bucket.assert_not_called()
+
+ def test_ensure_bucket_exists_creates_bucket_when_missing(self):
+ """Test that missing bucket triggers creation."""
+ large_binary_manager._s3_client = None
+
+ with patch("boto3.client") as mock_boto3_client:
+ mock_client = MagicMock()
+ mock_boto3_client.return_value = mock_client
+ # head_bucket raises NoSuchBucket exception
+ no_such_bucket = type("NoSuchBucket", (Exception,), {})
+ mock_client.exceptions.NoSuchBucket = no_such_bucket
+ mock_client.head_bucket.side_effect = no_such_bucket()
+
+ large_binary_manager._ensure_bucket_exists("test-bucket")
+ mock_client.head_bucket.assert_called_once_with(Bucket="test-bucket")
+ mock_client.create_bucket.assert_called_once_with(Bucket="test-bucket")
+
+ def test_create_generates_unique_uri(self):
+ """Test that create() generates a unique S3 URI."""
+ large_binary_manager._s3_client = None
+
+ with patch("boto3.client") as mock_boto3_client:
+ mock_client = MagicMock()
+ mock_boto3_client.return_value = mock_client
+ mock_client.head_bucket.return_value = None
+ mock_client.exceptions.NoSuchBucket = type("NoSuchBucket", (Exception,), {})
+
+ uri = large_binary_manager.create()
+
+ # Check URI format
+ assert uri.startswith("s3://")
+ assert uri.startswith(f"s3://{large_binary_manager.DEFAULT_BUCKET}/")
+ assert "objects/" in uri
+
+ # Verify bucket was checked/created
+ mock_client.head_bucket.assert_called_once_with(
+ Bucket=large_binary_manager.DEFAULT_BUCKET
+ )
+
+ def test_create_uses_default_bucket(self):
+ """Test that create() uses the default bucket."""
+ large_binary_manager._s3_client = None
+
+ with patch("boto3.client") as mock_boto3_client:
+ mock_client = MagicMock()
+ mock_boto3_client.return_value = mock_client
+ mock_client.head_bucket.return_value = None
+ mock_client.exceptions.NoSuchBucket = type("NoSuchBucket", (Exception,), {})
+
+ uri = large_binary_manager.create()
+ assert large_binary_manager.DEFAULT_BUCKET in uri
diff --git a/amber/src/main/python/pytexera/storage/test_large_binary_output_stream.py b/amber/src/main/python/pytexera/storage/test_large_binary_output_stream.py
new file mode 100644
index 00000000000..7ebcc9b4cfd
--- /dev/null
+++ b/amber/src/main/python/pytexera/storage/test_large_binary_output_stream.py
@@ -0,0 +1,238 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+import time
+from unittest.mock import patch, MagicMock
+from core.models.type.large_binary import largebinary
+from pytexera.storage.large_binary_output_stream import LargeBinaryOutputStream
+from pytexera.storage import large_binary_manager
+
+
+class TestLargeBinaryOutputStream:
+ @pytest.fixture
+ def large_binary(self):
+ """Create a test largebinary."""
+ return largebinary("s3://test-bucket/path/to/object")
+
+ def test_init_with_valid_large_binary(self, large_binary):
+ """Test initialization with a valid largebinary."""
+ stream = LargeBinaryOutputStream(large_binary)
+ assert stream._large_binary == large_binary
+ assert stream._bucket_name == "test-bucket"
+ assert stream._object_key == "path/to/object"
+ assert not stream._closed
+ assert stream._upload_thread is None
+
+ def test_init_with_none_raises_error(self):
+ """Test that initializing with None raises ValueError."""
+ with pytest.raises(ValueError, match="largebinary cannot be None"):
+ LargeBinaryOutputStream(None)
+
+ def test_write_starts_upload_thread(self, large_binary):
+ """Test that write() starts the upload thread."""
+ with (
+ patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client,
+ patch.object(
+ large_binary_manager, "_ensure_bucket_exists"
+ ) as mock_ensure_bucket,
+ ):
+ mock_s3 = MagicMock()
+ mock_get_s3_client.return_value = mock_s3
+ mock_ensure_bucket.return_value = None
+
+ stream = LargeBinaryOutputStream(large_binary)
+ assert stream._upload_thread is None
+
+ stream.write(b"test data")
+ assert stream._upload_thread is not None
+ # Thread may have already completed, so just check it was created
+ assert stream._upload_thread is not None
+
+ # Wait for thread to finish
+ stream.close()
+
+ def test_write_data(self, large_binary):
+ """Test writing data to the stream."""
+ with (
+ patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client,
+ patch.object(
+ large_binary_manager, "_ensure_bucket_exists"
+ ) as mock_ensure_bucket,
+ ):
+ mock_s3 = MagicMock()
+ mock_get_s3_client.return_value = mock_s3
+ mock_ensure_bucket.return_value = None
+
+ stream = LargeBinaryOutputStream(large_binary)
+ bytes_written = stream.write(b"test data")
+ assert bytes_written == len(b"test data")
+
+ stream.close()
+
+ def test_write_multiple_chunks(self, large_binary):
+ """Test writing multiple chunks of data."""
+ with (
+ patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client,
+ patch.object(
+ large_binary_manager, "_ensure_bucket_exists"
+ ) as mock_ensure_bucket,
+ ):
+ mock_s3 = MagicMock()
+ mock_get_s3_client.return_value = mock_s3
+ mock_ensure_bucket.return_value = None
+
+ stream = LargeBinaryOutputStream(large_binary)
+ stream.write(b"chunk1")
+ stream.write(b"chunk2")
+ stream.write(b"chunk3")
+
+ stream.close()
+
+ def test_writable(self, large_binary):
+ """Test writable() method."""
+ stream = LargeBinaryOutputStream(large_binary)
+ assert stream.writable() is True
+
+ stream.close()
+ assert stream.writable() is False
+
+ def test_seekable(self, large_binary):
+ """Test seekable() method (should always return False)."""
+ stream = LargeBinaryOutputStream(large_binary)
+ assert stream.seekable() is False
+
+ def test_closed_property(self, large_binary):
+ """Test closed property."""
+ stream = LargeBinaryOutputStream(large_binary)
+ assert stream.closed is False
+
+ stream.close()
+ assert stream.closed is True
+
+ def test_flush(self, large_binary):
+ """Test flush() method (should be a no-op)."""
+ stream = LargeBinaryOutputStream(large_binary)
+ # Should not raise any exception
+ stream.flush()
+
+ def test_close_completes_upload(self, large_binary):
+ """Test that close() completes the upload."""
+ with (
+ patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client,
+ patch.object(
+ large_binary_manager, "_ensure_bucket_exists"
+ ) as mock_ensure_bucket,
+ ):
+ mock_s3 = MagicMock()
+ mock_get_s3_client.return_value = mock_s3
+ mock_ensure_bucket.return_value = None
+
+ stream = LargeBinaryOutputStream(large_binary)
+ stream.write(b"test data")
+
+ # Close should wait for upload to complete
+ stream.close()
+
+ # Verify upload_fileobj was called
+ assert mock_s3.upload_fileobj.called
+
+ def test_context_manager(self, large_binary):
+ """Test using as context manager."""
+ with (
+ patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client,
+ patch.object(
+ large_binary_manager, "_ensure_bucket_exists"
+ ) as mock_ensure_bucket,
+ ):
+ mock_s3 = MagicMock()
+ mock_get_s3_client.return_value = mock_s3
+ mock_ensure_bucket.return_value = None
+
+ with LargeBinaryOutputStream(large_binary) as stream:
+ stream.write(b"test data")
+ assert not stream._closed
+
+ # Stream should be closed after context exit
+ assert stream._closed
+
+ def test_write_after_close_raises_error(self, large_binary):
+ """Test that writing after close raises ValueError."""
+ stream = LargeBinaryOutputStream(large_binary)
+ stream.close()
+
+ with pytest.raises(ValueError, match="I/O operation on closed stream"):
+ stream.write(b"data")
+
+ def test_close_handles_upload_error(self, large_binary):
+ """Test that close() raises IOError if upload fails."""
+ with (
+ patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client,
+ patch.object(
+ large_binary_manager, "_ensure_bucket_exists"
+ ) as mock_ensure_bucket,
+ ):
+ mock_s3 = MagicMock()
+ mock_get_s3_client.return_value = mock_s3
+ mock_ensure_bucket.return_value = None
+ mock_s3.upload_fileobj.side_effect = Exception("Upload failed")
+
+ stream = LargeBinaryOutputStream(large_binary)
+ stream.write(b"test data")
+
+ with pytest.raises(IOError, match="Failed to complete upload"):
+ stream.close()
+
+ def test_write_after_upload_error_raises_error(self, large_binary):
+ """Test that writing after upload error raises IOError."""
+ with (
+ patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client,
+ patch.object(
+ large_binary_manager, "_ensure_bucket_exists"
+ ) as mock_ensure_bucket,
+ ):
+ mock_s3 = MagicMock()
+ mock_get_s3_client.return_value = mock_s3
+ mock_ensure_bucket.return_value = None
+ mock_s3.upload_fileobj.side_effect = Exception("Upload failed")
+
+ stream = LargeBinaryOutputStream(large_binary)
+ stream.write(b"test data")
+
+ # Wait a bit for the error to be set
+ time.sleep(0.1)
+
+ with pytest.raises(IOError, match="Background upload failed"):
+ stream.write(b"more data")
+
+ def test_multiple_close_calls(self, large_binary):
+ """Test that multiple close() calls are safe."""
+ with (
+ patch.object(large_binary_manager, "_get_s3_client") as mock_get_s3_client,
+ patch.object(
+ large_binary_manager, "_ensure_bucket_exists"
+ ) as mock_ensure_bucket,
+ ):
+ mock_s3 = MagicMock()
+ mock_get_s3_client.return_value = mock_s3
+ mock_ensure_bucket.return_value = None
+
+ stream = LargeBinaryOutputStream(large_binary)
+ stream.write(b"test data")
+ stream.close()
+ # Second close should not raise error
+ stream.close()
diff --git a/amber/src/main/python/pytexera/udf/__init__.py b/amber/src/main/python/pytexera/udf/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/amber/src/main/python/pytexera/udf/examples/__init__.py b/amber/src/main/python/pytexera/udf/examples/__init__.py
new file mode 100644
index 00000000000..34b7183567d
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/__init__.py
@@ -0,0 +1,23 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .echo_operator import EchoOperator
+from .echo_table_operator import EchoTableOperator
+from .join_operator import JoinOperator
+from .count_batch_operator import CountBatchOperator
+
+__all__ = ["EchoOperator", "EchoTableOperator", "JoinOperator", "CountBatchOperator"]
diff --git a/amber/src/main/python/pytexera/udf/examples/count_batch_operator.py b/amber/src/main/python/pytexera/udf/examples/count_batch_operator.py
new file mode 100644
index 00000000000..d4aed1746e7
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/count_batch_operator.py
@@ -0,0 +1,31 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from pytexera import *
+
+
+class CountBatchOperator(UDFBatchOperator):
+ BATCH_SIZE = 10
+
+ def __init__(self):
+ super().__init__()
+ self.count = 0
+
+ @overrides
+ def process_batch(self, batch: Batch, port: int) -> Iterator[Optional[BatchLike]]:
+ self.count += 1
+ yield batch
diff --git a/amber/src/main/python/pytexera/udf/examples/echo_operator.py b/amber/src/main/python/pytexera/udf/examples/echo_operator.py
new file mode 100644
index 00000000000..2b93f43bc45
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/echo_operator.py
@@ -0,0 +1,28 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from pytexera import *
+
+
+class EchoOperator(UDFOperatorV2):
+ @overrides
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ yield tuple_
+
+ @overrides
+ def on_finish(self, port: int) -> Iterator[Optional[TupleLike]]:
+ yield
diff --git a/amber/src/main/python/pytexera/udf/examples/echo_table_operator.py b/amber/src/main/python/pytexera/udf/examples/echo_table_operator.py
new file mode 100644
index 00000000000..3402e4434b7
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/echo_table_operator.py
@@ -0,0 +1,24 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from pytexera import *
+
+
+class EchoTableOperator(UDFTableOperator):
+ @overrides
+ def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:
+ yield table
diff --git a/amber/src/main/python/pytexera/udf/examples/generator_operator_binary.py b/amber/src/main/python/pytexera/udf/examples/generator_operator_binary.py
new file mode 100644
index 00000000000..bdb13387e75
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/generator_operator_binary.py
@@ -0,0 +1,28 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from pytexera import *
+
+
+class GeneratorOperatorBinary(UDFSourceOperator):
+ """
+ A simple generator operator that produces a single tuple with a binary attribute.
+ """
+
+ @overrides
+ def produce(self) -> Iterator[Union[TupleLike, TableLike, None]]:
+ yield {"test": [1, 2, 3]}
diff --git a/amber/src/main/python/pytexera/udf/examples/generator_operator_integer.py b/amber/src/main/python/pytexera/udf/examples/generator_operator_integer.py
new file mode 100644
index 00000000000..6be50846358
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/generator_operator_integer.py
@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from pytexera import *
+
+
+class GeneratorOperatorInteger(UDFSourceOperator):
+ """
+ A simple generator operator that produces tuples with an integer attribute.
+ """
+
+ @overrides
+ def produce(self) -> Iterator[Union[TupleLike, TableLike, None]]:
+ for i in [1, 2, 3]:
+ yield {"test": i}
diff --git a/amber/src/main/python/pytexera/udf/examples/join_operator.py b/amber/src/main/python/pytexera/udf/examples/join_operator.py
new file mode 100644
index 00000000000..340305fde2a
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/join_operator.py
@@ -0,0 +1,38 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from collections import defaultdict
+
+from pytexera import *
+
+
+# needs the dual-input-ports PythonUDF
+class JoinOperator(UDFOperatorV2):
+ @overrides
+ def open(self) -> None:
+ self.left_dict = defaultdict(list)
+
+ @overrides
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ if port == 0:
+ # building the hashmap
+ self.left_dict[tuple_["key"]].append(tuple_)
+ else:
+ # probing the hashmap
+ for left_tuple in self.left_dict.get(tuple_["key"], []):
+ # join and output
+ yield left_tuple + tuple_
diff --git a/amber/src/main/python/pytexera/udf/examples/rudf/r_table_operator.py b/amber/src/main/python/pytexera/udf/examples/rudf/r_table_operator.py
new file mode 100644
index 00000000000..e6ce072dd12
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/rudf/r_table_operator.py
@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Note: make sure R path is initialized in udf.conf and make sure that
+# the following packages (in R) are installed: arrow
+# --- Source Operator Examples ---
+r_table_source_simple_table = """
+function() {
+ df <- data.frame(
+ Training = c("Strength", "Stamina", "Other"),
+ Pulse = c(100L, 150L, 120L),
+ Duration = c(60, 30, 45)
+ )
+ return (df)
+}
+"""
+
+r_table_source_like_objects = """
+function() {
+ # Works with R lists, R vectors, R matrices,
+ # or anything that can converted to a data.frame
+
+ # Matrix
+ mdat <- matrix(c(1,2,3, 11,12,13), nrow = 2, ncol = 3, byrow = TRUE,
+ dimnames = list(c("row1", "row2"),
+ c("col1", "col2", "col3")))
+
+ # List
+ lst <- list(col1 = c(1,2), col2 = c(2,12), col3 = c(3,13))
+
+ # Vectors
+ col1_vec <- c(1,11)
+ col2_vec <- c(2,12)
+ col3_vec <- c(3,13)
+ df_from_vec <- data.frame(col1_vec, col2_vec, col3_vec)
+
+ return (mdat)
+}
+"""
+
+# --- UDF Operator Examples ---
+r_table_udf_echo_table = """
+function(table, port) {
+ return (table)
+}
+"""
+
+r_table_udf_add_row = """
+function(table, port) {
+ # Assuming table is:
+ # data.frame(
+ # Training = c("Strength", "Stamina", "Other"),
+ # Pulse = c(100L, 150L, 120L),
+ # Duration = c(60, 30, 45)
+ # )
+ new_row <- list(Training = "NEW", Pulse = 999L, Duration = 999)
+ new_df <- rbind(table, new_row)
+ return (new_df)
+}
+"""
+
+r_table_udf_extract_row = """
+function(table, port) {
+ # Assuming table is:
+ # data.frame(
+ # Training = c("Strength", "Stamina", "Other"),
+ # Pulse = c(100L, 150L, 120L),
+ # Duration = c(60, 30, 45)
+ # )
+ tuple <- table[1,]
+ return (tuple)
+}
+"""
diff --git a/amber/src/main/python/pytexera/udf/examples/rudf/r_tuple_operator.py b/amber/src/main/python/pytexera/udf/examples/rudf/r_tuple_operator.py
new file mode 100644
index 00000000000..4aeb8a010ce
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/rudf/r_tuple_operator.py
@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Note: make sure R path is initialized in udf.conf and make sure that
+# the following packages (in R) are installed: coro, arrow
+# --- Source Operator Examples ---
+r_tuple_source_zero_tuple = """
+library(coro)
+coro::generator(function() {
+ yield (list())
+ # yield (NULL) works too
+ })
+"""
+
+r_tuple_source_one_tuple = """
+library(coro)
+coro::generator(function() {
+ yield (list(
+ attr1 = 1L, # R integer
+ attr2 = "A", # R string
+ attr3 = TRUE, # R logical (boolean)
+ ))
+ })
+"""
+
+r_tuple_source_multiple_tuples = """
+library(coro)
+coro::generator(function() {
+ for (i in 1:5) {
+ yield (list(
+ attr1 = 1L, # R integer
+ attr2 = "A", # R string
+ attr3 = TRUE, # R logical (boolean)
+ ))
+ })
+"""
+
+# --- UDF Operator ---
+r_tuple_udf_zero_tuple = """
+library(coro)
+coro::generator(function(tuple, port) {
+ yield (list())
+ })
+"""
+
+r_tuple_udf_echo = """
+library(coro)
+coro::generator(function(tuple, port) {
+ yield (tuple)
+ })
+"""
+
+r_tuple_udf_echo_multiple_tuples = """
+library(coro)
+coro::generator(function(tuple, port) {
+ for (i in 1:5) {
+ yield (tuple)
+ })
+"""
diff --git a/amber/src/main/python/pytexera/udf/examples/test_count_batch_operator.py b/amber/src/main/python/pytexera/udf/examples/test_count_batch_operator.py
new file mode 100644
index 00000000000..9ab084a4050
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/test_count_batch_operator.py
@@ -0,0 +1,122 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import inspect
+import pytest
+from collections import deque
+
+from pytexera import *
+from .count_batch_operator import CountBatchOperator
+
+
+class TestCountBatchOperator:
+ @pytest.fixture
+ def count_batch_operator(self):
+ return CountBatchOperator()
+
+ def test_count_batch_operator(self, count_batch_operator):
+ count_batch_operator.open()
+ for i in range(27):
+ deque(
+ count_batch_operator.process_tuple(
+ Tuple({"test-1": "hello", "test-2": 10}), 0
+ )
+ )
+ deque(count_batch_operator.on_finish(0))
+ batch_counter = count_batch_operator.count
+ assert batch_counter == 3
+ count_batch_operator.close()
+
+ def test_count_batch_operator_simple(self, count_batch_operator):
+ count_batch_operator.open()
+ for i in range(20):
+ deque(
+ count_batch_operator.process_tuple(
+ Tuple({"test-1": "hello", "test-2": 10}), 0
+ )
+ )
+ deque(count_batch_operator.on_finish(0))
+ batch_counter = count_batch_operator.count
+ assert batch_counter == 2
+ count_batch_operator.close()
+
+ def test_count_batch_operator_medium(self, count_batch_operator):
+ count_batch_operator.open()
+ for i in range(27):
+ deque(
+ count_batch_operator.process_tuple(
+ Tuple({"test-1": "hello", "test-2": 10}), 0
+ )
+ )
+ deque(count_batch_operator.on_finish(0))
+ batch_counter = count_batch_operator.count
+ assert batch_counter == 3
+ count_batch_operator.close()
+
+ def test_count_batch_operator_hard(self, count_batch_operator):
+ count_batch_operator.open()
+ count_batch_operator.BATCH_SIZE = 10
+ for i in range(27):
+ deque(
+ count_batch_operator.process_tuple(
+ Tuple({"test-1": "hello", "test-2": 10}), 0
+ )
+ )
+ count_batch_operator.BATCH_SIZE = 5
+ for i in range(27):
+ deque(
+ count_batch_operator.process_tuple(
+ Tuple({"test-1": "hello", "test-2": 10}), 0
+ )
+ )
+ deque(count_batch_operator.on_finish(0))
+ batch_counter = count_batch_operator.count
+ assert batch_counter == 9
+ count_batch_operator.close()
+
+ def test_edge_case_string(self):
+ with pytest.raises(ValueError) as exc_info:
+ operator_string = str(inspect.getsource(CountBatchOperator))
+ operator_string = operator_string.replace(
+ "BATCH_SIZE = 10", 'BATCH_SIZE = "test"'
+ )
+ operator_string += "operator = CountBatchOperator()"
+ exec(operator_string)
+ assert (
+ exc_info.value.args[0]
+ == "BATCH_SIZE cannot be " + str(type("test")) + "."
+ )
+
+ def test_edge_case_non_positive(self, count_batch_operator):
+ with pytest.raises(ValueError) as exc_info:
+ operator_string = str(inspect.getsource(CountBatchOperator))
+ operator_string = operator_string.replace(
+ "BATCH_SIZE = 10", "BATCH_SIZE = -20"
+ )
+ operator_string += "operator = CountBatchOperator()"
+ exec(operator_string)
+ assert exc_info.value.args[0] == "BATCH_SIZE should be positive."
+
+ def test_edge_case_none(self, count_batch_operator):
+ with pytest.raises(ValueError) as exc_info:
+ operator_string = str(inspect.getsource(CountBatchOperator))
+ operator_string = operator_string.replace(
+ "BATCH_SIZE = 10", "BATCH_SIZE = None"
+ )
+ operator_string += "operator = CountBatchOperator()"
+ exec(operator_string)
+ assert exc_info.value.args[0] == "BATCH_SIZE cannot be None."
diff --git a/amber/src/main/python/pytexera/udf/examples/test_echo_operator.py b/amber/src/main/python/pytexera/udf/examples/test_echo_operator.py
new file mode 100644
index 00000000000..a70b8a0e7c1
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/test_echo_operator.py
@@ -0,0 +1,38 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from pytexera import Tuple
+from .echo_operator import EchoOperator
+
+
+class TestEchoOperator:
+ @pytest.fixture
+ def echo_operator(self):
+ return EchoOperator()
+
+ def test_echo_operator(self, echo_operator):
+ echo_operator.open()
+ tuple_ = Tuple({"test-1": "hello", "test-2": 10})
+
+ outputs = echo_operator.process_tuple(tuple_, 0)
+ output_tuple = next(outputs)
+
+ assert output_tuple == tuple_
+ with pytest.raises(StopIteration):
+ next(outputs)
diff --git a/amber/src/main/python/pytexera/udf/examples/test_echo_table_operator.py b/amber/src/main/python/pytexera/udf/examples/test_echo_table_operator.py
new file mode 100644
index 00000000000..d619f53a66b
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/test_echo_table_operator.py
@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+from collections import deque
+
+from core.models.table import all_output_to_tuple
+from pytexera import Tuple
+from .echo_table_operator import EchoTableOperator
+
+
+class TestEchoTableOperator:
+ @pytest.fixture
+ def echo_table_operator(self):
+ return EchoTableOperator()
+
+ def test_echo_table_operator(self, echo_table_operator):
+ echo_table_operator.open()
+ tuple_ = Tuple({"test-1": "hello", "test-2": 10})
+ print(tuple_)
+ deque(echo_table_operator.process_tuple(tuple_, 0))
+ outputs = echo_table_operator.on_finish(0)
+ output_tuple = next(all_output_to_tuple(next(outputs)))
+ assert output_tuple == tuple_
+ with pytest.raises(StopIteration):
+ next(outputs)
+ echo_table_operator.close()
diff --git a/amber/src/main/python/pytexera/udf/examples/test_generator_operator_binary.py b/amber/src/main/python/pytexera/udf/examples/test_generator_operator_binary.py
new file mode 100644
index 00000000000..4c7e5d8b407
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/test_generator_operator_binary.py
@@ -0,0 +1,34 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from pytexera import Tuple
+from .generator_operator_binary import GeneratorOperatorBinary
+
+
+class TestEchoOperator:
+ @pytest.fixture
+ def generator_operator_binary(self):
+ return GeneratorOperatorBinary()
+
+ def test_generator_operator_binary(self, generator_operator_binary):
+ generator_operator_binary.open()
+ outputs = generator_operator_binary.produce()
+ output_tuple = Tuple(next(outputs))
+ assert output_tuple == Tuple({"test": [1, 2, 3]})
+ generator_operator_binary.close()
diff --git a/amber/src/main/python/pytexera/udf/examples/test_generator_operator_integer.py b/amber/src/main/python/pytexera/udf/examples/test_generator_operator_integer.py
new file mode 100644
index 00000000000..3ab19064623
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/examples/test_generator_operator_integer.py
@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from pytexera import Tuple
+from .generator_operator_integer import GeneratorOperatorInteger
+
+
+class TestEchoOperator:
+ @pytest.fixture
+ def generator_operator_integer(self):
+ return GeneratorOperatorInteger()
+
+ def test_generator_operator_integer(self, generator_operator_integer):
+ generator_operator_integer.open()
+ outputs = generator_operator_integer.produce()
+ for i in [1, 2, 3]:
+ output_tuple = Tuple(next(outputs))
+ assert output_tuple == Tuple({"test": i})
+ generator_operator_integer.close()
diff --git a/amber/src/main/python/pytexera/udf/udf_operator.py b/amber/src/main/python/pytexera/udf/udf_operator.py
new file mode 100644
index 00000000000..003225c75c3
--- /dev/null
+++ b/amber/src/main/python/pytexera/udf/udf_operator.py
@@ -0,0 +1,156 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import abstractmethod
+from typing import Iterator, Optional, Union
+
+from pyamber import *
+
+
+class UDFOperatorV2(TupleOperatorV2):
+ """
+ Base class for tuple-oriented user-defined operators. A concrete implementation must
+ be provided upon using.
+ """
+
+ def open(self) -> None:
+ """
+ Open a context of the operator. Usually can be used for loading/initiating some
+ resources, such as a file, a model, or an API client.
+ """
+ pass
+
+ @abstractmethod
+ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ """
+ Process an input Tuple from the given link.
+
+ :param tuple_: Tuple, a Tuple from an input port to be processed.
+ :param port: int, input port index of the current Tuple.
+ :return: Iterator[Optional[TupleLike]], producing one TupleLike object at a
+ time, or None.
+
+ See .examples/ for example operators.
+ """
+ yield
+
+ def on_finish(self, port: int) -> Iterator[Optional[TupleLike]]:
+ """
+ Callback when one input port is exhausted.
+
+ :param port: int, input port index of the current exhausted port.
+ :return: Iterator[Optional[TupleLike]], producing one TupleLike object at a
+ time, or None.
+ """
+ yield
+
+ def close(self) -> None:
+ """
+ Close the context of the operator.
+ """
+ pass
+
+
+class UDFSourceOperator(SourceOperator):
+ def open(self) -> None:
+ """
+ Open a context of the operator. Usually can be used for loading/initiating some
+ resources, such as a file, a model, or an API client.
+ """
+ pass
+
+ @abstractmethod
+ def produce(self) -> Iterator[Optional[Union[TupleLike, TableLike]]]:
+ """
+ Produce Tuples or Tables. Used by the source operator only.
+
+ :return: Iterator[Union[TupleLike, TableLike, None]], producing
+ one TupleLike object, one TableLike object, or None, at a time.
+ """
+ yield
+
+ def close(self) -> None:
+ """
+ Close the context of the operator.
+ """
+ pass
+
+
+class UDFTableOperator(TableOperator):
+ """
+ Base class for table-oriented user-defined operators. A concrete implementation must
+ be provided upon using.
+ """
+
+ def open(self) -> None:
+ """
+ Open a context of the operator. Usually can be used for loading/initiating some
+ resources, such as a file, a model, or an API client.
+ """
+ pass
+
+ @abstractmethod
+ def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:
+ """
+ Process an input Table from the given link. The Table is represented as
+ pandas.DataFrame.
+
+ :param table: Table, a table to be processed.
+ :param port: int, input index of the current Table.
+ :return: Iterator[Optional[TableLike]], producing one TableLike object at a
+ time, or None.
+ """
+ yield
+
+ def close(self) -> None:
+ """
+ Close the context of the operator.
+ """
+ pass
+
+
+class UDFBatchOperator(BatchOperator):
+ """
+ Base class for batch-oriented user-defined operators. A concrete implementation must
+ be provided upon using.
+ """
+
+ def open(self) -> None:
+ """
+ Open a context of the operator. Usually can be used for loading/initiating some
+ resources, such as a file, a model, or an API client.
+ """
+ pass
+
+ @abstractmethod
+ def process_batch(self, batch: Batch, port: int) -> Iterator[Optional[BatchLike]]:
+ """
+ Process an input Batch from the given link. The Batch is represented as
+ pandas.DataFrame.
+
+ :param batch: Batch, a batch to be processed.
+ :param port: int, input index of the current Batch.
+ :return: Iterator[Optional[BatchLike]], producing one BatchLike object at a
+ time, or None.
+ """
+ yield
+
+ def close(self) -> None:
+ """
+ Close the context of the operator.
+ """
+ pass
diff --git a/amber/src/main/python/texera_run_python_worker.py b/amber/src/main/python/texera_run_python_worker.py
new file mode 100644
index 00000000000..8687298f819
--- /dev/null
+++ b/amber/src/main/python/texera_run_python_worker.py
@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import sys
+from loguru import logger
+
+from core.python_worker import PythonWorker
+from core.storage.storage_config import StorageConfig
+
+
+def init_loguru_logger(stream_log_level) -> None:
+ """
+ initialize the loguru's logger with the given configurations
+ :param stream_log_level: level to be output to stdout/stderr
+ :return:
+ """
+
+ # loguru has default configuration which includes stderr as the handler. In order to
+ # change the configuration, the easiest way is to remove any existing handlers and
+ # re-configure them.
+ logger.remove()
+
+ # set up stream handler, which outputs to stderr
+ logger.add(sys.stderr, level=stream_log_level)
+
+
+if __name__ == "__main__":
+ (
+ _,
+ worker_id,
+ output_port,
+ logger_level,
+ r_path,
+ iceberg_catalog_type,
+ iceberg_postgres_catalog_uri_without_scheme,
+ iceberg_postgres_catalog_username,
+ iceberg_postgres_catalog_password,
+ iceberg_rest_catalog_uri,
+ iceberg_rest_catalog_warehouse_name,
+ iceberg_table_namespace,
+ iceberg_file_storage_directory_path,
+ iceberg_table_commit_batch_size,
+ s3_endpoint,
+ s3_region,
+ s3_auth_username,
+ s3_auth_password,
+ ) = sys.argv
+ init_loguru_logger(logger_level)
+ StorageConfig.initialize(
+ iceberg_catalog_type,
+ iceberg_postgres_catalog_uri_without_scheme,
+ iceberg_postgres_catalog_username,
+ iceberg_postgres_catalog_password,
+ iceberg_rest_catalog_uri,
+ iceberg_rest_catalog_warehouse_name,
+ iceberg_table_namespace,
+ iceberg_file_storage_directory_path,
+ iceberg_table_commit_batch_size,
+ s3_endpoint,
+ s3_region,
+ s3_auth_username,
+ s3_auth_password,
+ )
+
+ # Setting R_HOME environment variable for R-UDF usage
+ if r_path:
+ import os
+
+ os.environ["R_HOME"] = r_path
+
+ PythonWorker(
+ worker_id=worker_id, host="localhost", output_port=int(output_port)
+ ).run()
diff --git a/amber/src/main/resources/cache.ccf b/amber/src/main/resources/cache.ccf
new file mode 100644
index 00000000000..b03d869250f
--- /dev/null
+++ b/amber/src/main/resources/cache.ccf
@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# DEFAULT CACHE REGION
+
+jcs.default=DC
+jcs.default.cacheattributes=org.apache.commons.jcs3.engine.CompositeCacheAttributes
+jcs.default.cacheattributes.MaxObjects=10
+jcs.default.cacheattributes.MemoryCacheName=org.apache.commons.jcs3.engine.memory.lru.LRUMemoryCache
+jcs.default.elementattributes.IsSpool=true
+
+jcs.auxiliary.DC=org.apache.commons.jcs3.auxiliary.disk.indexed.IndexedDiskCacheFactory
+jcs.auxiliary.DC.attributes=org.apache.commons.jcs3.auxiliary.disk.indexed.IndexedDiskCacheAttributes
+jcs.auxiliary.DC.attributes.DiskPath=/tmp/disk_cache
+jcs.auxiliary.DC.attributes.MaxKeySize=0
\ No newline at end of file
diff --git a/amber/src/main/resources/computing-unit-master-config.yml b/amber/src/main/resources/computing-unit-master-config.yml
new file mode 100644
index 00000000000..0dba594b8ae
--- /dev/null
+++ b/amber/src/main/resources/computing-unit-master-config.yml
@@ -0,0 +1,49 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+server:
+ applicationContextPath: /
+ applicationConnectors:
+ - type: http
+ port: 8085
+
+ # Disable the admin connectors if you don't need an admin interface
+ adminConnectors: []
+
+ # Optional: Minimize the request log configuration if not handling HTTP requests
+ requestLog:
+ type: classic
+ timeZone: UTC
+ appenders: []
+
+logging:
+ level: ${TEXERA_SERVICE_LOG_LEVEL:-INFO}
+ loggers:
+ "io.dropwizard": ${TEXERA_SERVICE_LOG_LEVEL:-INFO}
+ appenders:
+ - type: console
+ logFormat: "[%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n"
+ - type: file
+ currentLogFilename: logs/computing-unit-master.log
+ threshold: ALL
+ queueSize: 512
+ discardingThreshold: 0
+ archive: false
+ timeZone: UTC
+ logFormat: "[%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n"
+ bufferSize: 8KiB
+ immediateFlush: true
diff --git a/amber/src/main/resources/gmail/.gitkeep b/amber/src/main/resources/gmail/.gitkeep
new file mode 100644
index 00000000000..8b137891791
--- /dev/null
+++ b/amber/src/main/resources/gmail/.gitkeep
@@ -0,0 +1 @@
+
diff --git a/amber/src/main/resources/logback.xml b/amber/src/main/resources/logback.xml
new file mode 100644
index 00000000000..43afb5d44c7
--- /dev/null
+++ b/amber/src/main/resources/logback.xml
@@ -0,0 +1,56 @@
+
+
+
+
+ -->
+
+
+ [%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n
+
+
+
+
+
+
+ logs/amber-worker.log
+ true
+
+ logs/amber-worker-%d{yyyy-MM-dd}.log.gz
+
+
+ [%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n
+
+
+
+
+ 8192
+ true
+
+
+
+
+
+
+
+
+
+
+
diff --git a/amber/src/main/resources/texera-compiling-service-web-config.yml b/amber/src/main/resources/texera-compiling-service-web-config.yml
new file mode 100644
index 00000000000..ea2c1b9c1e9
--- /dev/null
+++ b/amber/src/main/resources/texera-compiling-service-web-config.yml
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+server:
+ # modify applicationContextPath if you want the root path to be the name of the application
+ # for example, set it to /twitter, then the url will become texera.ics.uci.edu:port/twitter
+ applicationContextPath: /
+ applicationConnectors:
+ - type: http
+ port: 9090
+ adminConnectors:
+ - type: http
+ port: 9091
+ requestLog:
+ type: classic
+ timeZone: UTC
+ appenders:
+ - type: file
+ currentLogFilename: logs/access.log
+ threshold: ALL
+ queueSize: 512
+ discardingThreshold: 0
+ archive: true
+ archivedLogFilenamePattern: logs/access-%d{yyyy-MM-dd}.log.gz
+ archivedFileCount: 7
+ bufferSize: 8KiB
+ immediateFlush: true
+logging:
+ level: ${TEXERA_SERVICE_LOG_LEVEL:-INFO}
+ loggers:
+ "io.dropwizard": ${TEXERA_SERVICE_LOG_LEVEL:-INFO}
+ appenders:
+ - type: console
+ logFormat: "[%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n"
+ - type: file
+ currentLogFilename: logs/texera-workflow-compiling-service.log
+ threshold: ALL
+ queueSize: 512
+ discardingThreshold: 0
+ archive: false
+ timeZone: UTC
+ logFormat: "[%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n"
+ bufferSize: 8KiB
+ immediateFlush: true
\ No newline at end of file
diff --git a/amber/src/main/resources/web-config.yml b/amber/src/main/resources/web-config.yml
new file mode 100644
index 00000000000..9fde1d078e8
--- /dev/null
+++ b/amber/src/main/resources/web-config.yml
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+server:
+ # modify applicationContextPath if you want the root path to be the name of the application
+ # for example, set it to /twitter, then the url will become texera.ics.uci.edu:port/twitter
+ applicationContextPath: /
+ applicationConnectors:
+ - type: http
+ port: 8080
+ adminConnectors:
+ - type: http
+ port: 8081
+ requestLog:
+ type: classic
+ timeZone: UTC
+ appenders:
+ - type: file
+ currentLogFilename: logs/access.log
+ threshold: ALL
+ queueSize: 512
+ discardingThreshold: 0
+ archive: true
+ archivedLogFilenamePattern: logs/access-%d{yyyy-MM-dd}.log.gz
+ archivedFileCount: 7
+ bufferSize: 8KiB
+ immediateFlush: true
+logging:
+ level: ${TEXERA_SERVICE_LOG_LEVEL:-INFO}
+ loggers:
+ "io.dropwizard": ${TEXERA_SERVICE_LOG_LEVEL:-INFO}
+ appenders:
+ - type: console
+ logFormat: "[%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n"
+ - type: file
+ currentLogFilename: logs/amber-server.log
+ threshold: ALL
+ queueSize: 512
+ discardingThreshold: 0
+ archive: false
+ timeZone: UTC
+ logFormat: "[%date{ISO8601}] [%level] [%logger] [%thread] - %msg %n"
+ bufferSize: 8KiB
+ immediateFlush: true
diff --git a/amber/src/main/scala/org/apache/texera/amber/clustering/ClusterListener.scala b/amber/src/main/scala/org/apache/texera/amber/clustering/ClusterListener.scala
new file mode 100644
index 00000000000..e3e9afab3fb
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/clustering/ClusterListener.scala
@@ -0,0 +1,151 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.clustering
+
+import org.apache.pekko.actor.{Actor, Address}
+import org.apache.pekko.cluster.Cluster
+import org.apache.pekko.cluster.ClusterEvent._
+import com.google.protobuf.timestamp.Timestamp
+import com.twitter.util.{Await, Future}
+import org.apache.texera.amber.clustering.ClusterListener.numWorkerNodesInCluster
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflowruntimestate.FatalErrorType.EXECUTION_FAILURE
+import org.apache.texera.amber.core.workflowruntimestate.WorkflowFatalError
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{
+ COMPLETED,
+ FAILED
+}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.error.ErrorUtils.getStackTraceWithAllCauses
+import org.apache.texera.web.SessionState
+import org.apache.texera.web.model.websocket.response.ClusterStatusUpdateEvent
+import org.apache.texera.web.service.{WorkflowExecutionService, WorkflowService}
+import org.apache.texera.web.storage.ExecutionStateStore.updateWorkflowState
+
+import java.time.Instant
+import scala.collection.mutable.ArrayBuffer
+
+object ClusterListener {
+ final case class GetAvailableNodeAddresses()
+
+ var numWorkerNodesInCluster = 0
+}
+
+class ClusterListener extends Actor with AmberLogging {
+
+ val actorId: ActorVirtualIdentity = ActorVirtualIdentity("ClusterListener")
+ val cluster: Cluster = Cluster(context.system)
+
+ // subscribe to cluster changes, re-subscribe when restart
+ override def preStart(): Unit = {
+ cluster.subscribe(
+ self,
+ initialStateMode = InitialStateAsEvents,
+ classOf[MemberEvent]
+ )
+ }
+
+ override def postStop(): Unit = cluster.unsubscribe(self)
+
+ def receive: Receive = {
+ case evt: MemberEvent =>
+ logger.info(s"received member event = $evt")
+ updateClusterStatus(evt)
+ case ClusterListener.GetAvailableNodeAddresses() =>
+ sender() ! getAllAddress.toArray
+ case other =>
+ println(other)
+ }
+
+ private def getAllAddress: Iterable[Address] = {
+ cluster.state.members
+ .map(_.address)
+ }
+
+ private def forcefullyStop(executionService: WorkflowExecutionService, cause: Throwable): Unit = {
+ executionService.client.shutdown()
+ executionService.executionStateStore.statsStore.updateState(stats =>
+ stats.withEndTimeStamp(System.currentTimeMillis())
+ )
+ executionService.executionStateStore.metadataStore.updateState { metadataStore =>
+ logger.error("forcefully stopping execution", cause)
+ updateWorkflowState(FAILED, metadataStore).addFatalErrors(
+ WorkflowFatalError(
+ EXECUTION_FAILURE,
+ Timestamp(Instant.now),
+ cause.toString,
+ getStackTraceWithAllCauses(cause),
+ "unknown operator"
+ )
+ )
+ }
+ }
+
+ private def updateClusterStatus(evt: MemberEvent): Unit = {
+ evt match {
+ case MemberRemoved(member, status) =>
+ logger.info("Cluster node " + member + " is down!")
+ val futures = new ArrayBuffer[Future[_]]
+ WorkflowService.getAllWorkflowServices.foreach { workflow =>
+ val executionService = workflow.executionService.getValue
+ if (
+ executionService != null && executionService.executionStateStore.metadataStore.getState.state != COMPLETED
+ ) {
+ if (ApplicationConfig.isFaultToleranceEnabled) {
+ logger.info(
+ s"Trigger recovery process for execution id = ${executionService.executionStateStore.metadataStore.getState.executionId.id}"
+ )
+ try {
+ futures.append(executionService.client.notifyNodeFailure(member.address))
+ } catch {
+ case t: Throwable =>
+ logger.warn(
+ s"execution ${executionService.workflowContext.executionId.id} cannot recover! forcing it to stop"
+ )
+ forcefullyStop(executionService, t)
+ }
+ } else {
+ logger.info(
+ s"Kill execution id = ${executionService.executionStateStore.metadataStore.getState.executionId.id}"
+ )
+ forcefullyStop(
+ executionService,
+ new RuntimeException("fault tolerance is not enabled")
+ )
+ }
+ }
+ }
+ Await.all(futures.toSeq: _*)
+ case other => //skip
+ }
+
+ numWorkerNodesInCluster = getAllAddress.size
+ SessionState.getAllSessionStates.foreach { state =>
+ state.send(ClusterStatusUpdateEvent(numWorkerNodesInCluster))
+ }
+
+ logger.info(
+ "---------Now we have " + numWorkerNodesInCluster + s" nodes in the cluster---------"
+ )
+
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/clustering/SingleNodeListener.scala b/amber/src/main/scala/org/apache/texera/amber/clustering/SingleNodeListener.scala
new file mode 100644
index 00000000000..282710524ae
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/clustering/SingleNodeListener.scala
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.clustering
+
+import org.apache.pekko.actor.{Actor, ActorLogging}
+import org.apache.texera.amber.clustering.ClusterListener.GetAvailableNodeAddresses
+
+class SingleNodeListener extends Actor with ActorLogging {
+ override def receive: Receive = {
+ case GetAvailableNodeAddresses() => sender() ! Array(context.self.path.address)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaActorRefMappingService.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaActorRefMappingService.scala
new file mode 100644
index 00000000000..6cad3147031
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaActorRefMappingService.scala
@@ -0,0 +1,141 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.common
+
+import org.apache.pekko.actor.ActorRef
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.{
+ CreditRequest,
+ GetActorRef,
+ NetworkMessage,
+ RegisterActorRef
+}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.virtualidentity.util.{CONTROLLER, SELF}
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+import scala.collection.mutable
+
+class AkkaActorRefMappingService(actorService: AkkaActorService) extends AmberLogging {
+
+ override def actorId: ActorVirtualIdentity = actorService.id
+
+ implicit val self: ActorRef = actorService.self
+
+ private val actorRefMapping: mutable.HashMap[ActorVirtualIdentity, ActorRef] = mutable.HashMap()
+ private val queriedActorVirtualIdentities = new mutable.HashSet[ActorVirtualIdentity]()
+ private val toNotifyOnRegistration =
+ new mutable.HashMap[ActorVirtualIdentity, mutable.Set[ActorRef]]()
+ private val messageStash =
+ new mutable.HashMap[ActorVirtualIdentity, mutable.Queue[NetworkMessage]]
+ actorRefMapping(SELF) = actorService.self
+
+ def getActorRef(id: ActorVirtualIdentity): ActorRef = {
+ actorRefMapping(id)
+ }
+
+ def askForCredit(channelId: ChannelIdentity): Unit = {
+ val id = channelId.toWorkerId
+ if (actorRefMapping.contains(id)) {
+ actorRefMapping(id) ! CreditRequest(channelId)
+ }
+ }
+
+ def hasActorRef(id: ActorVirtualIdentity): Boolean = {
+ actorRefMapping.contains(id)
+ }
+
+ def forwardToActor(msg: NetworkMessage): Unit = {
+ val id = msg.internalMessage.channelId.toWorkerId
+ if (actorRefMapping.contains(id)) {
+ actorRefMapping(id) ! msg
+ } else {
+ val stash = messageStash.getOrElseUpdate(id, new mutable.Queue[NetworkMessage]())
+ stash.enqueue(msg)
+ retrieveActorRef(id, Set())
+ }
+ }
+
+ def removeActorRef(id: ActorVirtualIdentity): Unit = {
+ if (actorRefMapping.contains(id)) {
+ val ref = actorRefMapping.remove(id).get
+ logger.warn(s"actor $id is not reachable anymore, it might have crashed. old ref = $ref")
+ }
+ }
+
+ def registerActorRef(id: ActorVirtualIdentity, ref: ActorRef): Unit = {
+ if (!actorRefMapping.contains(id)) {
+ logger.info(s"register ${VirtualIdentityUtils.toShorterString(id)} -> $ref")
+ actorRefMapping(id) = ref
+ if (messageStash.contains(id)) {
+ val stash = messageStash(id)
+ while (stash.nonEmpty) {
+ ref ! stash.dequeue()
+ }
+ }
+ }
+ if (toNotifyOnRegistration.contains(id)) {
+ toNotifyOnRegistration(id).foreach { toNotify =>
+ toNotify ! RegisterActorRef(id, ref)
+ }
+ toNotifyOnRegistration.remove(id)
+ }
+ }
+
+ def retrieveActorRef(id: ActorVirtualIdentity, replyTo: Set[ActorRef]): Unit = {
+ if (actorRefMapping.contains(id)) {
+ replyTo.foreach { actor =>
+ actor ! RegisterActorRef(id, actorRefMapping(id))
+ }
+ } else if (actorId != CONTROLLER) {
+ // propagation stops at controller
+ if (!queriedActorVirtualIdentities.contains(id)) {
+ try {
+ actorService.parent ! GetActorRef(id, replyTo + actorService.self)
+ queriedActorVirtualIdentities.add(id)
+ } catch {
+ case e: Throwable =>
+ logger.warn(
+ s"Failed to fetch actorRef for ${VirtualIdentityUtils.toShorterString(id)} parentRef = " + actorService.parent
+ )
+ }
+ }
+ } else {
+ // on controller, wait for actor ref registration.
+ logger.warn(s"unknown identifier: ${VirtualIdentityUtils.toShorterString(id)}")
+ val toNotifySet = toNotifyOnRegistration.getOrElseUpdate(id, mutable.HashSet[ActorRef]())
+ replyTo.foreach(toNotifySet.add)
+ }
+ }
+
+ def clearQueriedActorRefs(): Unit = {
+ queriedActorVirtualIdentities.clear()
+ }
+
+ def findActorVirtualIdentity(ref: ActorRef): Option[ActorVirtualIdentity] = {
+ actorRefMapping
+ .find {
+ case (_, actorRef) =>
+ actorRef == ref
+ }
+ .map(_._1)
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaActorService.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaActorService.scala
new file mode 100644
index 00000000000..10a6d7a38c6
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaActorService.scala
@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.common
+
+import org.apache.pekko
+import pekko.actor.{ActorContext, ActorRef, Address, Cancellable, Props}
+import pekko.util.Timeout
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.common.FutureBijection._
+
+import scala.concurrent.ExecutionContext
+import scala.concurrent.duration.{DurationInt, FiniteDuration}
+
+class AkkaActorService(val id: ActorVirtualIdentity, actorContext: ActorContext) {
+
+ implicit def ec: ExecutionContext = actorContext.dispatcher
+
+ implicit val timeout: Timeout = 5.seconds
+ implicit val self: ActorRef = actorContext.self
+
+ def parent: ActorRef = actorContext.parent
+
+ var getAvailableNodeAddressesFunc: () => Array[Address] = () => Array.empty
+
+ def getClusterNodeAddresses: Array[Address] = {
+ getAvailableNodeAddressesFunc()
+ }
+
+ def actorOf(props: Props): ActorRef = {
+ actorContext.actorOf(props)
+ }
+
+ def scheduleOnce(delay: FiniteDuration, callable: () => Unit): Cancellable = {
+ actorContext.system.scheduler.scheduleOnce(delay) {
+ callable()
+ }
+ }
+
+ def scheduleWithFixedDelay(
+ initialDelay: FiniteDuration,
+ delay: FiniteDuration,
+ callable: () => Unit
+ ): Cancellable = {
+ actorContext.system.scheduler.scheduleWithFixedDelay(initialDelay, delay)(() => callable())
+ }
+
+ def sendToSelfOnce(delay: FiniteDuration, msg: Any): Cancellable = {
+ actorContext.system.scheduler.scheduleOnce(delay, actorContext.self, msg)
+ }
+
+ def sendToSelfWithFixedDelay(
+ initialDelay: FiniteDuration,
+ delay: FiniteDuration,
+ msg: Any
+ ): Cancellable = {
+ actorContext.system.scheduler.scheduleWithFixedDelay(
+ initialDelay,
+ delay,
+ actorContext.self,
+ msg
+ )
+ }
+
+ def ask(ref: ActorRef, message: Any): com.twitter.util.Future[Any] = {
+ pekko.pattern.ask(ref, message).asTwitter()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaMessageTransferService.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaMessageTransferService.scala
new file mode 100644
index 00000000000..3401e3ff639
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AkkaMessageTransferService.scala
@@ -0,0 +1,182 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.common
+
+import org.apache.pekko.actor.Cancellable
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkMessage
+import org.apache.texera.amber.engine.architecture.messaginglayer.{CongestionControl, FlowControl}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+
+import scala.collection.mutable
+import scala.concurrent.duration.DurationInt
+
+class AkkaMessageTransferService(
+ actorService: AkkaActorService,
+ refService: AkkaActorRefMappingService,
+ handleBackpressure: Boolean => Unit
+) extends AmberLogging {
+
+ override def actorId: ActorVirtualIdentity = actorService.id
+
+ var resendHandle: Cancellable = Cancellable.alreadyCancelled
+ var creditPollingHandle: Cancellable = Cancellable.alreadyCancelled
+
+ // add congestion control and flow control here
+ val channelToCC = new mutable.HashMap[ChannelIdentity, CongestionControl]()
+ val channelToFC = new mutable.HashMap[ChannelIdentity, FlowControl]()
+ val messageIDToIdentity = new mutable.LongMap[ChannelIdentity]
+
+ private var backpressured = false
+
+ /** keeps track of every outgoing message.
+ * Each message is identified by this monotonic increasing ID.
+ * It's different from the sequence number and it will only
+ * be used by the output gate.
+ */
+ private var networkMessageID = 0L
+
+ def initialize(): Unit = {
+ resendHandle = actorService.scheduleWithFixedDelay(30.seconds, 30.seconds, checkResend)
+ val pollingInterval = ApplicationConfig.creditPollingIntervalInMs.millis
+ creditPollingHandle =
+ actorService.scheduleWithFixedDelay(pollingInterval, pollingInterval, checkCreditPolling)
+ }
+
+ def stop(): Unit = {
+ resendHandle.cancel()
+ creditPollingHandle.cancel()
+ }
+
+ private def checkCreditPolling(): Unit = {
+ channelToFC.foreach {
+ case (channel, fc) =>
+ if (fc.isOverloaded) {
+ refService.askForCredit(channel)
+ }
+ }
+ }
+
+ def send(msg: WorkflowFIFOMessage): Unit = {
+ val networkMessage = NetworkMessage(networkMessageID, msg)
+ messageIDToIdentity(networkMessageID) = msg.channelId
+ networkMessageID += 1
+ forwardToFlowControl(
+ networkMessage,
+ out => forwardToCongestionControl(out, refService.forwardToActor)
+ )
+ }
+
+ private def forwardToFlowControl(
+ msg: NetworkMessage,
+ chainedStep: NetworkMessage => Unit
+ ): Unit = {
+ if (msg.internalMessage.channelId.isControl) {
+ // skip flow control for all control channels
+ chainedStep(msg)
+ } else {
+ val flowControl =
+ channelToFC.getOrElseUpdate(msg.internalMessage.channelId, new FlowControl())
+ flowControl.getMessagesToSend(msg).foreach { msg =>
+ chainedStep(msg)
+ }
+ checkForBackPressure()
+ }
+ }
+
+ private def forwardToCongestionControl(
+ msg: NetworkMessage,
+ chainedStep: NetworkMessage => Unit
+ ): Unit = {
+ val congestionControl =
+ channelToCC.getOrElseUpdate(msg.internalMessage.channelId, new CongestionControl())
+ if (congestionControl.canSend) {
+ congestionControl.markMessageInTransit(msg)
+ chainedStep(msg)
+ } else {
+ congestionControl.enqueueMessage(msg)
+ }
+ }
+
+ def receiveAck(msgId: Long, ackedCredit: Long, queuedCredit: Long): Unit = {
+ if (!messageIDToIdentity.contains(msgId)) {
+ return
+ }
+ val channelId = messageIDToIdentity.remove(msgId).get
+ val congestionControl = channelToCC.getOrElseUpdate(channelId, new CongestionControl())
+ congestionControl.ack(msgId)
+ congestionControl.getBufferedMessagesToSend.foreach { msg =>
+ congestionControl.markMessageInTransit(msg)
+ refService.forwardToActor(msg)
+ }
+ if (channelToFC.contains(channelId)) {
+ channelToFC(channelId).decreaseInflightCredit(ackedCredit)
+ updateChannelCreditFromReceiver(channelId, queuedCredit)
+ }
+ }
+
+ def getAllUnAckedMessages: Iterable[WorkflowFIFOMessage] = {
+ val fcMessages = channelToFC.values.flatMap { fc =>
+ fc.getMessagesToSend.map(_.internalMessage)
+ }
+ val ccMessages = channelToCC.values.flatMap { cc =>
+ cc.getAllMessages.map(_.internalMessage)
+ }
+ fcMessages ++ ccMessages
+ }
+
+ def updateChannelCreditFromReceiver(channelId: ChannelIdentity, queuedCredit: Long): Unit = {
+ val flowControl = channelToFC.getOrElseUpdate(channelId, new FlowControl())
+ flowControl.updateQueuedCredit(queuedCredit)
+ flowControl.getMessagesToSend.foreach(out =>
+ forwardToCongestionControl(out, refService.forwardToActor)
+ )
+ checkForBackPressure()
+ }
+
+ private def checkForBackPressure(): Unit = {
+ val existOverloadedChannel = channelToFC.values.exists(_.isOverloaded)
+ if (backpressured == existOverloadedChannel) {
+ return
+ }
+ backpressured = existOverloadedChannel
+ logger.debug(s"current backpressure status = $backpressured channel credits = ${channelToFC
+ .map(c => c._1 -> c._2.getCredit)}")
+ handleBackpressure(backpressured)
+ }
+
+ private def checkResend(): Unit = {
+ refService.clearQueriedActorRefs()
+ channelToCC.foreach {
+ case (channel, cc) =>
+ val msgsNeedResend = cc.getTimedOutInTransitMessages
+ if (msgsNeedResend.nonEmpty) {
+ logger.debug(s"output for $channel: ${cc.getStatusReport}")
+ }
+ if (refService.hasActorRef(channel.fromWorkerId)) {
+ msgsNeedResend.foreach { msg =>
+ refService.forwardToActor(msg)
+ }
+ }
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AmberProcessor.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AmberProcessor.scala
new file mode 100644
index 00000000000..f1c8136fd85
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AmberProcessor.scala
@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.common
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.messaginglayer.{
+ NetworkInputGateway,
+ NetworkOutputGateway
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.ControlInvocation
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.ReturnInvocation
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage
+import org.apache.texera.amber.engine.architecture.worker.managers.StatisticsManager
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DirectControlMessagePayload,
+ WorkflowFIFOMessage
+}
+import org.apache.texera.amber.engine.common.rpc.{AsyncRPCClient, AsyncRPCServer}
+
+abstract class AmberProcessor(
+ val actorId: ActorVirtualIdentity,
+ @transient var outputHandler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit
+) extends AmberLogging
+ with Serializable {
+
+ /** FIFO & exactly once */
+ val inputGateway: NetworkInputGateway = new NetworkInputGateway(this.actorId)
+
+ // 1. Unified Output
+ val outputGateway: NetworkOutputGateway =
+ new NetworkOutputGateway(
+ this.actorId,
+ msg => {
+ // done by the same thread
+ outputHandler(Right(msg))
+ }
+ )
+ // 2. RPC Layer
+ val asyncRPCClient = new AsyncRPCClient(inputGateway, outputGateway, actorId)
+ val asyncRPCServer: AsyncRPCServer =
+ new AsyncRPCServer(outputGateway, actorId)
+
+ // statistics manager
+ val statisticsManager: StatisticsManager = new StatisticsManager()
+
+ def processDCM(
+ channelId: ChannelIdentity,
+ payload: DirectControlMessagePayload
+ ): Unit = {
+ val controlProcessingStartTime = System.nanoTime();
+ payload match {
+ case invocation: ControlInvocation =>
+ asyncRPCServer.receive(invocation, channelId.fromWorkerId)
+ case ret: ReturnInvocation =>
+ asyncRPCClient.logControlReply(ret, channelId)
+ asyncRPCClient.fulfillPromise(ret)
+ }
+ statisticsManager.increaseControlProcessingTime(System.nanoTime() - controlProcessingStartTime)
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/ExecutorDeployment.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/ExecutorDeployment.scala
new file mode 100644
index 00000000000..cf41297c981
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/ExecutorDeployment.scala
@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.common
+
+import org.apache.pekko.actor.{Address, Deploy}
+import org.apache.pekko.remote.RemoteScope
+import org.apache.texera.amber.core.workflow.{PhysicalOp, PreferController, RoundRobinPreference}
+import org.apache.texera.amber.engine.architecture.controller.execution.OperatorExecution
+import org.apache.texera.amber.engine.architecture.deploysemantics.AddressInfo
+import org.apache.texera.amber.engine.architecture.pythonworker.PythonWorkflowWorker
+import org.apache.texera.amber.engine.architecture.scheduling.config.OperatorConfig
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ FaultToleranceConfig,
+ StateRestoreConfig,
+ WorkerReplayInitialization
+}
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+object ExecutorDeployment {
+
+ def createWorkers(
+ op: PhysicalOp,
+ controllerActorService: AkkaActorService,
+ operatorExecution: OperatorExecution,
+ operatorConfig: OperatorConfig,
+ stateRestoreConfig: Option[StateRestoreConfig],
+ replayLoggingConfig: Option[FaultToleranceConfig]
+ ): Unit = {
+
+ val addressInfo = AddressInfo(
+ controllerActorService.getClusterNodeAddresses,
+ controllerActorService.self.path.address
+ )
+
+ operatorConfig.workerConfigs.foreach(workerConfig => {
+ val workerId = workerConfig.workerId
+ val workerIndex = VirtualIdentityUtils.getWorkerIndex(workerId)
+ val locationPreference = op.locationPreference.getOrElse(RoundRobinPreference)
+ val preferredAddress: Address = locationPreference match {
+ case PreferController =>
+ addressInfo.controllerAddress
+ case RoundRobinPreference =>
+ assert(
+ addressInfo.allAddresses.nonEmpty,
+ "Execution failed to start, no available computation nodes"
+ )
+ addressInfo.allAddresses(workerIndex % addressInfo.allAddresses.length)
+ }
+
+ val workflowWorker = if (op.isPythonBased) {
+ PythonWorkflowWorker.props(workerConfig)
+ } else {
+ WorkflowWorker.props(
+ workerConfig,
+ WorkerReplayInitialization(
+ stateRestoreConfig,
+ replayLoggingConfig
+ )
+ )
+ }
+ // Note: At this point, we don't know if the actor is fully initialized.
+ // Thus, the ActorRef returned from `controllerActorService.actorOf` is ignored.
+ controllerActorService.actorOf(
+ workflowWorker.withDeploy(Deploy(scope = RemoteScope(preferredAddress)))
+ )
+ operatorExecution.initWorkerExecution(workerId)
+ })
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/ProcessingStepCursor.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/ProcessingStepCursor.scala
new file mode 100644
index 00000000000..ac92781737a
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/ProcessingStepCursor.scala
@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.common
+
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+import org.apache.texera.amber.engine.architecture.common.ProcessingStepCursor.INIT_STEP
+
+object ProcessingStepCursor {
+ // step value before processing any incoming message
+ // processing first message will have step = 0
+ val INIT_STEP: Long = -1L
+}
+
+class ProcessingStepCursor {
+ private var currentStepCounter: Long = INIT_STEP
+ private var currentChannel: ChannelIdentity = _
+
+ def setCurrentChannel(channelId: ChannelIdentity): Unit = {
+ currentChannel = channelId
+ }
+
+ def getStep: Long = currentStepCounter
+
+ def getChannel: ChannelIdentity = currentChannel
+
+ def stepIncrement(): Unit = {
+ currentStepCounter += 1
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/WorkflowActor.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/WorkflowActor.scala
new file mode 100644
index 00000000000..5ce64a0a3e0
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/WorkflowActor.scala
@@ -0,0 +1,253 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.common
+
+import org.apache.pekko.actor.{Actor, ActorRef, Address, Stash}
+import org.apache.pekko.pattern.ask
+import org.apache.pekko.util.Timeout
+import org.apache.texera.amber.clustering.ClusterListener.GetAvailableNodeAddresses
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor._
+import org.apache.texera.amber.engine.architecture.logreplay.{
+ ReplayLogGenerator,
+ ReplayLogManager,
+ ReplayLogRecord,
+ ReplayOrderEnforcer
+}
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ FaultToleranceConfig,
+ MainThreadDelegateMessage,
+ StateRestoreConfig,
+ TriggerSend
+}
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.engine.common.{AmberLogging, CheckpointState}
+
+import scala.concurrent.Await
+import scala.concurrent.duration.DurationInt
+
+object WorkflowActor {
+
+ /** Ack for NetworkMessage
+ *
+ * @param messageId Long, id of the received network message
+ * @param ackedCredit Long, received size of the message, used to subtract sender's inflight credit
+ * @param queuedCredit Long, receiver queue's size
+ */
+ final case class NetworkAck(messageId: Long, ackedCredit: Long, queuedCredit: Long)
+
+ final case class MessageBecomesDeadLetter(message: NetworkMessage)
+
+ /** Identifier <-> ActorRef related messages
+ */
+ final case class GetActorRef(id: ActorVirtualIdentity, replyTo: Set[ActorRef])
+
+ final case class RegisterActorRef(id: ActorVirtualIdentity, ref: ActorRef)
+
+ /** All outgoing message should be eventually NetworkMessage
+ *
+ * @param messageId Long, id for a NetworkMessage, used for FIFO and ExactlyOnce
+ * @param internalMessage WorkflowMessage, the message payload
+ */
+ final case class NetworkMessage(messageId: Long, internalMessage: WorkflowFIFOMessage)
+
+ // sent from network communicator to next worker to poll for credit information
+ final case class CreditRequest(channelId: ChannelIdentity)
+
+ final case class CreditResponse(channelId: ChannelIdentity, credit: Long)
+}
+
+abstract class WorkflowActor(
+ replayLogConfOpt: Option[FaultToleranceConfig],
+ val actorId: ActorVirtualIdentity
+) extends Actor
+ with Stash
+ with AmberLogging {
+
+ //
+ // Akka related components:
+ //
+ val actorService: AkkaActorService = new AkkaActorService(actorId, this.context)
+ actorService.getAvailableNodeAddressesFunc = () => {
+ implicit val timeout: Timeout = 5.seconds
+ Await
+ .result(
+ context.actorSelection("/user/cluster-info") ? GetAvailableNodeAddresses(),
+ 5.seconds
+ )
+ .asInstanceOf[Array[Address]]
+ }
+ val actorRefMappingService: AkkaActorRefMappingService = new AkkaActorRefMappingService(
+ actorService
+ )
+ actorRefMappingService.registerActorRef(actorId, self)
+ val transferService: AkkaMessageTransferService =
+ new AkkaMessageTransferService(actorService, actorRefMappingService, handleBackpressure)
+
+ logger.info(s"worker replay log writing conf: $replayLogConfOpt")
+
+ val logStorage: SequentialRecordStorage[ReplayLogRecord] =
+ SequentialRecordStorage.getStorage(replayLogConfOpt.map(_.writeTo))
+ val logManager: ReplayLogManager =
+ ReplayLogManager.createLogManager(logStorage, getLogName, sendMessageFromLogWriterToActor)
+
+ def getLogName: String = actorId.name.replace("Worker:", "")
+
+ def sendMessageFromLogWriterToActor(
+ msg: Either[MainThreadDelegateMessage, WorkflowFIFOMessage]
+ ): Unit = {
+ // limitation: TriggerSend will be processed after input messages before it.
+ msg match {
+ case Left(value) => self ! value
+ case Right(value) => self ! TriggerSend(value)
+ }
+ }
+
+ def handleTriggerSend: Receive = {
+ case TriggerSend(msg) =>
+ transferService.send(msg)
+ }
+
+ def receiveActorRefRelatedMessages: Receive = {
+ case GetActorRef(actorId, replyTo) =>
+ actorRefMappingService.retrieveActorRef(actorId, replyTo)
+ case RegisterActorRef(actorId, ref) =>
+ actorRefMappingService.registerActorRef(actorId, ref)
+ }
+
+ // actor behavior for FIFO messages
+ def receiveMessageAndAck: Receive = {
+ case NetworkMessage(id, workflowMsg @ WorkflowFIFOMessage(channel, _, _)) =>
+ actorRefMappingService.registerActorRef(channel.fromWorkerId, sender())
+ try {
+ handleInputMessage(id, workflowMsg)
+ } catch {
+ case e: Throwable =>
+ logger.warn("actor failed due to exception", e)
+ throw e
+ }
+ case NetworkAck(id, ackedCredit, queuedCredit) =>
+ transferService.receiveAck(id, ackedCredit, queuedCredit)
+ }
+
+ def receiveCreditMessages: Receive = {
+ case CreditRequest(channel) =>
+ sender() ! CreditResponse(channel, getQueuedCredit(channel))
+ case CreditResponse(channel, credit) =>
+ transferService.updateChannelCreditFromReceiver(channel, credit)
+ }
+
+ def receiveDeadLetterMessage: Receive = {
+ case MessageBecomesDeadLetter(msg) =>
+ val dest = msg.internalMessage.channelId.toWorkerId
+ if (dest == actorId) {
+ actorService.scheduleOnce(
+ 100.millis,
+ () => {
+ logger.warn(s"sending message to self failed, retry sending $msg to self directly.")
+ self ! msg
+ }
+ )
+ } else {
+ actorRefMappingService.removeActorRef(dest)
+ }
+ }
+
+ def handleInputMessage(id: Long, workflowMsg: WorkflowFIFOMessage): Unit
+
+ //
+ //flow control:
+ //
+ def getQueuedCredit(channelId: ChannelIdentity): Long
+
+ def handleBackpressure(isBackpressured: Boolean): Unit
+
+ //
+ //Actor lifecycle: Initialization
+ //
+ def initState(): Unit
+
+ def loadFromCheckpoint(chkpt: CheckpointState): Unit
+
+ def setupReplay(
+ amberProcessor: AmberProcessor,
+ stateRestoreConf: StateRestoreConfig,
+ onComplete: () => Unit
+ ): Unit = {
+ val logStorageToRead =
+ SequentialRecordStorage.getStorage[ReplayLogRecord](Some(stateRestoreConf.readFrom))
+ val replayTo = stateRestoreConf.replayDestination
+ if (logStorageToRead.containsFolder(replayTo.toString)) {
+ // checkpoint found
+ val chkptStorage = SequentialRecordStorage.getStorage[CheckpointState](
+ Some(stateRestoreConf.readFrom.resolve(replayTo.toString))
+ )
+ val chkpt = chkptStorage.getReader(getLogName).mkRecordIterator().next()
+ loadFromCheckpoint(chkpt)
+ } else {
+ // do replay from scratch
+ val (processSteps, messages) =
+ ReplayLogGenerator.generate(logStorageToRead, getLogName, replayTo)
+ logger.info(
+ s"setting up replay, " +
+ s"read from ${stateRestoreConf.readFrom} " +
+ s"current step = ${logManager.getStep} " +
+ s"target step = $replayTo " +
+ s"# of log record to replay = ${processSteps.size}"
+ )
+ val orderEnforcer = new ReplayOrderEnforcer(
+ logManager,
+ processSteps,
+ startStep = logManager.getStep,
+ onComplete
+ )
+ amberProcessor.inputGateway.addEnforcer(orderEnforcer)
+ messages.foreach(message =>
+ amberProcessor.inputGateway.getChannel(message.channelId).acceptMessage(message)
+ )
+ }
+ }
+
+ override def preStart(): Unit = {
+ try {
+ transferService.initialize()
+ initState()
+ context.parent ! RegisterActorRef(actorId, context.self)
+ } catch {
+ case t: Throwable =>
+ logger.warn("actor initialization failed due to exception", t)
+ throw t
+ }
+ }
+
+ override def receive: Receive = {
+ receiveActorRefRelatedMessages orElse
+ handleTriggerSend orElse
+ receiveMessageAndAck orElse
+ receiveCreditMessages orElse
+ receiveDeadLetterMessage
+ }
+
+ override def postStop(): Unit = {
+ transferService.stop()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ClientEvent.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ClientEvent.scala
new file mode 100644
index 00000000000..1092af15e77
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ClientEvent.scala
@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessagePayload
+import org.apache.texera.amber.engine.common.executionruntimestate.OperatorMetrics
+
+trait ClientEvent extends WorkflowFIFOMessagePayload
+
+case class ExecutionStateUpdate(state: WorkflowAggregatedState) extends ClientEvent
+
+case class ExecutionStatsUpdate(operatorMetrics: Map[String, OperatorMetrics]) extends ClientEvent
+
+case class RuntimeStatisticsPersist(operatorMetrics: Map[String, OperatorMetrics])
+ extends ClientEvent
+
+case class ReportCurrentProcessingTuple(
+ operatorID: String,
+ tuple: Array[(Tuple, ActorVirtualIdentity)]
+) extends ClientEvent
+
+case class WorkerAssignmentUpdate(workerMapping: Map[String, Seq[String]]) extends ClientEvent
+
+final case class FatalError(e: Throwable, fromActor: Option[ActorVirtualIdentity] = None)
+ extends ClientEvent
+
+case class UpdateExecutorCompleted(id: ActorVirtualIdentity) extends ClientEvent
+
+final case class ReplayStatusUpdate(id: ActorVirtualIdentity, status: Boolean) extends ClientEvent
+
+final case class WorkflowRecoveryStatus(isRecovering: Boolean) extends ClientEvent
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/Controller.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/Controller.scala
new file mode 100644
index 00000000000..daa977d8575
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/Controller.scala
@@ -0,0 +1,251 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.pekko.actor.SupervisorStrategy.Stop
+import org.apache.pekko.actor.{AllForOneStrategy, Props, SupervisorStrategy}
+import org.apache.texera.web.model.websocket.response.RegionUpdateEvent
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+import org.apache.texera.amber.core.workflow.{PhysicalPlan, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkAck
+import org.apache.texera.amber.engine.architecture.common.{ExecutorDeployment, WorkflowActor}
+import org.apache.texera.amber.engine.architecture.controller.execution.OperatorExecution
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ ControlInvocation,
+ EmbeddedControlMessage
+}
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ FaultToleranceConfig,
+ StateRestoreConfig
+}
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DirectControlMessagePayload,
+ WorkflowFIFOMessage
+}
+import org.apache.texera.amber.engine.common.virtualidentity.util.{CLIENT, CONTROLLER, SELF}
+import org.apache.texera.amber.engine.common.{CheckpointState, SerializedState}
+import org.apache.texera.web.SessionState
+
+import scala.concurrent.duration.DurationInt
+
+object ControllerConfig {
+ def default: ControllerConfig =
+ ControllerConfig(
+ statusUpdateIntervalMs = Option(ApplicationConfig.getStatusUpdateIntervalInMs),
+ runtimeStatisticsPersistenceIntervalMs =
+ Option(ApplicationConfig.getRuntimeStatisticsPersistenceIntervalInMs),
+ stateRestoreConfOpt = None,
+ faultToleranceConfOpt = None
+ )
+}
+
+final case class ControllerConfig(
+ statusUpdateIntervalMs: Option[Long],
+ runtimeStatisticsPersistenceIntervalMs: Option[Long],
+ stateRestoreConfOpt: Option[StateRestoreConfig],
+ faultToleranceConfOpt: Option[FaultToleranceConfig]
+)
+
+object Controller {
+
+ def props(
+ workflowContext: WorkflowContext,
+ physicalPlan: PhysicalPlan,
+ controllerConfig: ControllerConfig = ControllerConfig.default
+ ): Props =
+ Props(
+ new Controller(
+ workflowContext,
+ physicalPlan,
+ controllerConfig
+ )
+ )
+}
+
+class Controller(
+ workflowContext: WorkflowContext,
+ physicalPlan: PhysicalPlan,
+ controllerConfig: ControllerConfig
+) extends WorkflowActor(
+ controllerConfig.faultToleranceConfOpt,
+ CONTROLLER
+ ) {
+
+ actorRefMappingService.registerActorRef(CLIENT, context.parent)
+ val controllerTimerService = new ControllerTimerService(controllerConfig, actorService)
+ var cp = new ControllerProcessor(
+ workflowContext,
+ controllerConfig,
+ actorId,
+ logManager.sendCommitted
+ )
+
+ // manages the lifecycle of entire replay process
+ // triggers onStart callback when the first worker/controller marks itself as recovering.
+ // triggers onComplete callback when all worker/controller finishes recovering.
+ private val globalReplayManager = new GlobalReplayManager(
+ () => {
+ //onStart
+ context.parent ! WorkflowRecoveryStatus(true)
+ },
+ () => {
+ //onComplete
+ context.parent ! WorkflowRecoveryStatus(false)
+ }
+ )
+
+ override def initState(): Unit = {
+ attachRuntimeServicesToCPState()
+ cp.workflowScheduler.updateSchedule(physicalPlan)
+
+ val regions: List[(Long, List[String])] =
+ cp.workflowScheduler.getSchedule.getRegions.map { region =>
+ (region.id.id, region.physicalOps.map(_.id.logicalOpId.id).toList)
+ }
+
+ SessionState.getAllSessionStates.foreach { state =>
+ state.send(RegionUpdateEvent(regions))
+ }
+
+ val controllerRestoreConf = controllerConfig.stateRestoreConfOpt
+ if (controllerRestoreConf.isDefined) {
+ globalReplayManager.markRecoveryStatus(CONTROLLER, isRecovering = true)
+ setupReplay(
+ cp,
+ controllerRestoreConf.get,
+ () => {
+ globalReplayManager.markRecoveryStatus(CONTROLLER, isRecovering = false)
+ }
+ )
+ processMessages()
+ }
+ }
+
+ override def handleInputMessage(id: Long, workflowMsg: WorkflowFIFOMessage): Unit = {
+ val channel = cp.inputGateway.getChannel(workflowMsg.channelId)
+ channel.acceptMessage(workflowMsg)
+ sender() ! NetworkAck(id, getInMemSize(workflowMsg), getQueuedCredit(workflowMsg.channelId))
+ processMessages()
+ }
+
+ def processMessages(): Unit = {
+ var waitingForInput = false
+ while (!waitingForInput) {
+ cp.inputGateway.tryPickChannel match {
+ case Some(channel) =>
+ val msg = channel.take
+ val msgToLog = Some(msg).filter(_.payload.isInstanceOf[DirectControlMessagePayload])
+ logManager.withFaultTolerant(msg.channelId, msgToLog) {
+ msg.payload match {
+ case payload: DirectControlMessagePayload => cp.processDCM(msg.channelId, payload)
+ case _: EmbeddedControlMessage => // skip ECM
+ case p => throw new RuntimeException(s"controller cannot handle $p")
+ }
+ }
+ case None =>
+ waitingForInput = true
+ }
+ }
+ }
+
+ def handleDirectInvocation: Receive = {
+ case c: ControlInvocation =>
+ // only client and self can send direction invocations
+ val source = if (sender() == self) {
+ SELF
+ } else {
+ CLIENT
+ }
+ val controlChannelId = ChannelIdentity(source, SELF, isControl = true)
+ val channel = cp.inputGateway.getChannel(controlChannelId)
+ channel.acceptMessage(
+ WorkflowFIFOMessage(controlChannelId, channel.getCurrentSeq, c)
+ )
+ processMessages()
+ }
+
+ def handleReplayMessages: Receive = {
+ case ReplayStatusUpdate(id, status) =>
+ globalReplayManager.markRecoveryStatus(id, status)
+ }
+
+ override def receive: Receive = {
+ super.receive orElse handleDirectInvocation orElse handleReplayMessages
+ }
+
+ /** flow-control */
+ override def getQueuedCredit(channelId: ChannelIdentity): Long = {
+ 0 // no queued credit for controller
+ }
+
+ override def handleBackpressure(isBackpressured: Boolean): Unit = {}
+
+ // Use AllForOneStrategy to stop all children on any fatal error and report it to the client.
+ override val supervisorStrategy: SupervisorStrategy =
+ AllForOneStrategy(maxNrOfRetries = 0, withinTimeRange = 1.minute) {
+ case e: Throwable =>
+ val failedWorker = actorRefMappingService.findActorVirtualIdentity(sender())
+ logger.error(s"Encountered fatal error from $failedWorker, amber is shutting done.", e)
+ cp.asyncRPCClient.sendToClient(
+ FatalError(e, failedWorker)
+ ) // only place to actively report fatal error
+ Stop
+ }
+
+ private def attachRuntimeServicesToCPState(): Unit = {
+ cp.setupActorService(actorService)
+ cp.setupTimerService(controllerTimerService)
+ cp.setupActorRefService(actorRefMappingService)
+ cp.setupLogManager(logManager)
+ cp.setupTransferService(transferService)
+ }
+
+ override def loadFromCheckpoint(chkpt: CheckpointState): Unit = {
+ val cpState: ControllerProcessor = chkpt.load(SerializedState.CP_STATE_KEY)
+ val outputMessages: Array[WorkflowFIFOMessage] = chkpt.load(SerializedState.OUTPUT_MSG_KEY)
+ cp = cpState
+ cp.outputHandler = logManager.sendCommitted
+ attachRuntimeServicesToCPState()
+ // revive all workers.
+ cp.workflowExecution.getRunningRegionExecutions.foreach { regionExecution =>
+ regionExecution.getAllOperatorExecutions.foreach {
+ case (opId, opExecution) =>
+ val op = physicalPlan.getOperator(opId)
+ ExecutorDeployment.createWorkers(
+ op,
+ actorService,
+ OperatorExecution(), //use dummy value here
+ regionExecution.region.resourceConfig.get.operatorConfigs(opId),
+ controllerConfig.stateRestoreConfOpt,
+ controllerConfig.faultToleranceConfOpt
+ )
+ }
+ }
+ outputMessages.foreach(transferService.send)
+ cp.asyncRPCClient.sendToClient(
+ ExecutionStatsUpdate(
+ cp.workflowExecution.getAllRegionExecutionsStats
+ )
+ )
+ globalReplayManager.markRecoveryStatus(CONTROLLER, isRecovering = false)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerAsyncRPCHandlerInitializer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerAsyncRPCHandlerInitializer.scala
new file mode 100644
index 00000000000..2902173364e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerAsyncRPCHandlerInitializer.scala
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.controller.promisehandlers._
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.AsyncRPCContext
+import org.apache.texera.amber.engine.architecture.rpc.controllerservice.ControllerServiceFs2Grpc
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCHandlerInitializer
+
+class ControllerAsyncRPCHandlerInitializer(
+ val cp: ControllerProcessor
+) extends AsyncRPCHandlerInitializer(cp.asyncRPCClient, cp.asyncRPCServer)
+ with ControllerServiceFs2Grpc[Future, AsyncRPCContext]
+ with AmberLogging
+ with LinkWorkersHandler
+ with WorkerExecutionCompletedHandler
+ with WorkerStateUpdatedHandler
+ with PauseHandler
+ with QueryWorkerStatisticsHandler
+ with ResumeHandler
+ with StartWorkflowHandler
+ with PortCompletedHandler
+ with ConsoleMessageHandler
+ with RetryWorkflowHandler
+ with EvaluatePythonExpressionHandler
+ with DebugCommandHandler
+ with TakeGlobalCheckpointHandler
+ with EmbeddedControlMessageHandler
+ with RetrieveWorkflowStateHandler
+ with ReconfigurationHandler {
+ val actorId: ActorVirtualIdentity = cp.actorId
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerProcessor.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerProcessor.scala
new file mode 100644
index 00000000000..7a8e94cf3a7
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerProcessor.scala
@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflow.WorkflowContext
+import org.apache.texera.amber.engine.architecture.common.{
+ AkkaActorRefMappingService,
+ AkkaActorService,
+ AkkaMessageTransferService,
+ AmberProcessor
+}
+import org.apache.texera.amber.engine.architecture.controller.execution.WorkflowExecution
+import org.apache.texera.amber.engine.architecture.logreplay.ReplayLogManager
+import org.apache.texera.amber.engine.architecture.scheduling.WorkflowExecutionCoordinator
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+
+class ControllerProcessor(
+ workflowContext: WorkflowContext,
+ controllerConfig: ControllerConfig,
+ actorId: ActorVirtualIdentity,
+ outputHandler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit
+) extends AmberProcessor(actorId, outputHandler) {
+
+ val workflowExecution: WorkflowExecution = WorkflowExecution()
+ val workflowScheduler: WorkflowScheduler =
+ new WorkflowScheduler(workflowContext, actorId)
+ val workflowExecutionCoordinator: WorkflowExecutionCoordinator = new WorkflowExecutionCoordinator(
+ () => this.workflowScheduler.getNextRegions,
+ workflowExecution,
+ controllerConfig,
+ asyncRPCClient
+ )
+
+ private val initializer = new ControllerAsyncRPCHandlerInitializer(this)
+
+ @transient var controllerTimerService: ControllerTimerService = _
+
+ def setupTimerService(controllerTimerService: ControllerTimerService): Unit = {
+ this.controllerTimerService = controllerTimerService
+ }
+
+ @transient var transferService: AkkaMessageTransferService = _
+
+ def setupTransferService(transferService: AkkaMessageTransferService): Unit = {
+ this.transferService = transferService
+ }
+
+ @transient var actorService: AkkaActorService = _
+
+ def setupActorService(akkaActorService: AkkaActorService): Unit = {
+ this.actorService = akkaActorService
+ }
+
+ @transient var actorRefService: AkkaActorRefMappingService = _
+
+ def setupActorRefService(actorRefService: AkkaActorRefMappingService): Unit = {
+ this.actorRefService = actorRefService
+ this.workflowExecutionCoordinator.setupActorRefService(this.actorRefService)
+ }
+
+ @transient var logManager: ReplayLogManager = _
+
+ def setupLogManager(logManager: ReplayLogManager): Unit = {
+ this.logManager = logManager
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerTimerService.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerTimerService.scala
new file mode 100644
index 00000000000..a778a27c46c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerTimerService.scala
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.pekko.actor.Cancellable
+import org.apache.texera.amber.engine.architecture.common.AkkaActorService
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ QueryStatisticsRequest,
+ StatisticsUpdateTarget
+}
+import org.apache.texera.amber.engine.architecture.rpc.controllerservice.ControllerServiceGrpc.METHOD_CONTROLLER_INITIATE_QUERY_STATISTICS
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
+import org.apache.texera.amber.engine.common.virtualidentity.util.SELF
+
+import scala.concurrent.duration.{DurationInt, FiniteDuration, MILLISECONDS}
+
+class ControllerTimerService(
+ controllerConfig: ControllerConfig,
+ akkaActorService: AkkaActorService
+) {
+ var statusUpdateAskHandle: Option[Cancellable] = None
+ var runtimeStatisticsAskHandle: Option[Cancellable] = None
+
+ private def enableTimer(
+ intervalMs: Option[Long],
+ updateTarget: StatisticsUpdateTarget,
+ handleOpt: Option[Cancellable]
+ ): Option[Cancellable] = {
+ if (intervalMs.nonEmpty && handleOpt.isEmpty) {
+ Option(
+ akkaActorService.sendToSelfWithFixedDelay(
+ 0.milliseconds,
+ FiniteDuration.apply(intervalMs.get, MILLISECONDS),
+ ControlInvocation(
+ METHOD_CONTROLLER_INITIATE_QUERY_STATISTICS,
+ QueryStatisticsRequest(Seq.empty, updateTarget),
+ AsyncRPCContext(SELF, SELF),
+ 0
+ )
+ )
+ )
+ } else {
+ handleOpt
+ }
+ }
+
+ private def disableTimer(handleOpt: Option[Cancellable]): Option[Cancellable] = {
+ if (handleOpt.nonEmpty) {
+ handleOpt.get.cancel()
+ Option.empty
+ } else {
+ handleOpt
+ }
+ }
+
+ def enableStatusUpdate(): Unit = {
+ statusUpdateAskHandle = enableTimer(
+ controllerConfig.statusUpdateIntervalMs,
+ StatisticsUpdateTarget.UI_ONLY,
+ statusUpdateAskHandle
+ )
+ }
+
+ def enableRuntimeStatisticsCollection(): Unit = {
+ runtimeStatisticsAskHandle = enableTimer(
+ controllerConfig.runtimeStatisticsPersistenceIntervalMs,
+ StatisticsUpdateTarget.PERSISTENCE_ONLY,
+ runtimeStatisticsAskHandle
+ )
+ }
+
+ def disableStatusUpdate(): Unit = {
+ statusUpdateAskHandle = disableTimer(statusUpdateAskHandle)
+ }
+
+ def disableRuntimeStatisticsCollection(): Unit = {
+ runtimeStatisticsAskHandle = disableTimer(runtimeStatisticsAskHandle)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/GlobalReplayManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/GlobalReplayManager.scala
new file mode 100644
index 00000000000..b8dc3fb6072
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/GlobalReplayManager.scala
@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+
+import scala.collection.mutable
+
+class GlobalReplayManager(onRecoveryStart: () => Unit, onRecoveryComplete: () => Unit) {
+ private val recovering = mutable.HashSet[ActorVirtualIdentity]()
+
+ def markRecoveryStatus(vid: ActorVirtualIdentity, isRecovering: Boolean): Unit = {
+ val globalRecovering = recovering.nonEmpty
+ if (isRecovering) {
+ recovering.add(vid)
+ } else {
+ recovering.remove(vid)
+ }
+ if (!globalRecovering && recovering.nonEmpty) {
+ onRecoveryStart()
+ }
+ if (globalRecovering && recovering.isEmpty) {
+ onRecoveryComplete()
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/Workflow.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/Workflow.scala
new file mode 100644
index 00000000000..4c2220ae9aa
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/Workflow.scala
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.texera.amber.core.workflow.{PhysicalPlan, WorkflowContext}
+import org.apache.texera.workflow.LogicalPlan
+
+case class Workflow(
+ context: WorkflowContext,
+ logicalPlan: LogicalPlan,
+ physicalPlan: PhysicalPlan
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/WorkflowScheduler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/WorkflowScheduler.scala
new file mode 100644
index 00000000000..b1acb3c0650
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/WorkflowScheduler.scala
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflow.{PhysicalPlan, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.scheduling.{
+ CostBasedScheduleGenerator,
+ Region,
+ Schedule
+}
+
+class WorkflowScheduler(
+ workflowContext: WorkflowContext,
+ actorId: ActorVirtualIdentity
+) extends java.io.Serializable {
+ var physicalPlan: PhysicalPlan = _
+ private var schedule: Schedule = _
+
+ def getSchedule: Schedule = schedule
+
+ /**
+ * Update the schedule to be executed, based on the given physicalPlan.
+ */
+ def updateSchedule(physicalPlan: PhysicalPlan): Unit = {
+ // generate a schedule using a region plan generator.
+ val (generatedSchedule, updatedPhysicalPlan) =
+ // CostBasedRegionPlanGenerator considers costs to try to find an optimal plan.
+ new CostBasedScheduleGenerator(
+ workflowContext,
+ physicalPlan,
+ actorId
+ ).generate()
+ this.schedule = generatedSchedule
+ this.physicalPlan = updatedPhysicalPlan
+ }
+
+ def getNextRegions: Set[Region] = if (!schedule.hasNext) Set() else schedule.next()
+
+ def hasPendingRegions: Boolean = schedule != null && schedule.hasNext
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/ChannelExecution.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/ChannelExecution.scala
new file mode 100644
index 00000000000..960cad507ea
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/ChannelExecution.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+case class ChannelExecution()
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/ExecutionUtils.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/ExecutionUtils.scala
new file mode 100644
index 00000000000..7ee9bc04735
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/ExecutionUtils.scala
@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.worker.statistics.{
+ PortTupleMetricsMapping,
+ TupleMetrics
+}
+import org.apache.texera.amber.engine.common.executionruntimestate.{
+ OperatorMetrics,
+ OperatorStatistics
+}
+
+object ExecutionUtils {
+
+ /**
+ * Handle the case when a logical operator has two physical operators within a same region (e.g., Aggregate operator)
+ */
+ def aggregateMetrics(metrics: Iterable[OperatorMetrics]): OperatorMetrics = {
+ if (metrics.isEmpty) {
+ // Return a default OperatorMetrics if metrics are empty
+ return OperatorMetrics(
+ WorkflowAggregatedState.UNINITIALIZED,
+ OperatorStatistics(Seq.empty, Seq.empty, 0, 0, 0, 0)
+ )
+ }
+
+ val aggregatedState = aggregateStates(
+ metrics.map(_.operatorState),
+ WorkflowAggregatedState.COMPLETED,
+ WorkflowAggregatedState.TERMINATED,
+ WorkflowAggregatedState.RUNNING,
+ WorkflowAggregatedState.UNINITIALIZED,
+ WorkflowAggregatedState.PAUSED,
+ WorkflowAggregatedState.READY
+ )
+
+ def sumMetrics(
+ extractor: OperatorMetrics => Iterable[PortTupleMetricsMapping]
+ ): Seq[PortTupleMetricsMapping] = {
+ val filteredMetrics = metrics.flatMap(extractor).filterNot(_.portId.internal)
+ aggregatePortMetrics(filteredMetrics)
+ }
+
+ val inputMetricsSum = sumMetrics(_.operatorStatistics.inputMetrics)
+ val outputMetricsSum = sumMetrics(_.operatorStatistics.outputMetrics)
+
+ val numWorkersSum = metrics.map(_.operatorStatistics.numWorkers).sum
+ val dataProcessingTimeSum = metrics.map(_.operatorStatistics.dataProcessingTime).sum
+ val controlProcessingTimeSum = metrics.map(_.operatorStatistics.controlProcessingTime).sum
+ val idleTimeSum = metrics.map(_.operatorStatistics.idleTime).sum
+
+ OperatorMetrics(
+ aggregatedState,
+ OperatorStatistics(
+ inputMetricsSum,
+ outputMetricsSum,
+ numWorkersSum,
+ dataProcessingTimeSum,
+ controlProcessingTimeSum,
+ idleTimeSum
+ )
+ )
+ }
+
+ def aggregateStates[T](
+ states: Iterable[T],
+ completedState: T,
+ terminatedState: T,
+ runningState: T,
+ uninitializedState: T,
+ pausedState: T,
+ readyState: T
+ ): WorkflowAggregatedState = {
+ states match {
+ case _ if states.isEmpty => WorkflowAggregatedState.UNINITIALIZED
+ case _ if states.forall(_ == completedState) => WorkflowAggregatedState.COMPLETED
+ case _ if states.forall(_ == terminatedState) => WorkflowAggregatedState.COMPLETED
+ case _ if states.exists(_ == runningState) => WorkflowAggregatedState.RUNNING
+ case _ =>
+ val unCompletedStates = states.filter(_ != completedState)
+ if (unCompletedStates.forall(_ == uninitializedState)) {
+ WorkflowAggregatedState.UNINITIALIZED
+ } else if (unCompletedStates.forall(_ == pausedState)) {
+ WorkflowAggregatedState.PAUSED
+ } else if (unCompletedStates.forall(_ == readyState)) {
+ WorkflowAggregatedState.RUNNING
+ } else {
+ WorkflowAggregatedState.UNKNOWN
+ }
+ }
+ }
+
+ def aggregatePortMetrics(
+ metrics: Iterable[PortTupleMetricsMapping]
+ ): Seq[PortTupleMetricsMapping] = {
+ metrics
+ .groupBy(_.portId)
+ .view
+ .map {
+ case (portId, mappings) =>
+ val totalCount = mappings.map(_.tupleMetrics.count).sum
+ val totalSize = mappings.map(_.tupleMetrics.size).sum
+ PortTupleMetricsMapping(portId, TupleMetrics(totalCount, totalSize))
+ }
+ .toSeq
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/LinkExecution.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/LinkExecution.scala
new file mode 100644
index 00000000000..434636e16f6
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/LinkExecution.scala
@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+
+import scala.collection.mutable
+
+case class LinkExecution() {
+ private val channelExecutions: mutable.Map[ChannelIdentity, ChannelExecution] = mutable.HashMap()
+
+ def initChannelExecution(channelId: ChannelIdentity): Unit = {
+ assert(!channelExecutions.contains(channelId))
+ channelExecutions(channelId) = ChannelExecution()
+ }
+
+ def getAllChannelExecutions: Iterable[(ChannelIdentity, ChannelExecution)] = channelExecutions
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/OperatorExecution.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/OperatorExecution.scala
new file mode 100644
index 00000000000..5a7f57083a4
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/OperatorExecution.scala
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.controller.execution.ExecutionUtils.aggregateStates
+import org.apache.texera.amber.engine.architecture.deploysemantics.layer.WorkerExecution
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.worker.statistics.{
+ PortTupleMetricsMapping,
+ WorkerState
+}
+import org.apache.texera.amber.engine.common.executionruntimestate.{
+ OperatorMetrics,
+ OperatorStatistics
+}
+
+import java.util
+import scala.jdk.CollectionConverters._
+
+case class OperatorExecution() {
+
+ private val workerExecutions =
+ new util.concurrent.ConcurrentHashMap[ActorVirtualIdentity, WorkerExecution]()
+
+ /**
+ * Initializes a `WorkerExecution` for the specified workerId and adds it to the workerExecutions map.
+ * If a `WorkerExecution` for the given workerId already exists, an AssertionError is thrown.
+ * After successfully adding the new `WorkerExecution`, it retrieves and returns the newly added instance.
+ *
+ * @param workerId The `ActorVirtualIdentity` representing the unique identity of the worker.
+ * @return The `WorkerExecution` instance associated with the specified workerId.
+ * @throws AssertionError if a `WorkerExecution` already exists for the given workerId.
+ */
+ def initWorkerExecution(workerId: ActorVirtualIdentity): WorkerExecution = {
+ assert(
+ !workerExecutions.contains(workerId),
+ s"WorkerExecution already exists for workerId: $workerId"
+ )
+ workerExecutions.put(workerId, WorkerExecution())
+ getWorkerExecution(workerId)
+ }
+
+ /**
+ * Retrieves the `WorkerExecution` instance associated with the specified workerId.
+ */
+ def getWorkerExecution(workerId: ActorVirtualIdentity): WorkerExecution =
+ workerExecutions.get(workerId)
+
+ /**
+ * Retrieves the set of all workerIds for which `WorkerExecution` instances have been initialized.
+ */
+ def getWorkerIds: Set[ActorVirtualIdentity] = workerExecutions.keys.asScala.toSet
+
+ def getState: WorkflowAggregatedState = {
+ val workerStates = workerExecutions.values.asScala.map(_.getState)
+ aggregateStates(
+ workerStates,
+ WorkerState.COMPLETED,
+ WorkerState.TERMINATED,
+ WorkerState.RUNNING,
+ WorkerState.UNINITIALIZED,
+ WorkerState.PAUSED,
+ WorkerState.READY
+ )
+ }
+
+ private[this] def computeOperatorPortStats(
+ workerPortStats: Iterable[PortTupleMetricsMapping]
+ ): Seq[PortTupleMetricsMapping] = {
+ ExecutionUtils.aggregatePortMetrics(workerPortStats)
+ }
+
+ def getStats: OperatorMetrics = {
+ val workerRawStats = workerExecutions.values.asScala.map(_.getStats)
+ val inputMetrics = workerRawStats.flatMap(_.inputTupleMetrics)
+ val outputMetrics = workerRawStats.flatMap(_.outputTupleMetrics)
+ OperatorMetrics(
+ getState,
+ OperatorStatistics(
+ inputMetrics = computeOperatorPortStats(inputMetrics),
+ outputMetrics = computeOperatorPortStats(outputMetrics),
+ getWorkerIds.size,
+ dataProcessingTime = workerRawStats.map(_.dataProcessingTime).sum,
+ controlProcessingTime = workerRawStats.map(_.controlProcessingTime).sum,
+ idleTime = workerRawStats.map(_.idleTime).sum
+ )
+ )
+ }
+
+ def isInputPortCompleted(portId: PortIdentity): Boolean = {
+ workerExecutions
+ .values()
+ .asScala
+ .map(workerExecution => workerExecution.getInputPortExecution(portId))
+ .forall(_.completed)
+ }
+
+ def isOutputPortCompleted(portId: PortIdentity): Boolean = {
+ workerExecutions
+ .values()
+ .asScala
+ .map(workerExecution => workerExecution.getOutputPortExecution(portId))
+ .forall(_.completed)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/RegionExecution.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/RegionExecution.scala
new file mode 100644
index 00000000000..d5939c2e3b1
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/RegionExecution.scala
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+import com.rits.cloning.Cloner
+import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
+import org.apache.texera.amber.core.workflow.PhysicalLink
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.scheduling.Region
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerStatistics
+import org.apache.texera.amber.engine.common.executionruntimestate.OperatorMetrics
+
+import scala.collection.mutable
+
+object Cloning {
+ val cloner = new Cloner()
+ // prevent cloner from cloning scala Nil, which it cannot handle properly
+ cloner.dontClone(classOf[WorkerStatistics])
+}
+
+case class RegionExecution(region: Region) {
+
+ private val operatorExecutions: mutable.Map[PhysicalOpIdentity, OperatorExecution] =
+ mutable.HashMap()
+
+ private val linkExecutions: mutable.Map[PhysicalLink, LinkExecution] = mutable.HashMap()
+
+ /**
+ * Initializes and retrieves an `OperatorExecution` for a given physical operatorId.
+ * Optionally, an OperatorExecution instance (from other regionExecutions) can
+ * be provided to make a copy.
+ * If an existing `OperatorExecution` is not provided, it creates a new one.
+ * An assertion error is thrown if initialization is attempted for an already existing
+ * operatorId.
+ *
+ * @param physicalOpId The physical operatorId for which to initialize or retrieve the execution.
+ * @param inheritOperatorExecution An optional `OperatorExecution` to make a copy.
+ * @return The `OperatorExecution` associated with the given physical operatorId.
+ * @throws AssertionError if the `OperatorExecution` has already been initialized.
+ */
+ def initOperatorExecution(
+ physicalOpId: PhysicalOpIdentity,
+ inheritOperatorExecution: Option[OperatorExecution] = None
+ ): OperatorExecution = {
+ assert(!operatorExecutions.contains(physicalOpId), "OperatorExecution already exists.")
+
+ operatorExecutions.getOrElseUpdate(
+ physicalOpId,
+ inheritOperatorExecution
+ .map(operatorExecution => Cloning.cloner.deepClone(operatorExecution))
+ .getOrElse(OperatorExecution())
+ )
+ }
+
+ /**
+ * Retrieves an `OperatorExecution` for the specified operatorId.
+ *
+ * @param opId The ID of the operator whose execution is to be retrieved.
+ * @return The `OperatorExecution` associated with the specified ID.
+ */
+ def getOperatorExecution(opId: PhysicalOpIdentity): OperatorExecution = operatorExecutions(opId)
+
+ /**
+ * Checks if an `OperatorExecution` exists for the specified operatorId.
+ *
+ * @param opId The identifier of the operator to check.
+ * @return True if an execution exists for the operatorId, false otherwise.
+ */
+ def hasOperatorExecution(opId: PhysicalOpIdentity): Boolean = operatorExecutions.contains(opId)
+
+ /**
+ * Retrieves all `OperatorExecutions` stored.
+ */
+ def getAllOperatorExecutions: Iterable[(PhysicalOpIdentity, OperatorExecution)] =
+ operatorExecutions
+
+ /**
+ * Initializes a `LinkExecution` for a given physical link. Creates a new `LinkExecution`
+ * if one does not already exist for the link.
+ * An assertion error is thrown if initialization is attempted for an already existing link.
+ *
+ * @param link The `PhysicalLink` for which to initialize the `LinkExecution`.
+ * @return The newly initialized `LinkExecution`.
+ * @throws AssertionError if the `LinkExecution` has already been initialized for the link.
+ */
+ def initLinkExecution(link: PhysicalLink): LinkExecution = {
+ assert(!linkExecutions.contains(link))
+ linkExecutions.getOrElseUpdate(link, new LinkExecution())
+ }
+
+ /**
+ * Retrieves all `LinkExecutions` stored.
+ */
+ def getAllLinkExecutions: Iterable[(PhysicalLink, LinkExecution)] = linkExecutions
+
+ def getStats: Map[PhysicalOpIdentity, OperatorMetrics] = {
+ operatorExecutions.map {
+ case (physicalOpId, operatorExecution) =>
+ physicalOpId -> operatorExecution.getStats
+ }.toMap
+ }
+
+ def isCompleted: Boolean = getState == WorkflowAggregatedState.COMPLETED
+
+ def getState: WorkflowAggregatedState = {
+ if (
+ region.getPorts.forall(globalPortId => {
+ val operatorExecution = this.getOperatorExecution(globalPortId.opId)
+ if (globalPortId.input) operatorExecution.isInputPortCompleted(globalPortId.portId)
+ else operatorExecution.isOutputPortCompleted(globalPortId.portId)
+ })
+ ) {
+ WorkflowAggregatedState.COMPLETED
+ } else {
+ WorkflowAggregatedState.RUNNING
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkerPortExecution.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkerPortExecution.scala
new file mode 100644
index 00000000000..0b7c3a39397
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkerPortExecution.scala
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+case class WorkerPortExecution() {
+ var completed: Boolean = false
+
+ def setCompleted(): Unit = {
+ completed = true
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkflowExecution.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkflowExecution.scala
new file mode 100644
index 00000000000..c1e44bd5cc8
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkflowExecution.scala
@@ -0,0 +1,183 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
+import org.apache.texera.amber.engine.architecture.controller.execution.ExecutionUtils.aggregateMetrics
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState._
+import org.apache.texera.amber.engine.architecture.scheduling.{Region, RegionIdentity}
+import org.apache.texera.amber.engine.common.executionruntimestate.OperatorMetrics
+
+import scala.collection.mutable
+
+case class WorkflowExecution() {
+
+ // region executions are stored with LinkedHashMap to maintain their creation order.
+ private val regionExecutions: mutable.LinkedHashMap[RegionIdentity, RegionExecution] =
+ mutable.LinkedHashMap()
+
+ /**
+ * Initializes or retrieves a `RegionExecution` for a given `Region`. If not already
+ * initialized, it creates and returns a new `RegionExecution`; otherwise, an assertion
+ * error is thrown if re-initialization is attempted.
+ *
+ * @param region The `Region` for which to initialize or retrieve the `RegionExecution`.
+ * @return The `RegionExecution` associated with the given `Region`.
+ * @throws AssertionError if the `RegionExecution` has already been initialized.
+ */
+ def initRegionExecution(region: Region): RegionExecution = {
+ // ensure the region execution hasn't been initialized already.
+ assert(
+ !regionExecutions.contains(region.id),
+ s"RegionExecution of ${region.id} already initialized."
+ )
+ regionExecutions.getOrElseUpdate(region.id, RegionExecution(region))
+ }
+
+ def restartRegionExecution(region: Region): RegionExecution = {
+ regionExecutions.get(region.id).foreach { existingRegionExecution =>
+ assert(
+ existingRegionExecution.isCompleted,
+ s"Cannot restart running RegionExecution of ${region.id}."
+ )
+ }
+ val regionExecution = RegionExecution(region)
+ regionExecutions.put(region.id, regionExecution)
+ regionExecution
+ }
+
+ /**
+ * Retrieves a specific `RegionExecution` by its identifier.
+ *
+ * @param regionId The unique identifier of the region for which the execution is to be retrieved.
+ * @return The `RegionExecution` associated with the specified `regionId`.
+ */
+ def getRegionExecution(regionId: RegionIdentity): RegionExecution = regionExecutions(regionId)
+
+ def hasRegionExecution(regionId: RegionIdentity): Boolean = regionExecutions.contains(regionId)
+
+ /**
+ * Retrieves all `RegionExecutions` that are currently in running state,
+ * preserving the order in which they were created.
+ *
+ * This method filters the executions to include only those that have not completed.
+ *
+ * @return An `Iterable` of `RegionExecution` objects that are in running state.
+ */
+ def getRunningRegionExecutions: Iterable[RegionExecution] = {
+ regionExecutions.values.filterNot(_.isCompleted)
+ }
+
+ /**
+ * Retrieve the runtime stats of all `RegionExecutions`
+ *
+ * @return A `Map` with key being `Logical Operator ID` and the value being operator runtime statistics
+ */
+ def getAllRegionExecutionsStats: Map[String, OperatorMetrics] = {
+ val allRegionExecutions: Iterable[RegionExecution] = getAllRegionExecutions
+
+ val statsMap: Map[PhysicalOpIdentity, OperatorMetrics] = allRegionExecutions.flatMap {
+ regionExecution =>
+ regionExecution.getStats.map {
+ case (physicalOpIdentity, operatorMetrics) =>
+ (physicalOpIdentity, operatorMetrics)
+ }
+ }.toMap
+
+ val aggregatedStats: Map[String, OperatorMetrics] =
+ statsMap.groupBy(_._1.logicalOpId.id).map {
+ case (logicalOpId, stats) =>
+ (logicalOpId, aggregateMetrics(stats.values))
+ }
+ aggregatedStats
+ }
+
+ /**
+ * Retrieves all `RegionExecutions`, preserving the order in which they were created.
+ *
+ * This method provides access to all executions, regardless of their state.
+ *
+ * @return An `Iterable` of all `RegionExecution` objects in the order they were added.
+ */
+ def getAllRegionExecutions: Iterable[RegionExecution] = regionExecutions.values
+
+ /**
+ * Retrieves the latest `OperatorExecution` associated with the specified physical operatorId.
+ *
+ * This method searches through all `RegionExecutions` in reverse creation order to find the most recent
+ * `OperatorExecution` that matches the given physical operatorId. It assumes that each `RegionExecution`
+ * may contain zero or exactly one `OperatorExecution` instance, and it returns the latest one found that
+ * corresponds to the specified operatorId.
+ *
+ * @param physicalOpId The unique identifier of the physical operator for which the latest execution is
+ * to be retrieved.
+ * @return The latest `OperatorExecution` instance associated with the given physical operatorId.
+ * @throws NoSuchElementException if no `OperatorExecution` is found for the specified operatorId.
+ */
+ def getLatestOperatorExecution(physicalOpId: PhysicalOpIdentity): OperatorExecution = {
+ getLatestOperatorExecutionOption(physicalOpId).get
+ }
+
+ /**
+ * Returns the latest `OperatorExecution` for a physical operator if it has been initialized.
+ *
+ * This is the safe counterpart of `getLatestOperatorExecution` for callers that may traverse
+ * operators before their region is launched (e.g., full-graph stats queries while execution is still
+ * progressing through schedule levels).
+ */
+ def getLatestOperatorExecutionOption(
+ physicalOpId: PhysicalOpIdentity
+ ): Option[OperatorExecution] = {
+ regionExecutions.values.toSeq
+ .findLast(regionExecution => regionExecution.hasOperatorExecution(physicalOpId))
+ .map(_.getOperatorExecution(physicalOpId))
+ }
+
+ def isCompleted: Boolean = getState == WorkflowAggregatedState.COMPLETED
+
+ def getState: WorkflowAggregatedState = {
+ val regionStates = regionExecutions.values.map(_.getState)
+ if (regionStates.isEmpty) {
+ return WorkflowAggregatedState.UNINITIALIZED
+ }
+ if (regionStates.forall(_ == COMPLETED)) {
+ return WorkflowAggregatedState.COMPLETED
+ }
+ val unCompletedOpStates = regionExecutions.values
+ .filter(_.getState != COMPLETED)
+ .flatMap(_.getAllOperatorExecutions.map(_._2.getState))
+ .filter(_ != COMPLETED)
+ if (unCompletedOpStates.forall(_ == UNINITIALIZED)) {
+ return WorkflowAggregatedState.UNINITIALIZED
+ }
+ val runningOpStates = unCompletedOpStates.filter(_ != UNINITIALIZED)
+ if (runningOpStates.exists(_ == RUNNING)) {
+ WorkflowAggregatedState.RUNNING
+ } else if (runningOpStates.forall(_ == PAUSED)) {
+ WorkflowAggregatedState.PAUSED
+ } else if (runningOpStates.forall(_ == READY)) {
+ WorkflowAggregatedState.READY
+ } else {
+ WorkflowAggregatedState.UNKNOWN
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ConsoleMessageHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ConsoleMessageHandler.scala
new file mode 100644
index 00000000000..4030ad7d3ee
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ConsoleMessageHandler.scala
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ ConsoleMessageTriggeredRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+
+trait ConsoleMessageHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def consoleMessageTriggered(
+ msg: ConsoleMessageTriggeredRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ // forward message to frontend
+ sendToClient(msg.consoleMessage)
+ EmptyReturn()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/DebugCommandHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/DebugCommandHandler.scala
new file mode 100644
index 00000000000..2e18ec6ba14
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/DebugCommandHandler.scala
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ DebugCommandRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+
+trait DebugCommandHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def debugCommand(msg: DebugCommandRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = {
+ workerInterface.debugCommand(msg, mkContext(ActorVirtualIdentity(msg.workerId)))
+ EmptyReturn()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/EmbeddedControlMessageHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/EmbeddedControlMessageHandler.scala
new file mode 100644
index 00000000000..45a4491ad95
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/EmbeddedControlMessageHandler.scala
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ ControlInvocation,
+ PropagateEmbeddedControlMessageRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ ControlReturn,
+ PropagateEmbeddedControlMessageResponse
+}
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+trait EmbeddedControlMessageHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def propagateEmbeddedControlMessage(
+ msg: PropagateEmbeddedControlMessageRequest,
+ ctx: AsyncRPCContext
+ ): Future[PropagateEmbeddedControlMessageResponse] = {
+ // step1: create separate control commands for each target actor.
+ val inputSet = msg.targetOps.flatMap { target =>
+ cp.workflowExecution.getRunningRegionExecutions
+ .map(_.getOperatorExecution(target))
+ .flatMap(_.getWorkerIds.map { worker =>
+ worker -> createInvocation(msg.methodName, msg.command, worker)
+ })
+ }
+ // step 2: packing all control commands into one compound command.
+ val cmdMapping: Map[String, ControlInvocation] = inputSet.map {
+ case (workerId, (control, _)) => (workerId.name, control)
+ }.toMap
+ val futures: Set[Future[(ActorVirtualIdentity, ControlReturn)]] = inputSet.map {
+ case (workerId, (_, future)) => future.map(ret => (workerId, ret.asInstanceOf[ControlReturn]))
+ }.toSet
+
+ // step 3: convert scope DAG to channels.
+ val channelScope = cp.workflowExecution.getRunningRegionExecutions
+ .flatMap(regionExecution =>
+ regionExecution.getAllLinkExecutions
+ .map(_._2)
+ .flatMap(linkExecution => linkExecution.getAllChannelExecutions.map(_._1))
+ )
+ .filter(channelId => {
+ msg.scope
+ .contains(VirtualIdentityUtils.getPhysicalOpId(channelId.fromWorkerId)) &&
+ msg.scope
+ .contains(VirtualIdentityUtils.getPhysicalOpId(channelId.toWorkerId))
+ })
+ val controlChannels = msg.sourceOpToStartProp.flatMap { source =>
+ cp.workflowExecution.getLatestOperatorExecution(source).getWorkerIds.flatMap { worker =>
+ Seq(
+ ChannelIdentity(CONTROLLER, worker, isControl = true),
+ ChannelIdentity(worker, CONTROLLER, isControl = true)
+ )
+ }
+ }
+
+ val finalScope = channelScope ++ controlChannels
+
+ // step 4: start prop, send ECM through control channel with the compound command from sources.
+ msg.sourceOpToStartProp.foreach { source =>
+ cp.workflowExecution.getLatestOperatorExecution(source).getWorkerIds.foreach { worker =>
+ sendECM(
+ msg.id,
+ msg.ecmType,
+ finalScope.toSet,
+ cmdMapping,
+ ChannelIdentity(actorId, worker, isControl = true)
+ )
+ }
+ }
+
+ // step 5: wait for the ECM propagation.
+ Future.collect(futures.toList).map { ret =>
+ cp.logManager.markAsReplayDestination(msg.id)
+ PropagateEmbeddedControlMessageResponse(ret.map(x => (x._1.name, x._2)).toMap)
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/EvaluatePythonExpressionHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/EvaluatePythonExpressionHandler.scala
new file mode 100644
index 00000000000..6a235b2e77f
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/EvaluatePythonExpressionHandler.scala
@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.OperatorIdentity
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EvaluatePythonExpressionRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EvaluatePythonExpressionResponse
+
+trait EvaluatePythonExpressionHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def evaluatePythonExpression(
+ msg: EvaluatePythonExpressionRequest,
+ ctx: AsyncRPCContext
+ ): Future[EvaluatePythonExpressionResponse] = {
+ val logicalOpId = new OperatorIdentity(msg.operatorId)
+ val physicalOps = cp.workflowScheduler.physicalPlan.getPhysicalOpsOfLogicalOp(logicalOpId)
+ if (physicalOps.size != 1) {
+ val msg =
+ s"logical operator $logicalOpId has ${physicalOps.size} physical operators, expecting a single one"
+ throw new RuntimeException(msg)
+ }
+
+ val physicalOp = physicalOps.head
+ val opExecution = cp.workflowExecution.getLatestOperatorExecution(physicalOp.id)
+
+ Future
+ .collect(
+ opExecution.getWorkerIds
+ .map(worker => workerInterface.evaluatePythonExpression(msg, mkContext(worker)))
+ .toList
+ )
+ .map(evaluatedValues => {
+ EvaluatePythonExpressionResponse(evaluatedValues)
+ })
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/LinkWorkersHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/LinkWorkersHandler.scala
new file mode 100644
index 00000000000..f8a967aad73
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/LinkWorkersHandler.scala
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AddInputChannelRequest,
+ AddPartitioningRequest,
+ AsyncRPCContext,
+ LinkWorkersRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+
+/** add a data transfer partitioning to the sender workers and update input linking
+ * for the receiver workers of a link strategy.
+ *
+ * possible sender: controller, client
+ */
+trait LinkWorkersHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def linkWorkers(msg: LinkWorkersRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = {
+ val region = cp.workflowExecutionCoordinator.getRegionOfLink(msg.link)
+ val resourceConfig = region.resourceConfig.get
+ val linkConfig = resourceConfig.linkConfigs(msg.link)
+ val linkExecution =
+ cp.workflowExecution.getRegionExecution(region.id).initLinkExecution(msg.link)
+ val futures = linkConfig.channelConfigs
+ .map(_.channelId)
+ .flatMap(channelId => {
+ linkExecution.initChannelExecution(channelId)
+ Seq(
+ workerInterface.addPartitioning(
+ AddPartitioningRequest(msg.link, linkConfig.partitioning),
+ mkContext(channelId.fromWorkerId)
+ ),
+ workerInterface.addInputChannel(
+ AddInputChannelRequest(channelId, msg.link.toPortId),
+ mkContext(channelId.toWorkerId)
+ )
+ )
+ })
+
+ Future.collect(futures).map { _ =>
+ // returns when all has completed
+ EmptyReturn()
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PauseHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PauseHandler.scala
new file mode 100644
index 00000000000..35a85f56ae9
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PauseHandler.scala
@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerAsyncRPCHandlerInitializer,
+ ExecutionStateUpdate,
+ ExecutionStatsUpdate,
+ RuntimeStatisticsPersist
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ EmptyReturn,
+ WorkerMetricsResponse
+}
+
+import scala.collection.mutable
+
+/** pause the entire workflow
+ *
+ * possible sender: client, controller
+ */
+trait PauseHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def pauseWorkflow(request: EmptyRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = {
+ cp.controllerTimerService.disableStatusUpdate() // to be enabled in resume
+ cp.controllerTimerService.disableRuntimeStatisticsCollection() // to be enabled in resume
+ Future
+ .collect(
+ cp.workflowExecution.getRunningRegionExecutions
+ .flatMap(_.getAllOperatorExecutions)
+ .map {
+ case (physicalOpId, opExecution) =>
+ // create a buffer for the current input tuple
+ // since we need to show them on the frontend
+ val buffer = mutable.ArrayBuffer[(Tuple, ActorVirtualIdentity)]()
+ Future
+ .collect(
+ opExecution.getWorkerIds
+ // send pause to all workers
+ // pause message has no effect on completed or paused workers
+ .map { worker =>
+ val workerExecution = opExecution.getWorkerExecution(worker)
+ // send a pause message
+ workerInterface.pauseWorker(EmptyRequest(), mkContext(worker)).flatMap {
+ resp =>
+ workerExecution.update(System.nanoTime(), resp.state)
+ workerInterface
+ .queryStatistics(EmptyRequest(), mkContext(worker))
+ // get the stats and current input tuple from the worker
+ .map {
+ case WorkerMetricsResponse(metrics) =>
+ workerExecution.update(System.nanoTime(), metrics.workerStatistics)
+ }
+ }
+ }.toSeq
+ )
+ }
+ .toSeq
+ )
+ .map { _ =>
+ // update frontend workflow status and persist statistics
+ val stats = cp.workflowExecution.getAllRegionExecutionsStats
+ sendToClient(ExecutionStatsUpdate(stats))
+ sendToClient(RuntimeStatisticsPersist(stats))
+ sendToClient(ExecutionStateUpdate(cp.workflowExecution.getState))
+ logger.info(s"workflow paused")
+ }
+ EmptyReturn()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PortCompletedHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PortCompletedHandler.scala
new file mode 100644
index 00000000000..810c098c417
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PortCompletedHandler.scala
@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.core.workflow.GlobalPortIdentity
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerAsyncRPCHandlerInitializer,
+ FatalError
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ PortCompletedRequest,
+ QueryStatisticsRequest,
+ StatisticsUpdateTarget
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+/** Notify the completion of a port:
+ * - For input port, it means the worker has finished consuming and processing all the data
+ * through this port, including all possible links to this port.
+ * - For output port, it means the worker has finished sending all the data through this port.
+ *
+ * possible sender: worker
+ */
+trait PortCompletedHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def portCompleted(
+ msg: PortCompletedRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ controllerInterface
+ .controllerInitiateQueryStatistics(
+ QueryStatisticsRequest(
+ scala.Seq(ctx.sender),
+ StatisticsUpdateTarget.BOTH_UI_AND_PERSISTENCE
+ ),
+ CONTROLLER
+ )
+ .map { _ =>
+ val globalPortId = GlobalPortIdentity(
+ VirtualIdentityUtils.getPhysicalOpId(ctx.sender),
+ msg.portId,
+ input = msg.input
+ )
+ cp.workflowExecutionCoordinator.getRegionOfPortId(globalPortId) match {
+ case Some(region) =>
+ val regionExecution = cp.workflowExecution.getRegionExecution(region.id)
+ val operatorExecution =
+ regionExecution.getOperatorExecution(VirtualIdentityUtils.getPhysicalOpId(ctx.sender))
+ val workerExecution = operatorExecution.getWorkerExecution(ctx.sender)
+
+ // set the port on this worker to be completed
+ (if (msg.input) workerExecution.getInputPortExecution(msg.portId)
+ else workerExecution.getOutputPortExecution(msg.portId)).setCompleted()
+
+ // check if the port on this operator is completed
+ val isPortCompleted =
+ if (msg.input) operatorExecution.isInputPortCompleted(msg.portId)
+ else operatorExecution.isOutputPortCompleted(msg.portId)
+
+ if (isPortCompleted) {
+ cp.workflowExecutionCoordinator
+ .coordinateRegionExecutors(cp.actorService)
+ // Since this message is sent from a worker, any exception from the above code will be returned to that worker.
+ // Additionally, a fatal error is sent to the client, indicating that the region cannot be scheduled.
+ .onFailure {
+ case err: WorkflowRuntimeException =>
+ sendToClient(FatalError(err, err.relatedWorkerId))
+ case other =>
+ sendToClient(FatalError(other, None))
+ }
+ }
+ case None => // currently "start" and "end" ports are not part of a region, thus no region can be found.
+ // do nothing.
+ }
+ EmptyReturn()
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/QueryWorkerStatisticsHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/QueryWorkerStatisticsHandler.scala
new file mode 100644
index 00000000000..6551579f719
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/QueryWorkerStatisticsHandler.scala
@@ -0,0 +1,191 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerAsyncRPCHandlerInitializer,
+ ExecutionStatsUpdate,
+ RuntimeStatisticsPersist
+}
+import org.apache.texera.amber.engine.architecture.deploysemantics.layer.WorkerExecution
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest,
+ QueryStatisticsRequest,
+ StatisticsUpdateTarget
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.COMPLETED
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ EmptyReturn,
+ WorkerMetricsResponse
+}
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+/** Get statistics from all the workers
+ *
+ * possible sender: controller(by statusUpdateAskHandle)
+ */
+trait QueryWorkerStatisticsHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ private var globalQueryStatsOngoing = false
+
+ // Minimum of the two timer intervals converted to nanoseconds.
+ // A full-graph worker query is skipped and served from cache when the last completed
+ // query falls within this window, avoiding redundant worker RPCs.
+ private val minQueryIntervalNs: Long =
+ Math.min(
+ ApplicationConfig.getStatusUpdateIntervalInMs,
+ ApplicationConfig.getRuntimeStatisticsPersistenceIntervalInMs
+ ) * 1_000_000L
+
+ // Nanosecond timestamp of the last completed full-graph worker stats query.
+ @volatile private var lastWorkerQueryTimestampNs: Long = 0L
+
+ // Reads the current cached stats and forwards them to the appropriate client sink(s).
+ private def forwardStats(updateTarget: StatisticsUpdateTarget): Unit = {
+ val stats = cp.workflowExecution.getAllRegionExecutionsStats
+ updateTarget match {
+ case StatisticsUpdateTarget.UI_ONLY =>
+ sendToClient(ExecutionStatsUpdate(stats))
+ case StatisticsUpdateTarget.PERSISTENCE_ONLY =>
+ sendToClient(RuntimeStatisticsPersist(stats))
+ case StatisticsUpdateTarget.BOTH_UI_AND_PERSISTENCE |
+ StatisticsUpdateTarget.Unrecognized(_) =>
+ sendToClient(ExecutionStatsUpdate(stats))
+ sendToClient(RuntimeStatisticsPersist(stats))
+ }
+ }
+
+ override def controllerInitiateQueryStatistics(
+ msg: QueryStatisticsRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ // Avoid issuing concurrent full-graph statistics queries.
+ // If a global query is already in progress, skip this request.
+ if (globalQueryStatsOngoing && msg.filterByWorkers.isEmpty) {
+ // A query is already in-flight: serve the last completed query's cached data,
+ // or drop silently if no prior query has finished yet.
+ if (lastWorkerQueryTimestampNs > 0) forwardStats(msg.updateTarget)
+ return EmptyReturn()
+ }
+
+ var opFilter: Set[PhysicalOpIdentity] = Set.empty
+ // Only enforce the single-query restriction for full-graph queries.
+ if (msg.filterByWorkers.isEmpty) {
+ if (System.nanoTime() - lastWorkerQueryTimestampNs < minQueryIntervalNs) {
+ // Cache is still fresh: the faster timer already queried workers recently.
+ forwardStats(msg.updateTarget)
+ return EmptyReturn()
+ }
+ globalQueryStatsOngoing = true
+ } else {
+ // Map the filtered worker IDs (if any) to their corresponding physical operator IDs
+ val initialOps: Set[PhysicalOpIdentity] =
+ msg.filterByWorkers.map(VirtualIdentityUtils.getPhysicalOpId).toSet
+
+ // Include all transitive upstream operators in the filter set
+ opFilter = {
+ val visited = scala.collection.mutable.Set.empty[PhysicalOpIdentity]
+ val toVisit = scala.collection.mutable.Queue.from(initialOps)
+
+ while (toVisit.nonEmpty) {
+ val current = toVisit.dequeue()
+ if (visited.add(current)) {
+ val upstreamOps = cp.workflowScheduler.physicalPlan.getUpstreamPhysicalOpIds(current)
+ toVisit.enqueueAll(upstreamOps)
+ }
+ }
+
+ visited.toSet
+ }
+ }
+
+ // Traverse the physical plan in reverse topological order (sink to source),
+ // grouped by layers of parallel operators.
+ val layers = cp.workflowScheduler.physicalPlan.layeredReversedTopologicalOrder
+
+ // Accumulator to collect all (exec, wid, state, stats) results
+ val collectedResults =
+ scala.collection.mutable.ArrayBuffer.empty[(WorkerExecution, WorkerMetricsResponse, Long)]
+
+ // Recursively process each operator layer sequentially (top-down in reverse topo order)
+ def processLayers(layers: Seq[Set[PhysicalOpIdentity]]): Future[Unit] =
+ layers match {
+ case Nil =>
+ // All layers have been processed
+ Future.Done
+
+ case layer +: rest =>
+ // Issue statistics queries to all eligible workers in the current layer
+ val futures = layer.toSeq.flatMap { opId =>
+ // Skip operators not included in the filtered subset (if any)
+ if (opFilter.nonEmpty && !opFilter.contains(opId)) {
+ Seq.empty
+ } else {
+ cp.workflowExecution.getLatestOperatorExecutionOption(opId) match {
+ // Operator region has not been initialized yet; skip in this polling round.
+ case None => Seq.empty
+ case Some(exec) =>
+ // Skip completed operators
+ if (exec.getState == COMPLETED) {
+ Seq.empty
+ } else {
+ // Select all workers for this operator
+ val workerIds = exec.getWorkerIds
+
+ // Send queryStatistics to each worker and update internal state on reply
+ workerIds.map { wid =>
+ workerInterface.queryStatistics(EmptyRequest(), wid).map { resp =>
+ collectedResults.addOne(
+ (exec.getWorkerExecution(wid), resp, System.nanoTime())
+ )
+ }
+ }
+ }
+ }
+ }
+ }
+
+ // After all worker queries in this layer complete, process the next layer
+ Future.collect(futures).flatMap(_ => processLayers(rest))
+ }
+
+ // Start processing all layers and forward stats to the appropriate sink(s) on completion.
+ processLayers(layers).map { _ =>
+ collectedResults.foreach {
+ case (wExec, resp, timestamp) =>
+ wExec.update(timestamp, resp.metrics.workerState, resp.metrics.workerStatistics)
+ }
+ forwardStats(msg.updateTarget)
+ // Record the completion timestamp before releasing the lock so that any timer
+ // firing in between sees a valid cache entry rather than triggering a redundant query.
+ if (globalQueryStatsOngoing) {
+ lastWorkerQueryTimestampNs = System.nanoTime()
+ globalQueryStatsOngoing = false
+ }
+ EmptyReturn()
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ReconfigurationHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ReconfigurationHandler.scala
new file mode 100644
index 00000000000..7653f873c13
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ReconfigurationHandler.scala
@@ -0,0 +1,141 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerAsyncRPCHandlerInitializer,
+ UpdateExecutorCompleted
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.ALL_ALIGNMENT
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ ControlInvocation,
+ WorkflowReconfigureRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ControlReturn, EmptyReturn}
+import org.apache.texera.amber.engine.common.FriesReconfigurationAlgorithm
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.util.VirtualIdentityUtils
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_UPDATE_EXECUTOR
+
+import scala.collection.mutable
+
+trait ReconfigurationHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def reconfigureWorkflow(
+ msg: WorkflowReconfigureRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ if (
+ msg.reconfiguration.exists(req =>
+ cp.workflowScheduler.physicalPlan.getOperator(req.targetOpId).isSourceOperator
+ )
+ ) {
+ throw new IllegalStateException(
+ "Reconfiguration cannot be applied to source operators"
+ )
+ }
+ val futures = mutable.ArrayBuffer[Future[_]]()
+ val friesComponents =
+ FriesReconfigurationAlgorithm.getReconfigurations(cp.workflowExecutionCoordinator, msg)
+ friesComponents.foreach { friesComponent =>
+ if (friesComponent.scope.size == 1) {
+ val updateExecutorRequest = friesComponent.reconfigurations.head
+ val workerIds = cp.workflowExecution
+ .getLatestOperatorExecution(updateExecutorRequest.targetOpId)
+ .getWorkerIds
+ workerIds.foreach { worker =>
+ futures.append(
+ notifyOnComplete(
+ workerInterface.updateExecutor(updateExecutorRequest, mkContext(worker)),
+ worker
+ )
+ )
+ }
+ } else {
+ val channelScope = cp.workflowExecution.getRunningRegionExecutions
+ .flatMap(regionExecution =>
+ regionExecution.getAllLinkExecutions
+ .map(_._2)
+ .flatMap(linkExecution => linkExecution.getAllChannelExecutions.map(_._1))
+ )
+ .filter(channelId => {
+ friesComponent.scope
+ .contains(VirtualIdentityUtils.getPhysicalOpId(channelId.fromWorkerId)) &&
+ friesComponent.scope
+ .contains(VirtualIdentityUtils.getPhysicalOpId(channelId.toWorkerId))
+ })
+ val controlChannels = friesComponent.sources.flatMap { source =>
+ cp.workflowExecution.getLatestOperatorExecution(source).getWorkerIds.flatMap { worker =>
+ Seq(
+ ChannelIdentity(CONTROLLER, worker, isControl = true),
+ ChannelIdentity(worker, CONTROLLER, isControl = true)
+ )
+ }
+ }
+ val finalScope = channelScope ++ controlChannels
+ val workerCommands: Seq[(ActorVirtualIdentity, ControlInvocation, Future[ControlReturn])] =
+ friesComponent.reconfigurations.flatMap { updateReq =>
+ val workers =
+ cp.workflowExecution.getLatestOperatorExecution(updateReq.targetOpId).getWorkerIds
+ workers.map { worker =>
+ val (invocation, future) =
+ createInvocation(METHOD_UPDATE_EXECUTOR.getBareMethodName, updateReq, worker)
+ (worker, invocation, future)
+ }
+ }.toSeq
+ val cmdMapping: Map[String, ControlInvocation] = workerCommands.map {
+ case (worker, invocation, _) => worker.name -> invocation
+ }.toMap
+ futures ++= workerCommands.map {
+ case (worker, _, future) => notifyOnComplete(future, worker)
+ }
+ friesComponent.sources.foreach { source =>
+ cp.workflowExecution.getLatestOperatorExecution(source).getWorkerIds.foreach { worker =>
+ sendECM(
+ EmbeddedControlMessageIdentity(msg.reconfigurationId),
+ ALL_ALIGNMENT,
+ finalScope.toSet,
+ cmdMapping,
+ ChannelIdentity(actorId, worker, isControl = true)
+ )
+ }
+ }
+ }
+ }
+ Future.collect(futures.toList).map { _ =>
+ EmptyReturn()
+ }
+ }
+
+ // After a worker's updateExecutor completes, notify the client so the
+ // ExecutionReconfigurationService can advance completedReconfigurations
+ // and emit ModifyLogicCompletedEvent on the websocket.
+ private def notifyOnComplete[T](future: Future[T], worker: ActorVirtualIdentity): Future[T] =
+ future.onSuccess(_ => sendToClient(UpdateExecutorCompleted(worker)))
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ResumeHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ResumeHandler.scala
new file mode 100644
index 00000000000..c94ba91c205
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/ResumeHandler.scala
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerAsyncRPCHandlerInitializer,
+ ExecutionStatsUpdate,
+ RuntimeStatisticsPersist
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+/** resume the entire workflow
+ *
+ * possible sender: controller, client
+ */
+trait ResumeHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def resumeWorkflow(msg: EmptyRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = {
+ // send all workers resume
+ // resume message has no effect on non-paused workers
+ Future
+ .collect(
+ cp.workflowExecution.getRunningRegionExecutions
+ .flatMap(_.getAllOperatorExecutions.map(_._2))
+ .flatMap(_.getWorkerIds)
+ .map { workerId =>
+ workerInterface.resumeWorker(EmptyRequest(), mkContext(workerId)).map { resp =>
+ cp.workflowExecution
+ .getLatestOperatorExecution(VirtualIdentityUtils.getPhysicalOpId(workerId))
+ .getWorkerExecution(workerId)
+ .update(System.nanoTime(), resp.state)
+ }
+ }
+ .toSeq
+ )
+ .map { _ =>
+ // update frontend status and persist statistics
+ val stats = cp.workflowExecution.getAllRegionExecutionsStats
+ sendToClient(ExecutionStatsUpdate(stats))
+ sendToClient(RuntimeStatisticsPersist(stats))
+ cp.controllerTimerService
+ .enableStatusUpdate() //re-enabled it since it is disabled in pause
+ cp.controllerTimerService
+ .enableRuntimeStatisticsCollection() //re-enabled it since it is disabled in pause
+ EmptyReturn()
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/RetrieveWorkflowStateHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/RetrieveWorkflowStateHandler.scala
new file mode 100644
index 00000000000..b6e75cd823b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/RetrieveWorkflowStateHandler.scala
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.EmbeddedControlMessageIdentity
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.NO_ALIGNMENT
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest,
+ PropagateEmbeddedControlMessageRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ RetrieveWorkflowStateResponse,
+ StringResponse
+}
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_RETRIEVE_STATE
+import org.apache.texera.amber.engine.common.virtualidentity.util.SELF
+
+import java.time.Instant
+
+trait RetrieveWorkflowStateHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def retrieveWorkflowState(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[RetrieveWorkflowStateResponse] = {
+ val targetOps = cp.workflowScheduler.physicalPlan.operators.map(_.id).toSeq
+ val msg = PropagateEmbeddedControlMessageRequest(
+ cp.workflowExecution.getRunningRegionExecutions
+ .flatMap(_.getAllOperatorExecutions.map(_._1))
+ .toSeq,
+ EmbeddedControlMessageIdentity("RetrieveWorkflowState_" + Instant.now().toString),
+ NO_ALIGNMENT,
+ targetOps,
+ targetOps,
+ EmptyRequest(),
+ METHOD_RETRIEVE_STATE.getBareMethodName
+ )
+ controllerInterface
+ .propagateEmbeddedControlMessage(
+ msg,
+ mkContext(SELF)
+ )
+ .map { ret =>
+ RetrieveWorkflowStateResponse(ret.returns.map {
+ case (actorId, value) =>
+ val finalret = value match {
+ case s: StringResponse => s.value
+ case other =>
+ ""
+ }
+ (actorId, finalret)
+ })
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/RetryWorkflowHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/RetryWorkflowHandler.scala
new file mode 100644
index 00000000000..a2e1e257f8a
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/RetryWorkflowHandler.scala
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest,
+ RetryWorkflowRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+
+/** retry the execution of the entire workflow
+ *
+ * possible sender: controller, client
+ */
+trait RetryWorkflowHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def retryWorkflow(
+ msg: RetryWorkflowRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ // if it is a PythonWorker, prepare for retry
+ // retry message has no effect on completed workers
+ Future
+ .collect(
+ msg.workers
+ .map(worker => workerInterface.retryCurrentTuple(EmptyRequest(), worker))
+ )
+ .unit
+
+ // resume all workers
+ controllerInterface.resumeWorkflow(EmptyRequest(), mkContext(CONTROLLER))
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/StartWorkflowHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/StartWorkflowHandler.scala
new file mode 100644
index 00000000000..7d938dbedde
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/StartWorkflowHandler.scala
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.StartWorkflowResponse
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.RUNNING
+
+/** start the workflow by starting the source workers
+ * note that this SHOULD only be called once per workflow
+ *
+ * possible sender: client
+ */
+trait StartWorkflowHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def startWorkflow(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[StartWorkflowResponse] = {
+ if (cp.workflowExecution.getState.isUninitialized) {
+ cp.workflowExecutionCoordinator
+ .coordinateRegionExecutors(cp.actorService)
+ .map(_ => {
+ cp.controllerTimerService.enableStatusUpdate()
+ cp.controllerTimerService.enableRuntimeStatisticsCollection()
+ StartWorkflowResponse(RUNNING)
+ })
+ } else {
+ StartWorkflowResponse(cp.workflowExecution.getState)
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/TakeGlobalCheckpointHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/TakeGlobalCheckpointHandler.scala
new file mode 100644
index 00000000000..128c00fc562
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/TakeGlobalCheckpointHandler.scala
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.NO_ALIGNMENT
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.TakeGlobalCheckpointResponse
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_PREPARE_CHECKPOINT
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.engine.common.virtualidentity.util.SELF
+import org.apache.texera.amber.engine.common.{CheckpointState, SerializedState}
+
+import java.net.URI
+
+trait TakeGlobalCheckpointHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def takeGlobalCheckpoint(
+ msg: TakeGlobalCheckpointRequest,
+ ctx: AsyncRPCContext
+ ): Future[TakeGlobalCheckpointResponse] = {
+ var estimationOnly = msg.estimationOnly
+ val destinationURI = new URI(msg.destination)
+ @transient val storage =
+ SequentialRecordStorage.getStorage[CheckpointState](Some(destinationURI))
+ if (storage.containsFolder(msg.checkpointId.toString)) {
+ logger.info("skip checkpoint since its already taken")
+ estimationOnly = true
+ }
+ val uri = destinationURI.resolve(msg.checkpointId.toString)
+ var totalSize = 0L
+ val physicalOpIdsToTakeCheckpoint = cp.workflowScheduler.physicalPlan.operators.map(_.id)
+ controllerInterface
+ .propagateEmbeddedControlMessage(
+ PropagateEmbeddedControlMessageRequest(
+ cp.workflowExecution.getAllRegionExecutions
+ .flatMap(_.getAllOperatorExecutions.map(_._1))
+ .toSeq,
+ msg.checkpointId,
+ NO_ALIGNMENT,
+ cp.workflowScheduler.physicalPlan.operators.map(_.id).toSeq,
+ physicalOpIdsToTakeCheckpoint.toSeq,
+ PrepareCheckpointRequest(msg.checkpointId, estimationOnly),
+ METHOD_PREPARE_CHECKPOINT.getBareMethodName
+ ),
+ mkContext(SELF)
+ )
+ .flatMap { ret =>
+ Future
+ .collect(ret.returns.map {
+ case (workerId, _) =>
+ val destActor = ActorVirtualIdentity(workerId)
+ workerInterface
+ .finalizeCheckpoint(
+ FinalizeCheckpointRequest(msg.checkpointId, uri.toString),
+ mkContext(destActor)
+ )
+ .onSuccess { resp =>
+ totalSize += resp.size
+ }
+ .onFailure { err =>
+ throw err // TODO: handle failures.
+ }
+ }.toSeq)
+ .map { _ =>
+ logger.info("Start to take checkpoint")
+ val chkpt = new CheckpointState()
+ if (!estimationOnly) {
+ // serialize CP state
+ chkpt.save(SerializedState.CP_STATE_KEY, this.cp)
+ logger.info(
+ s"Serialized CP state, current workflow state = ${cp.workflowExecution.getState}"
+ )
+ // get all output messages from cp.transferService
+ chkpt.save(
+ SerializedState.OUTPUT_MSG_KEY,
+ this.cp.transferService.getAllUnAckedMessages.toArray
+ )
+ val storage = SequentialRecordStorage.getStorage[CheckpointState](Some(uri))
+ val writer = storage.getWriter(actorId.name)
+ writer.writeRecord(chkpt)
+ writer.flush()
+ writer.close()
+ }
+ totalSize += chkpt.size()
+ logger.info(s"global checkpoint finalized, total size = $totalSize")
+ TakeGlobalCheckpointResponse(totalSize)
+ }
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/WorkerExecutionCompletedHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/WorkerExecutionCompletedHandler.scala
new file mode 100644
index 00000000000..c3b3ddb234b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/WorkerExecutionCompletedHandler.scala
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerAsyncRPCHandlerInitializer,
+ ExecutionStateUpdate
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest,
+ QueryStatisticsRequest,
+ StatisticsUpdateTarget
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.common.virtualidentity.util.SELF
+
+/** indicate a worker has completed its execution
+ * i.e. received and processed all data from upstreams
+ * note that this doesn't mean all the output of this worker
+ * has been received by the downstream workers.
+ *
+ * possible sender: worker
+ */
+trait WorkerExecutionCompletedHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def workerExecutionCompleted(
+ msg: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+
+ // after worker execution is completed, query statistics immediately one last time
+ // because the worker might be killed before the next query statistics interval
+ // and the user sees the last update before completion
+ val statsRequest =
+ controllerInterface.controllerInitiateQueryStatistics(
+ QueryStatisticsRequest(Seq(ctx.sender), StatisticsUpdateTarget.BOTH_UI_AND_PERSISTENCE),
+ mkContext(SELF)
+ )
+
+ Future
+ .collect(Seq(statsRequest))
+ .flatMap(_ => {
+ // if entire workflow is completed, clean up
+ val isWorkflowTerminal =
+ cp.workflowExecution.isCompleted &&
+ !cp.workflowScheduler.hasPendingRegions &&
+ !cp.workflowExecutionCoordinator.hasUnfinishedRegionCoordinators
+ if (isWorkflowTerminal) {
+ // after query result come back: send completed event, cleanup ,and kill workflow
+ sendToClient(ExecutionStateUpdate(cp.workflowExecution.getState))
+ cp.controllerTimerService.disableStatusUpdate()
+ cp.controllerTimerService.disableRuntimeStatisticsCollection()
+ }
+ })
+ EmptyReturn()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/WorkerStateUpdatedHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/WorkerStateUpdatedHandler.scala
new file mode 100644
index 00000000000..5ee98a4918d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/WorkerStateUpdatedHandler.scala
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerAsyncRPCHandlerInitializer,
+ ExecutionStatsUpdate,
+ RuntimeStatisticsPersist
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ WorkerStateUpdatedRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+/** indicate the state change of a worker
+ *
+ * possible sender: worker
+ */
+trait WorkerStateUpdatedHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def workerStateUpdated(
+ msg: WorkerStateUpdatedRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ val physicalOpId = VirtualIdentityUtils.getPhysicalOpId(ctx.sender)
+ // set the state
+ cp.workflowExecution.getRunningRegionExecutions
+ .find(_.hasOperatorExecution(physicalOpId))
+ .map(_.getOperatorExecution(physicalOpId))
+ .foreach(operatorExecution =>
+ operatorExecution.getWorkerExecution(ctx.sender).update(System.nanoTime(), msg.state)
+ )
+ val stats = cp.workflowExecution.getAllRegionExecutionsStats
+ sendToClient(ExecutionStatsUpdate(stats))
+ sendToClient(RuntimeStatisticsPersist(stats))
+ EmptyReturn()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/AddressInfo.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/AddressInfo.scala
new file mode 100644
index 00000000000..f23e54d736d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/AddressInfo.scala
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.deploysemantics
+
+import org.apache.pekko.actor.Address
+
+// Holds worker and controller node addresses.
+case class AddressInfo(
+ allAddresses: Array[Address], // e.g., Node 1, Node 2, Node 3
+ controllerAddress: Address // Controller node
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/DeployStrategy.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/DeployStrategy.scala
new file mode 100644
index 00000000000..079d253c848
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/DeployStrategy.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.deploysemantics.deploystrategy
+
+import org.apache.pekko.actor.Address
+
+trait DeployStrategy extends Serializable {
+
+ def initialize(available: Array[Address]): Unit
+
+ def next(): Address
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/OneOnEach.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/OneOnEach.scala
new file mode 100644
index 00000000000..62cf288263b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/OneOnEach.scala
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.deploysemantics.deploystrategy
+
+import org.apache.pekko.actor.Address
+
+object OneOnEach {
+ def apply() = new OneOnEach()
+}
+
+class OneOnEach extends DeployStrategy {
+ var available: Array[Address] = _
+ var index = 0
+
+ override def initialize(available: Array[Address]): Unit = {
+ this.available = available
+ }
+
+ override def next(): Address = {
+ val i = index
+ if (i >= available.length) {
+ throw new IndexOutOfBoundsException()
+ }
+ index += 1
+ available(i)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/RandomDeployment.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/RandomDeployment.scala
new file mode 100644
index 00000000000..aebab32fca9
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/RandomDeployment.scala
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.deploysemantics.deploystrategy
+
+import org.apache.pekko.actor.Address
+
+object RandomDeployment {
+ def apply() = new RandomDeployment()
+}
+
+class RandomDeployment extends DeployStrategy {
+ var available: Array[Address] = _
+
+ override def initialize(available: Array[Address]): Unit = {
+ this.available = available
+ }
+
+ override def next(): Address = {
+ available(util.Random.nextInt(available.length))
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/RoundRobinDeployment.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/RoundRobinDeployment.scala
new file mode 100644
index 00000000000..3fee912d7d9
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/deploystrategy/RoundRobinDeployment.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.deploysemantics.deploystrategy
+
+import org.apache.pekko.actor.Address
+
+object RoundRobinDeployment {
+ def apply() = new RoundRobinDeployment()
+}
+
+class RoundRobinDeployment extends DeployStrategy {
+ var available: Array[Address] = _
+ var index = 0
+
+ override def initialize(available: Array[Address]): Unit = {
+ this.available = available
+ }
+
+ override def next(): Address = {
+ val i = index
+ index = (index + 1) % available.length
+ available(i)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/layer/WorkerExecution.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/layer/WorkerExecution.scala
new file mode 100644
index 00000000000..55e1e309181
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/deploysemantics/layer/WorkerExecution.scala
@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.deploysemantics.layer
+
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.controller.execution.WorkerPortExecution
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.UNINITIALIZED
+import org.apache.texera.amber.engine.architecture.worker.statistics.{WorkerState, WorkerStatistics}
+
+import scala.collection.mutable
+
+case class WorkerExecution() extends Serializable {
+
+ private val inputPortExecutions: mutable.HashMap[PortIdentity, WorkerPortExecution] =
+ mutable.HashMap()
+ private val outputPortExecutions: mutable.HashMap[PortIdentity, WorkerPortExecution] =
+ mutable.HashMap()
+
+ private var state: WorkerState = UNINITIALIZED
+ private var stats: WorkerStatistics = {
+ WorkerStatistics(Seq.empty, Seq.empty, 0, 0, 0)
+ }
+ private var lastUpdateTimeStamp = 0L
+
+ /**
+ * Updates both the worker state and statistics if the provided timestamp is newer
+ * than the last recorded update timestamp. This ensures that only the most recent
+ * data is reflected in the execution state.
+ *
+ * @param timeStamp the nanosecond-timestamp of this update
+ * @param state the new WorkerState to set
+ * @param stats the new WorkerStatistics to set
+ */
+ def update(timeStamp: Long, state: WorkerState, stats: WorkerStatistics): Unit = {
+ if (this.lastUpdateTimeStamp < timeStamp) {
+ this.stats = stats
+ this.state = state
+ this.lastUpdateTimeStamp = timeStamp
+ }
+ }
+
+ /**
+ * Updates only the worker state if the provided timestamp is newer than the
+ * last recorded update timestamp.
+ *
+ * @param timeStamp the nanosecond-timestamp of this update
+ * @param state the new WorkerState to set
+ */
+ def update(timeStamp: Long, state: WorkerState): Unit = {
+ if (this.lastUpdateTimeStamp < timeStamp) {
+ this.state = state
+ this.lastUpdateTimeStamp = timeStamp
+ }
+ }
+
+ /**
+ * Updates only the worker statistics if the provided timestamp is newer than the
+ * last recorded update timestamp.
+ *
+ * @param timeStamp the nanosecond-timestamp of this update
+ * @param stats the new WorkerStatistics to set
+ */
+ def update(timeStamp: Long, stats: WorkerStatistics): Unit = {
+ if (this.lastUpdateTimeStamp < timeStamp) {
+ this.stats = stats
+ this.lastUpdateTimeStamp = timeStamp
+ }
+ }
+
+ def getState: WorkerState = state
+
+ def getStats: WorkerStatistics = stats
+
+ def getInputPortExecution(portId: PortIdentity): WorkerPortExecution = {
+ if (!inputPortExecutions.contains(portId)) {
+ inputPortExecutions(portId) = new WorkerPortExecution()
+ }
+ inputPortExecutions(portId)
+
+ }
+
+ def getOutputPortExecution(portId: PortIdentity): WorkerPortExecution = {
+ if (!outputPortExecutions.contains(portId)) {
+ outputPortExecutions(portId) = new WorkerPortExecution()
+ }
+ outputPortExecutions(portId)
+
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/AsyncReplayLogWriter.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/AsyncReplayLogWriter.scala
new file mode 100644
index 00000000000..39969340d39
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/AsyncReplayLogWriter.scala
@@ -0,0 +1,111 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.logreplay
+
+import com.google.common.collect.Queues
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage.SequentialRecordWriter
+
+import java.util
+import java.util.concurrent.CompletableFuture
+import scala.collection.mutable.ListBuffer
+import scala.jdk.CollectionConverters.ListHasAsScala
+
+class AsyncReplayLogWriter(
+ handler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit,
+ writer: SequentialRecordWriter[ReplayLogRecord]
+) extends Thread {
+ private val drained =
+ new util.ArrayList[
+ Either[ReplayLogRecord, Either[MainThreadDelegateMessage, WorkflowFIFOMessage]]
+ ]()
+ private val writerQueue =
+ Queues.newLinkedBlockingQueue[
+ Either[ReplayLogRecord, Either[MainThreadDelegateMessage, WorkflowFIFOMessage]]
+ ]()
+ private var stopped = false
+ private val logInterval =
+ ApplicationConfig.faultToleranceLogFlushIntervalInMs
+ private val gracefullyStopped = new CompletableFuture[Unit]()
+
+ def putLogRecords(records: Array[ReplayLogRecord]): Unit = {
+ assert(!stopped)
+ records.foreach(x => {
+ writerQueue.put(Left(x))
+ })
+ }
+
+ def putOutput(output: Either[MainThreadDelegateMessage, WorkflowFIFOMessage]): Unit = {
+ assert(!stopped)
+ writerQueue.put(Right(output))
+ }
+
+ def terminate(): Unit = {
+ stopped = true
+ writerQueue.put(Left(TerminateSignal))
+ gracefullyStopped.get()
+ }
+
+ override def run(): Unit = {
+ var internalStop = false
+ while (!internalStop) {
+ if (logInterval > 0) {
+ Thread.sleep(logInterval)
+ }
+ internalStop = drainWriterQueueAndProcess()
+ }
+ writer.close()
+ gracefullyStopped.complete(())
+ }
+
+ private def drainWriterQueueAndProcess(): Boolean = {
+ var stop = false
+ if (writerQueue.drainTo(drained) == 0) {
+ drained.add(writerQueue.take())
+ }
+ var drainedScala = drained.asScala
+ if (drainedScala.last == Left(TerminateSignal)) {
+ drainedScala = drainedScala.dropRight(1)
+ stop = true
+ }
+
+ val (replayLogRecords, workflowFIFOMessages) =
+ drainedScala.foldLeft(
+ (
+ ListBuffer[ReplayLogRecord](),
+ ListBuffer[Either[MainThreadDelegateMessage, WorkflowFIFOMessage]]()
+ )
+ ) {
+ case ((accLogs, accMsgs), Left(logRecord)) => (accLogs += logRecord, accMsgs)
+ case ((accLogs, accMsgs), Right(fifoMessage)) => (accLogs, accMsgs += fifoMessage)
+ }
+ // write logs first
+ replayLogRecords.foreach(replayLogRecord => writer.writeRecord(replayLogRecord))
+ writer.flush()
+ // send messages after logs are written
+ workflowFIFOMessages.foreach(workflowFIFOMessage => handler(workflowFIFOMessage))
+
+ drained.clear()
+ stop
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/EmptyReplayLogger.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/EmptyReplayLogger.scala
new file mode 100644
index 00000000000..74be96acdbd
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/EmptyReplayLogger.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.logreplay
+
+import org.apache.texera.amber.core.virtualidentity.{
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+
+class EmptyReplayLogger extends ReplayLogger {
+
+ override def drainCurrentLogRecords(step: Long): Array[ReplayLogRecord] = {
+ Array.empty
+ }
+
+ def markAsReplayDestination(id: EmbeddedControlMessageIdentity): Unit = {}
+
+ override def logCurrentStepWithMessage(
+ step: Long,
+ channelId: ChannelIdentity,
+ msg: Option[WorkflowFIFOMessage]
+ ): Unit = {}
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/OrderEnforcer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/OrderEnforcer.scala
new file mode 100644
index 00000000000..b18151992fd
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/OrderEnforcer.scala
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.logreplay
+
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+
+trait OrderEnforcer {
+ var isCompleted: Boolean
+
+ def canProceed(channelId: ChannelIdentity): Boolean
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogGenerator.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogGenerator.scala
new file mode 100644
index 00000000000..d159263db33
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogGenerator.scala
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.logreplay
+
+import org.apache.texera.amber.core.virtualidentity.EmbeddedControlMessageIdentity
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+
+import scala.collection.mutable
+
+object ReplayLogGenerator {
+ def generate(
+ logStorage: SequentialRecordStorage[ReplayLogRecord],
+ logFileName: String,
+ replayTo: EmbeddedControlMessageIdentity
+ ): (mutable.Queue[ProcessingStep], mutable.Queue[WorkflowFIFOMessage]) = {
+ val logs = logStorage.getReader(logFileName).mkRecordIterator()
+ val steps = mutable.Queue[ProcessingStep]()
+ val messages = mutable.Queue[WorkflowFIFOMessage]()
+ logs.foreach {
+ case s: ProcessingStep =>
+ steps.enqueue(s)
+ case MessageContent(message) =>
+ messages.enqueue(message)
+ case ReplayDestination(id) =>
+ if (id == replayTo) {
+ // we only need log record upto this point
+ return (steps, messages)
+ }
+ case other =>
+ throw new RuntimeException(s"cannot handle $other in the log")
+ }
+ (steps, messages)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogManager.scala
new file mode 100644
index 00000000000..a894ab3c8ef
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogManager.scala
@@ -0,0 +1,139 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.logreplay
+
+import org.apache.texera.amber.core.virtualidentity.{
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.architecture.common.ProcessingStepCursor
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage.SequentialRecordWriter
+import org.apache.texera.amber.engine.common.storage.{EmptyRecordStorage, SequentialRecordStorage}
+
+//In-mem formats:
+sealed trait ReplayLogRecord extends Serializable
+
+case class MessageContent(message: WorkflowFIFOMessage) extends ReplayLogRecord
+
+case class ProcessingStep(channelId: ChannelIdentity, step: Long) extends ReplayLogRecord
+
+case class ReplayDestination(id: EmbeddedControlMessageIdentity) extends ReplayLogRecord
+
+case object TerminateSignal extends ReplayLogRecord
+
+object ReplayLogManager {
+ def createLogManager(
+ logStorage: SequentialRecordStorage[ReplayLogRecord],
+ logFileName: String,
+ handler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit
+ ): ReplayLogManager = {
+ logStorage match {
+ case _: EmptyRecordStorage[ReplayLogRecord] =>
+ new EmptyReplayLogManagerImpl(handler)
+ case other =>
+ val manager = new ReplayLogManagerImpl(handler)
+ manager.setupWriter(other.getWriter(logFileName))
+ manager
+ }
+ }
+}
+
+trait ReplayLogManager {
+
+ protected val cursor = new ProcessingStepCursor()
+
+ def setupWriter(logWriter: SequentialRecordWriter[ReplayLogRecord]): Unit
+
+ def sendCommitted(msg: Either[MainThreadDelegateMessage, WorkflowFIFOMessage]): Unit
+
+ def terminate(): Unit
+
+ def getStep: Long = cursor.getStep
+
+ def markAsReplayDestination(id: EmbeddedControlMessageIdentity): Unit
+
+ def withFaultTolerant(
+ channelId: ChannelIdentity,
+ message: Option[WorkflowFIFOMessage]
+ )(code: => Unit): Unit = {
+ cursor.setCurrentChannel(channelId)
+ try {
+ code
+ } catch {
+ case t: Throwable => throw t
+ } finally {
+ cursor.stepIncrement()
+ }
+ }
+
+}
+
+class EmptyReplayLogManagerImpl(
+ handler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit
+) extends ReplayLogManager {
+ override def setupWriter(
+ logWriter: SequentialRecordStorage.SequentialRecordWriter[ReplayLogRecord]
+ ): Unit = {}
+
+ override def sendCommitted(msg: Either[MainThreadDelegateMessage, WorkflowFIFOMessage]): Unit = {
+ handler(msg)
+ }
+
+ override def terminate(): Unit = {}
+
+ override def markAsReplayDestination(id: EmbeddedControlMessageIdentity): Unit = {}
+}
+
+class ReplayLogManagerImpl(handler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit)
+ extends ReplayLogManager {
+
+ private val replayLogger = new ReplayLoggerImpl()
+
+ private var writer: AsyncReplayLogWriter = _
+
+ override def withFaultTolerant(
+ channelId: ChannelIdentity,
+ message: Option[WorkflowFIFOMessage]
+ )(code: => Unit): Unit = {
+ replayLogger.logCurrentStepWithMessage(cursor.getStep, channelId, message)
+ super.withFaultTolerant(channelId, message)(code)
+ }
+
+ override def markAsReplayDestination(id: EmbeddedControlMessageIdentity): Unit = {
+ replayLogger.markAsReplayDestination(id)
+ }
+
+ override def setupWriter(logWriter: SequentialRecordWriter[ReplayLogRecord]): Unit = {
+ writer = new AsyncReplayLogWriter(handler, logWriter)
+ writer.start()
+ }
+
+ override def sendCommitted(msg: Either[MainThreadDelegateMessage, WorkflowFIFOMessage]): Unit = {
+ writer.putLogRecords(replayLogger.drainCurrentLogRecords(cursor.getStep))
+ writer.putOutput(msg)
+ }
+
+ override def terminate(): Unit = {
+ writer.terminate()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogger.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogger.scala
new file mode 100644
index 00000000000..9fbd8bf1a0b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLogger.scala
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.logreplay
+
+import org.apache.texera.amber.core.virtualidentity.{
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+
+abstract class ReplayLogger {
+
+ def logCurrentStepWithMessage(
+ step: Long,
+ channelId: ChannelIdentity,
+ msg: Option[WorkflowFIFOMessage]
+ ): Unit
+
+ def markAsReplayDestination(id: EmbeddedControlMessageIdentity): Unit
+
+ def drainCurrentLogRecords(step: Long): Array[ReplayLogRecord]
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLoggerImpl.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLoggerImpl.scala
new file mode 100644
index 00000000000..42f5b9e2069
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayLoggerImpl.scala
@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.logreplay
+
+import org.apache.texera.amber.core.virtualidentity.{
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.architecture.common.ProcessingStepCursor.INIT_STEP
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+
+import scala.collection.mutable
+
+class ReplayLoggerImpl extends ReplayLogger {
+
+ private val tempLogs = mutable.ArrayBuffer[ReplayLogRecord]()
+
+ private var currentChannelId: ChannelIdentity = _
+
+ private var lastStep = INIT_STEP
+
+ /**
+ * Records the current processing step along with an associated message.
+ * This method also monitors the channel information. If the new channel matches the last recorded channel
+ * and there is no associated message for this step, the logging operation is bypassed.
+ * Otherwise, it appends a ProcessingStep log record with the message content, provided the message exists.
+ *
+ * @param step The current processing step.
+ * @param channel The channel ID associated with the processing step.
+ * @param message An optional message associated with the processing step.
+ */
+ override def logCurrentStepWithMessage(
+ step: Long,
+ channelId: ChannelIdentity,
+ message: Option[WorkflowFIFOMessage]
+ ): Unit = {
+ if (currentChannelId == channelId && message.isEmpty) {
+ return
+ }
+ currentChannelId = channelId
+ lastStep = step
+ tempLogs.append(ProcessingStep(channelId, step))
+ if (message.isDefined) {
+ tempLogs.append(MessageContent(message.get))
+ }
+ }
+
+ /**
+ * Called when the data processor attempts to output a message.
+ * This method retrieves all accumulated log records and passes them to the writer thread for persistence.
+ * It ensures the processing up to the current processing step is captured in the log records.
+ *
+ * @param step The current processing step.
+ * @return An array of ReplayLogRecord containing all the log records up to the current step.
+ */
+ def drainCurrentLogRecords(step: Long): Array[ReplayLogRecord] = {
+ if (lastStep != step) {
+ lastStep = step
+ tempLogs.append(ProcessingStep(currentChannelId, step))
+ }
+ val result = tempLogs.toArray
+ tempLogs.clear()
+ result
+ }
+
+ def markAsReplayDestination(id: EmbeddedControlMessageIdentity): Unit = {
+ tempLogs.append(ReplayDestination(id))
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayOrderEnforcer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayOrderEnforcer.scala
new file mode 100644
index 00000000000..6441aae2c20
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/logreplay/ReplayOrderEnforcer.scala
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.logreplay
+
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+
+import scala.collection.mutable
+
+class ReplayOrderEnforcer(
+ logManager: ReplayLogManager,
+ channelStepOrder: mutable.Queue[ProcessingStep],
+ startStep: Long,
+ private var onComplete: () => Unit
+) extends OrderEnforcer {
+ private var currentChannelId: ChannelIdentity = _
+
+ private def triggerOnComplete(): Unit = {
+ if (!isCompleted) {
+ return
+ }
+ if (onComplete != null) {
+ onComplete()
+ onComplete = null // make sure the onComplete is called only once.
+ }
+ }
+
+ // restore replay progress by dropping some of the entries
+ while (channelStepOrder.nonEmpty && channelStepOrder.head.step <= startStep) {
+ forwardNext()
+ }
+
+ var isCompleted: Boolean = channelStepOrder.isEmpty
+
+ triggerOnComplete()
+
+ private def forwardNext(): Unit = {
+ if (channelStepOrder.nonEmpty) {
+ val nextStep = channelStepOrder.dequeue()
+ currentChannelId = nextStep.channelId
+ }
+ }
+
+ def canProceed(channelId: ChannelIdentity): Boolean = {
+ val step = logManager.getStep
+ // release the next log record if the step matches
+ // Note: To remove duplicate step orders caused by checkpoints
+ // sending out a MainThreadDelegateMessage, we use a while loop.
+ while (channelStepOrder.nonEmpty && channelStepOrder.head.step == step) {
+ forwardNext()
+ }
+ // To terminate replay:
+ // no next log record with step > current step, which means further processing is not logged.
+ if (channelStepOrder.isEmpty) {
+ isCompleted = true
+ triggerOnComplete()
+ }
+ // only proceed if the current channel ID matches the channel ID of the log record
+ currentChannelId == channelId
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/AmberFIFOChannel.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/AmberFIFOChannel.scala
new file mode 100644
index 00000000000..d81b4239ba7
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/AmberFIFOChannel.scala
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
+
+import java.util.concurrent.atomic.AtomicLong
+import scala.collection.mutable
+
+/* The abstracted FIFO/exactly-once logic */
+class AmberFIFOChannel(val channelId: ChannelIdentity) extends AmberLogging {
+
+ override def actorId: ActorVirtualIdentity = channelId.toWorkerId
+
+ private val ofoMap = new mutable.HashMap[Long, WorkflowFIFOMessage]
+ private var current = 0L
+ private var enabled = true
+ private val fifoQueue = new mutable.ListBuffer[WorkflowFIFOMessage]
+ private val holdCredit = new AtomicLong()
+ private var portId: Option[PortIdentity] = None
+
+ def acceptMessage(msg: WorkflowFIFOMessage): Unit = {
+ val seq = msg.sequenceNumber
+ val payload = msg.payload
+ if (isDuplicated(seq)) {
+ logger.debug(
+ s"received duplicated message $payload with seq = $seq while current seq = $current"
+ )
+ } else if (isAhead(seq)) {
+ logger.debug(s"received ahead message $payload with seq = $seq while current seq = $current")
+ stash(seq, msg)
+ } else {
+ enforceFIFO(msg)
+ }
+ }
+
+ def getCurrentSeq: Long = current
+
+ private def isDuplicated(sequenceNumber: Long): Boolean =
+ sequenceNumber < current || ofoMap.contains(sequenceNumber)
+
+ private def isAhead(sequenceNumber: Long): Boolean = sequenceNumber > current
+
+ private def stash(sequenceNumber: Long, data: WorkflowFIFOMessage): Unit = {
+ ofoMap(sequenceNumber) = data
+ }
+
+ private def enforceFIFO(data: WorkflowFIFOMessage): Unit = {
+ fifoQueue.append(data)
+ holdCredit.getAndAdd(getInMemSize(data))
+ current += 1
+ while (ofoMap.contains(current)) {
+ val msg = ofoMap(current)
+ fifoQueue.append(msg)
+ holdCredit.getAndAdd(getInMemSize(msg))
+ ofoMap.remove(current)
+ current += 1
+ }
+ }
+
+ def take: WorkflowFIFOMessage = {
+ val msg = fifoQueue.remove(0)
+ holdCredit.getAndAdd(-getInMemSize(msg))
+ msg
+ }
+
+ def hasMessage: Boolean = fifoQueue.nonEmpty
+
+ def enable(isEnabled: Boolean): Unit = {
+ this.enabled = isEnabled
+ }
+
+ def isEnabled: Boolean = enabled
+
+ def getTotalMessageSize: Long = {
+ if (fifoQueue.nonEmpty) {
+ fifoQueue.map(getInMemSize(_)).sum
+ } else {
+ 0
+ }
+ }
+
+ def getTotalStashedSize: Long =
+ if (ofoMap.nonEmpty) {
+ ofoMap.values.map(getInMemSize(_)).sum
+ } else {
+ 0
+ }
+
+ def getQueuedCredit: Long = {
+ holdCredit.get()
+ }
+
+ def setPortId(portId: PortIdentity): Unit = {
+ this.portId = Some(portId)
+ }
+
+ def getPortId: PortIdentity = {
+ this.portId.get
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/CongestionControl.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/CongestionControl.scala
new file mode 100644
index 00000000000..337be75cf19
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/CongestionControl.scala
@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkMessage
+
+import scala.collection.mutable
+
+class CongestionControl {
+
+ // ack should be received within 3s,
+ // otherwise, network congestion occurs.
+ final val ackTimeLimit = 3000
+
+ // if the ack for a message is not received after 60s,
+ // we trigger the resending logic.
+ // Note that the resend is not guaranteed to happen
+ // after sending the message for 60s
+ final val resendTimeLimit = 60000 // 60s
+
+ // slow start threshold
+ // after windowSize > ssThreshold,
+ // we increment windowSize by one every time,
+ // otherwise, we multiply windowSize by 2.
+ private var ssThreshold = 16
+
+ // initial window size = 1
+ // it represents how many messages can be sent concurrently
+ private var windowSize = 1
+
+ private val toBeSent = new mutable.Queue[NetworkMessage]
+ private val inTransit = new mutable.LongMap[NetworkMessage]()
+ private val sentTime = new mutable.LongMap[Long]()
+
+ // Note that toBeSent buffer is always empty if inTransit.size < windowSize
+ def canSend: Boolean = inTransit.size < windowSize
+
+ def enqueueMessage(data: NetworkMessage): Unit = {
+ toBeSent.enqueue(data)
+ }
+
+ def ack(id: Long): Unit = {
+ if (!inTransit.contains(id)) return
+ inTransit.remove(id)
+ if (System.currentTimeMillis() - sentTime(id) < ackTimeLimit) {
+ if (windowSize < ssThreshold) {
+ windowSize = Math.min(windowSize * 2, ssThreshold)
+ } else {
+ windowSize += 1
+ }
+ } else {
+ ssThreshold /= 2
+ if (ssThreshold < 1) {
+ ssThreshold = 1
+ }
+ windowSize = ssThreshold
+ }
+ sentTime.remove(id)
+ }
+
+ private def dequeueN[T](queue: mutable.Queue[T], n: Int): Seq[T] = {
+ var count = 0
+ queue.dequeueAll { _ =>
+ count += 1
+ count <= n
+ }
+ }
+
+ def getBufferedMessagesToSend: Iterable[NetworkMessage] = {
+ val count = windowSize - inTransit.size
+ dequeueN(toBeSent, count)
+ }
+
+ def markMessageInTransit(data: NetworkMessage): Unit = {
+ inTransit(data.messageId) = data
+ sentTime(data.messageId) = System.currentTimeMillis()
+ }
+
+ def getTimedOutInTransitMessages: Iterable[NetworkMessage] = {
+ val timeCap = System.currentTimeMillis() - resendTimeLimit
+ sentTime.collect {
+ case (id, timeStamp) if timeStamp < timeCap =>
+ inTransit(id)
+ }
+ }
+
+ def getInTransitMessages: Iterable[NetworkMessage] = {
+ inTransit.values
+ }
+
+ def getAllMessages: Iterable[NetworkMessage] = {
+ val intransitMsg = inTransit.values
+ val toBeSentMsg = toBeSent
+ intransitMsg.toSet.union(toBeSentMsg.toSet)
+ }
+
+ def getStatusReport: String = {
+ s"current window size = ${windowSize} \t in transit = ${inTransit.size} \t waiting = ${toBeSent.size}"
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/DeadLetterMonitorActor.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/DeadLetterMonitorActor.scala
new file mode 100644
index 00000000000..323761fca11
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/DeadLetterMonitorActor.scala
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.pekko.actor.{Actor, DeadLetter}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.{
+ MessageBecomesDeadLetter,
+ NetworkMessage
+}
+
+class DeadLetterMonitorActor extends Actor {
+ override def receive: Receive = {
+ case d: DeadLetter =>
+ d.message match {
+ case msg: NetworkMessage =>
+ // d.sender is the NetworkSenderActor
+ d.sender ! MessageBecomesDeadLetter(msg)
+ case other =>
+ // skip for now
+ }
+ case _ =>
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/FlowControl.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/FlowControl.scala
new file mode 100644
index 00000000000..d4b24dad1d8
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/FlowControl.scala
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkMessage
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
+
+import scala.collection.mutable
+import scala.util.control.Breaks.{break, breakable}
+
+/**
+ * We implement credit-based flow control. Suppose a sender worker S sends data in batches to a receiving worker R
+ * using the network communicator actor NC. The different parts of flow control work as follows:
+ *
+ * 1. A worker has a fixed amount of credits for each of its sender workers. When R is expensive, its internal queue
+ * starts getting filled with data from S. This leads to a decrease in credits available for S.
+ *
+ * 2. R sends the credits available in the NetworkAck() being sent to S. This includes acks for data messages and
+ * control messages. The responsibility to decrease data sending lies on S now.
+ *
+ * 3. Upon receiving NetworkAck(), S saves the credit information and then does two things:
+ * a) Tell parent to enable/disable backpressure: If the `buffer in NC` + `credit available in R` is less
+ * than the `backlog of data in NC`, backpressure needs to be enabled. For this, the NC sends a control message
+ * to S to enable backpressure (pause data processing) and adds R to `overloaded` list. On the other hand, if
+ * `buffer in NC` + `credit available in R` is enough, then R is removed from the overloaded list (if present). If
+ * the `overloaded` list is empty, then NC sends a request to S to disable backpressure (resume processing).
+ *
+ * b) It looks at its backlog and sends an amount of data, less than credits available, to congestion control.
+ *
+ * 3. If R sends a credit of 0, then S won't send any data as a response to NetworkAck(). This will lead to a problem
+ * because then there is no way for S to know when the data in its congestion control module can be sent. Thus,
+ * whenever S receives a credit of 0, it registers a periodic callback that serves as a trigger for it to send
+ * credit poll request to R. Then, R responds with a NetworkAck() for the credits.
+ *
+ * 4. In our current design, the term "Credit" refers to the message in memory size in bytes.
+ */
+class FlowControl {
+
+ private val maxByteAllowed = ApplicationConfig.maxCreditAllowedInBytesPerChannel
+ private var inflightCredit: Long = 0
+ private var queuedCredit: Long = 0
+ private val stashedMessages: mutable.Queue[NetworkMessage] = new mutable.Queue()
+ private var overloaded = false
+
+ def isOverloaded: Boolean = overloaded
+
+ /**
+ * Determines if an incoming message can be forwarded to the receiver based on the credits available.
+ */
+ def getMessagesToSend(msg: NetworkMessage): Iterable[NetworkMessage] = {
+ val creditNeeded = getInMemSize(msg.internalMessage)
+ // assume the biggest message can pass through flow control
+ assert(
+ creditNeeded <= maxByteAllowed,
+ s"Message $msg is too big to send through flow control, " +
+ s"max credit = $maxByteAllowed bytes " +
+ s"while the message size is $creditNeeded bytes."
+ )
+ if (stashedMessages.isEmpty) {
+ if (getCredit >= creditNeeded) {
+ inflightCredit += creditNeeded
+ Iterable(msg)
+ } else {
+ overloaded = true
+ stashedMessages.enqueue(msg)
+ Iterable.empty
+ }
+ } else {
+ stashedMessages.enqueue(msg)
+ getMessagesToSend
+ }
+ }
+
+ def getMessagesToSend: Iterable[NetworkMessage] = {
+ val toSend = mutable.ArrayBuffer[NetworkMessage]()
+ breakable {
+ while (stashedMessages.nonEmpty) {
+ val msg = stashedMessages.front
+ val creditNeeded = getInMemSize(msg.internalMessage)
+ if (getCredit >= creditNeeded) {
+ inflightCredit += creditNeeded
+ toSend.append(msg)
+ stashedMessages.dequeue()
+ } else {
+ break()
+ }
+ }
+ }
+ overloaded = stashedMessages.nonEmpty
+ toSend
+ }
+
+ def updateQueuedCredit(newCredit: Long): Unit = {
+ queuedCredit = newCredit
+ }
+
+ def decreaseInflightCredit(ackedCredit: Long): Unit = {
+ inflightCredit -= ackedCredit
+ }
+
+ def getCredit: Long = {
+ maxByteAllowed - inflightCredit - queuedCredit
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/InputGateway.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/InputGateway.scala
new file mode 100644
index 00000000000..8d0c313f806
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/InputGateway.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+import org.apache.texera.amber.engine.architecture.logreplay.OrderEnforcer
+
+trait InputGateway {
+
+ def tryPickControlChannel: Option[AmberFIFOChannel]
+
+ def tryPickChannel: Option[AmberFIFOChannel]
+
+ def getAllChannels: Iterable[AmberFIFOChannel]
+
+ def getAllDataChannels: Iterable[AmberFIFOChannel]
+
+ def getChannel(channelId: ChannelIdentity): AmberFIFOChannel
+
+ def getAllControlChannels: Iterable[AmberFIFOChannel]
+
+ def addEnforcer(enforcer: OrderEnforcer): Unit
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/InputManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/InputManager.scala
new file mode 100644
index 00000000000..4e297829431
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/InputManager.scala
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.tuple.{Schema, Tuple}
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.Partitioning
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.DPInputQueueElement
+import org.apache.texera.amber.engine.architecture.worker.managers.InputPortMaterializationReaderThread
+import org.apache.texera.amber.engine.common.AmberLogging
+
+import java.net.URI
+import java.util.concurrent.LinkedBlockingQueue
+import scala.collection.mutable
+
+class InputManager(
+ val actorId: ActorVirtualIdentity,
+ val inputMessageQueue: LinkedBlockingQueue[DPInputQueueElement]
+) extends AmberLogging {
+ private var inputBatch: Array[Tuple] = _
+ private var currentInputIdx: Int = -1
+ var currentChannelId: ChannelIdentity = _
+
+ private val ports: mutable.HashMap[PortIdentity, WorkerPort] = mutable.HashMap()
+
+ private val inputPortMaterializationReaderThreads
+ : mutable.HashMap[PortIdentity, List[InputPortMaterializationReaderThread]] =
+ mutable.HashMap()
+
+ def getAllPorts: Set[PortIdentity] = {
+ this.ports.keys.toSet
+ }
+
+ def addPort(
+ portId: PortIdentity,
+ schema: Schema,
+ urisToRead: List[URI],
+ partitionings: List[Partitioning]
+ ): Unit = {
+ assert(urisToRead.size == partitionings.size)
+ // each port can only be added and initialized once.
+ if (this.ports.contains(portId)) {
+ return
+ }
+ this.ports(portId) = WorkerPort(schema)
+
+ // if a materialization URI is provided, set up a materialization reader thread
+ setupInputPortMaterializationReaderThreads(portId, urisToRead, partitionings)
+ }
+
+ private def setupInputPortMaterializationReaderThreads(
+ portId: PortIdentity,
+ uris: List[URI],
+ partitionings: List[Partitioning]
+ ): Unit = {
+ if (uris.isEmpty) {
+ return
+ }
+ val readerThreads = uris.zip(partitionings).map {
+ case (uri, partitioning) =>
+ new InputPortMaterializationReaderThread(
+ uri = uri,
+ inputMessageQueue = this.inputMessageQueue,
+ workerActorId = this.actorId,
+ partitioning = partitioning
+ )
+ }
+
+ inputPortMaterializationReaderThreads(portId) = readerThreads
+ }
+
+ def getInputPortReaderThreads: Map[PortIdentity, List[InputPortMaterializationReaderThread]] = {
+ this.inputPortMaterializationReaderThreads.toMap
+ }
+
+ def startInputPortReaderThreads(): Unit = {
+ this.inputPortMaterializationReaderThreads
+ .filterNot {
+ // A completed port should not be started again
+ case (portId, _) => this.isPortCompleted(portId)
+ }
+ .values
+ .foreach(threadList =>
+ threadList.foreach(readerThread => {
+ try {
+ readerThread.start()
+ } catch {
+ case e: Exception =>
+ throw new RuntimeException(
+ s"Error starting input port materialization reader thread: ${e.getMessage}"
+ )
+ }
+ })
+ )
+ }
+
+ def getPort(portId: PortIdentity): WorkerPort = ports(portId)
+
+ /**
+ * For ports that read from materialization, the port completion is marked by the finish of the reader thread.
+ * For other ports that connect to upstream links, the completion is marked by the completion the port.
+ */
+ def isPortCompleted(portId: PortIdentity): Boolean = {
+ if (
+ !this.inputPortMaterializationReaderThreads
+ .contains(portId) || this.inputPortMaterializationReaderThreads(portId).isEmpty
+ ) {
+ this.getPort(portId).completed
+ } else {
+ val existingThread = this.inputPortMaterializationReaderThreads(portId).head
+ existingThread.finished
+ }
+ }
+
+ def hasUnfinishedInput: Boolean = inputBatch != null && currentInputIdx + 1 < inputBatch.length
+
+ def getNextTuple: Tuple = {
+ currentInputIdx += 1
+ inputBatch(currentInputIdx)
+ }
+
+ def getCurrentTuple: Tuple = {
+ if (inputBatch == null) {
+ null
+ } else if (inputBatch.isEmpty) {
+ null // TODO: create input exhausted
+ } else {
+ inputBatch(currentInputIdx)
+ }
+ }
+
+ def initBatch(channelId: ChannelIdentity, batch: Array[Tuple]): Unit = {
+ currentChannelId = channelId
+ inputBatch = batch
+ currentInputIdx = -1
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkInputGateway.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkInputGateway.scala
new file mode 100644
index 00000000000..aed5c36c4a0
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkInputGateway.scala
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.logreplay.OrderEnforcer
+import org.apache.texera.amber.engine.common.AmberLogging
+
+import scala.collection.mutable
+
+class NetworkInputGateway(val actorId: ActorVirtualIdentity)
+ extends AmberLogging
+ with Serializable
+ with InputGateway {
+
+ private val inputChannels =
+ new mutable.HashMap[ChannelIdentity, AmberFIFOChannel]()
+
+ @transient lazy private val enforcers = mutable.ListBuffer[OrderEnforcer]()
+
+ def tryPickControlChannel: Option[AmberFIFOChannel] = {
+ val ret = inputChannels
+ .find {
+ case (cid, channel) =>
+ cid.isControl && channel.isEnabled && channel.hasMessage && enforcers.forall(enforcer =>
+ enforcer.isCompleted || enforcer.canProceed(cid)
+ )
+ }
+ .map(_._2)
+
+ enforcers.filter(enforcer => enforcer.isCompleted).foreach(enforcer => enforcers -= enforcer)
+ ret
+ }
+
+ def tryPickChannel: Option[AmberFIFOChannel] = {
+ val control = tryPickControlChannel
+ val ret = if (control.isDefined) {
+ control
+ } else {
+ inputChannels
+ .find({
+ case (cid, channel) =>
+ !cid.isControl && channel.isEnabled && channel.hasMessage && enforcers
+ .forall(enforcer => enforcer.isCompleted || enforcer.canProceed(cid))
+ })
+ .map(_._2)
+ }
+ enforcers.filter(enforcer => enforcer.isCompleted).foreach(enforcer => enforcers -= enforcer)
+ ret
+ }
+
+ def getAllDataChannels: Iterable[AmberFIFOChannel] =
+ inputChannels.filter(!_._1.isControl).values
+
+ // this function is called by both main thread(for getting credit)
+ // and DP thread(for enqueuing messages) so a lock is required here
+ def getChannel(channelId: ChannelIdentity): AmberFIFOChannel = {
+ synchronized {
+ inputChannels.getOrElseUpdate(channelId, new AmberFIFOChannel(channelId))
+ }
+ }
+
+ def getAllControlChannels: Iterable[AmberFIFOChannel] =
+ inputChannels.filter(_._1.isControl).values
+
+ override def getAllChannels: Iterable[AmberFIFOChannel] = inputChannels.values
+
+ override def addEnforcer(enforcer: OrderEnforcer): Unit = {
+ enforcers += enforcer
+ }
+
+ def removeControlChannel(from: ActorVirtualIdentity): Unit = {
+ synchronized {
+ inputChannels.remove(ChannelIdentity(from, actorId, isControl = true))
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkOutputGateway.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkOutputGateway.scala
new file mode 100644
index 00000000000..ea7034e1d78
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkOutputGateway.scala
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataPayload,
+ DirectControlMessagePayload,
+ WorkflowFIFOMessage,
+ WorkflowFIFOMessagePayload
+}
+import org.apache.texera.amber.engine.common.virtualidentity.util.SELF
+
+import java.util.concurrent.atomic.AtomicLong
+import scala.collection.mutable
+
+/**
+ * NetworkOutput for generating sequence number when sending payloads
+ *
+ * @param actorId ActorVirtualIdentity for the sender
+ * @param handler actual sending logic
+ */
+class NetworkOutputGateway(
+ val actorId: ActorVirtualIdentity,
+ val handler: WorkflowFIFOMessage => Unit
+) extends AmberLogging
+ with Serializable {
+
+ private val idToSequenceNums = new mutable.HashMap[ChannelIdentity, AtomicLong]()
+
+ def addOutputChannel(channelId: ChannelIdentity): Unit = {
+ if (!idToSequenceNums.contains(channelId)) {
+ idToSequenceNums(channelId) = new AtomicLong()
+ }
+ }
+
+ private def sendToInternal(
+ to: ActorVirtualIdentity,
+ useControlChannel: Boolean,
+ payload: WorkflowFIFOMessagePayload
+ ): Unit = {
+ var receiverId = to
+ if (to == SELF) {
+ // selfID and VirtualIdentity.SELF should be one key
+ receiverId = actorId
+ }
+ val outChannelId = ChannelIdentity(actorId, receiverId, useControlChannel)
+ val seqNum = getSequenceNumber(outChannelId)
+ handler(WorkflowFIFOMessage(outChannelId, seqNum, payload))
+ }
+
+ def sendTo(to: ActorVirtualIdentity, payload: DirectControlMessagePayload): Unit = {
+ sendToInternal(to, useControlChannel = true, payload)
+ }
+
+ def sendTo(to: ActorVirtualIdentity, payload: DataPayload): Unit = {
+ sendToInternal(to, useControlChannel = false, payload)
+ }
+
+ def sendTo(channelIdentity: ChannelIdentity, payload: WorkflowFIFOMessagePayload): Unit = {
+ val destChannelId = if (channelIdentity.toWorkerId == SELF) {
+ // selfID and VirtualIdentity.SELF should be one key
+ ChannelIdentity(channelIdentity.fromWorkerId, actorId, channelIdentity.isControl)
+ } else {
+ channelIdentity
+ }
+ val seqNum = getSequenceNumber(destChannelId)
+ handler(WorkflowFIFOMessage(destChannelId, seqNum, payload))
+ }
+
+ def getFIFOState: Map[ChannelIdentity, Long] = idToSequenceNums.map(x => (x._1, x._2.get())).toMap
+
+ def getActiveChannels: Iterable[ChannelIdentity] = idToSequenceNums.keys
+
+ def getSequenceNumber(channelId: ChannelIdentity): Long = {
+ idToSequenceNums.getOrElseUpdate(channelId, new AtomicLong()).getAndIncrement()
+ }
+
+ def removeControlChannel(to: ActorVirtualIdentity): Unit = {
+ synchronized {
+ idToSequenceNums.remove(ChannelIdentity(actorId, to, isControl = true))
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OrderingEnforcer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OrderingEnforcer.scala
new file mode 100644
index 00000000000..d840738a45d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OrderingEnforcer.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import scala.collection.mutable
+
+/* The abstracted FIFO/exactly-once logic */
+class OrderingEnforcer[T] {
+
+ val ofoMap = new mutable.LongMap[T]
+ var current = 0L
+
+ def setCurrent(value: Long): Unit = {
+ current = value
+ }
+
+ def isDuplicated(sequenceNumber: Long): Boolean =
+ sequenceNumber < current || ofoMap.contains(sequenceNumber)
+
+ def isAhead(sequenceNumber: Long): Boolean = sequenceNumber > current
+
+ def stash(sequenceNumber: Long, data: T): Unit = {
+ ofoMap(sequenceNumber) = data
+ }
+
+ def enforceFIFO(data: T): List[T] = {
+ val res = mutable.ArrayBuffer[T](data)
+ current += 1
+ while (ofoMap.contains(current)) {
+ res.append(ofoMap(current))
+ ofoMap.remove(current)
+ current += 1
+ }
+ res.toList
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala
new file mode 100644
index 00000000000..4ab3d18056f
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala
@@ -0,0 +1,293 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.state.State
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.storage.model.BufferedItemWriter
+import org.apache.texera.amber.core.tuple._
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow.{PhysicalLink, PortIdentity}
+import org.apache.texera.amber.engine.architecture.messaginglayer.OutputManager.{
+ DPOutputIterator,
+ getBatchSize,
+ toPartitioner
+}
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitioners._
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings._
+import org.apache.texera.amber.engine.architecture.worker.managers.{
+ OutputPortResultWriterThread,
+ PortStorageWriterTerminateSignal
+}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+import java.net.URI
+import scala.collection.mutable
+
+object OutputManager {
+
+ // create a corresponding partitioner for the given partitioning policy
+ def toPartitioner(partitioning: Partitioning, actorId: ActorVirtualIdentity): Partitioner = {
+ val partitioner = partitioning match {
+ case oneToOnePartitioning: OneToOnePartitioning =>
+ OneToOnePartitioner(oneToOnePartitioning, actorId)
+ case roundRobinPartitioning: RoundRobinPartitioning =>
+ RoundRobinPartitioner(roundRobinPartitioning)
+ case hashBasedShufflePartitioning: HashBasedShufflePartitioning =>
+ HashBasedShufflePartitioner(hashBasedShufflePartitioning)
+ case rangeBasedShufflePartitioning: RangeBasedShufflePartitioning =>
+ RangeBasedShufflePartitioner(rangeBasedShufflePartitioning)
+ case broadcastPartitioning: BroadcastPartitioning =>
+ BroadcastPartitioner(broadcastPartitioning)
+ case _ => throw new RuntimeException(s"partitioning $partitioning not supported")
+ }
+ partitioner
+ }
+
+ def getBatchSize(partitioning: Partitioning): Int = {
+ partitioning match {
+ case p: OneToOnePartitioning => p.batchSize
+ case p: RoundRobinPartitioning => p.batchSize
+ case p: HashBasedShufflePartitioning => p.batchSize
+ case p: RangeBasedShufflePartitioning => p.batchSize
+ case p: BroadcastPartitioning => p.batchSize
+ case _ => throw new RuntimeException(s"partitioning $partitioning not supported")
+ }
+ }
+
+ class DPOutputIterator extends Iterator[(TupleLike, Option[PortIdentity])] {
+ val queue = new mutable.ListBuffer[(TupleLike, Option[PortIdentity])]
+ @transient var outputIter: Iterator[(TupleLike, Option[PortIdentity])] = Iterator.empty
+
+ def setTupleOutput(outputIter: Iterator[(TupleLike, Option[PortIdentity])]): Unit = {
+ if (outputIter != null) {
+ this.outputIter = outputIter
+ } else {
+ this.outputIter = Iterator.empty
+ }
+ }
+
+ override def hasNext: Boolean = outputIter.hasNext || queue.nonEmpty
+
+ override def next(): (TupleLike, Option[PortIdentity]) = {
+ if (outputIter.hasNext) {
+ outputIter.next()
+ } else {
+ queue.remove(0)
+ }
+ }
+
+ def appendSpecialTupleToEnd(tuple: TupleLike): Unit = {
+ queue.append((tuple, None))
+ }
+ }
+}
+
+/** This class is a container of all the transfer partitioners.
+ *
+ * @param actorId ActorVirtualIdentity of self.
+ * @param outputGateway DataOutputPort
+ */
+class OutputManager(
+ val actorId: ActorVirtualIdentity,
+ outputGateway: NetworkOutputGateway
+) extends AmberLogging {
+
+ val outputIterator: DPOutputIterator = new DPOutputIterator()
+ private val partitioners: mutable.Map[PhysicalLink, Partitioner] =
+ mutable.HashMap[PhysicalLink, Partitioner]()
+
+ private val ports: mutable.HashMap[PortIdentity, WorkerPort] = mutable.HashMap()
+
+ private val networkOutputBuffers =
+ mutable.HashMap[(PhysicalLink, ActorVirtualIdentity), NetworkOutputBuffer]()
+
+ private val outputPortResultWriterThreads
+ : mutable.HashMap[PortIdentity, OutputPortResultWriterThread] =
+ mutable.HashMap()
+
+ /**
+ * Add down stream operator and its corresponding Partitioner.
+ *
+ * @param partitioning Partitioning, describes how and whom to send to.
+ */
+ def addPartitionerWithPartitioning(
+ link: PhysicalLink,
+ partitioning: Partitioning
+ ): Unit = {
+ val partitioner = toPartitioner(partitioning, actorId)
+ partitioners.update(link, partitioner)
+ partitioner.allReceivers.foreach(receiver => {
+ val buffer = new NetworkOutputBuffer(receiver, outputGateway, getBatchSize(partitioning))
+ networkOutputBuffers.update((link, receiver), buffer)
+ outputGateway.addOutputChannel(ChannelIdentity(actorId, receiver, isControl = false))
+ })
+ }
+
+ /**
+ * Push one tuple to the downstream, will be batched by each transfer partitioning.
+ * Should ONLY be called by DataProcessor.
+ *
+ * @param tuple TupleLike to be passed.
+ * @param outputPortId Optionally specifies the output port from which the tuple should be emitted.
+ * If None, the tuple is broadcast to all output ports.
+ */
+ def passTupleToDownstream(
+ tuple: Tuple,
+ outputPortId: Option[PortIdentity] = None
+ ): Unit = {
+ (outputPortId match {
+ case Some(portId) => partitioners.filter(_._1.fromPortId == portId) // send to a specific port
+ case None => partitioners // send to all ports
+ }).foreach {
+ case (link, partitioner) =>
+ partitioner.getBucketIndex(tuple).foreach { bucketIndex =>
+ networkOutputBuffers((link, partitioner.allReceivers(bucketIndex))).addTuple(tuple)
+ }
+ }
+ }
+
+ /**
+ * Flushes the network output buffers based on the specified set of physical links.
+ *
+ * This method flushes the buffers associated with the network output. If the 'onlyFor' parameter
+ * is specified with a set of 'PhysicalLink's, only the buffers corresponding to those links are flushed.
+ * If 'onlyFor' is None, all network output buffers are flushed.
+ *
+ * @param onlyFor An optional set of 'ChannelID' indicating the specific buffers to flush.
+ * If None, all buffers are flushed. Default value is None.
+ */
+ def flush(onlyFor: Option[Set[ChannelIdentity]] = None): Unit = {
+ val buffersToFlush = onlyFor match {
+ case Some(channelIds) =>
+ networkOutputBuffers
+ .filter(out => {
+ val channel = ChannelIdentity(actorId, out._1._2, isControl = false)
+ channelIds.contains(channel)
+ })
+ .values
+ case None => networkOutputBuffers.values
+ }
+ buffersToFlush.foreach(_.flush())
+ }
+
+ def emitState(state: State): Unit = {
+ networkOutputBuffers.foreach(kv => kv._2.sendState(state))
+ }
+
+ def addPort(portId: PortIdentity, schema: Schema, storageURIOption: Option[URI]): Unit = {
+ // each port can only be added and initialized once.
+ if (this.ports.contains(portId)) {
+ return
+ }
+ this.ports(portId) = WorkerPort(schema)
+
+ // if a storage URI is provided, set up a storage writer thread
+ storageURIOption match {
+ case Some(storageUri) => setupOutputStorageWriterThread(portId, storageUri)
+ case None => // No need to add a writer
+ }
+ }
+
+ /**
+ * Optionally write the tuple to storage if the specified output port is determined by the scheduler to need storage.
+ * This method is not blocking because a separate thread is used to flush the tuple to storage in batch.
+ *
+ * @param tuple TupleLike to be written to storage.
+ * @param outputPortId If not specified, the tuple will be written to all output ports that need storage.
+ */
+ def saveTupleToStorageIfNeeded(
+ tuple: Tuple,
+ outputPortId: Option[PortIdentity] = None
+ ): Unit = {
+ (outputPortId match {
+ case Some(portId) =>
+ this.outputPortResultWriterThreads.get(portId) match {
+ case Some(_) => this.outputPortResultWriterThreads.filter(_._1 == portId)
+ case None => Map.empty
+ }
+ case None => this.outputPortResultWriterThreads
+ }).foreach({
+ case (portId, writerThread) =>
+ // write to storage in a separate thread
+ writerThread.queue.put(Left(tuple))
+ })
+ }
+
+ /**
+ * Singal the port storage writer to flush the remaining buffer and wait for commits to finish so that
+ * the output port is properly completed. If the output port does not need storage, no action will be done.
+ */
+ def closeOutputStorageWriterIfNeeded(outputPortId: PortIdentity): Unit = {
+ this.outputPortResultWriterThreads.get(outputPortId) match {
+ case Some(writerThread) =>
+ // Non-blocking call
+ writerThread.queue.put(Right(PortStorageWriterTerminateSignal))
+ // Blocking call
+ writerThread.join()
+ case None =>
+ }
+
+ }
+
+ def getPort(portId: PortIdentity): WorkerPort = ports(portId)
+
+ def hasUnfinishedOutput: Boolean = outputIterator.hasNext
+
+ def finalizeOutput(): Unit = {
+ this.ports.keys
+ .foreach(outputPortId =>
+ outputIterator.appendSpecialTupleToEnd(FinalizePort(outputPortId, input = false))
+ )
+ outputIterator.appendSpecialTupleToEnd(FinalizeExecutor())
+ }
+
+ /**
+ * This method is only used for ensuring correct region execution. Some operators may have input port dependency
+ * relationships, for which we currently use a two-phase region execution scheme. (See `RegionExecutionCoordinator`
+ * for details.)
+ * This logic will only be executed when the worker is part of an `executingDependeePort` region-execution phase.
+ * We currently assume that in this phase the operator (worker) will not output any data, hence no output ports.
+ * However we still need to keep this worker open for the next `executingNonDependeePort` phase.
+ *
+ * @return Whether this worker currently does not have any output port.
+ */
+ def isMissingOutputPort: Boolean = {
+ this.ports.isEmpty
+ }
+
+ def getSingleOutputPortIdentity: PortIdentity = {
+ assert(ports.size == 1, "expect 1 output port, got " + ports.size)
+ ports.head._1
+ }
+
+ private def setupOutputStorageWriterThread(portId: PortIdentity, storageUri: URI): Unit = {
+ val bufferedItemWriter = DocumentFactory
+ .openDocument(storageUri)
+ ._1
+ .writer(VirtualIdentityUtils.getWorkerIndex(actorId).toString)
+ .asInstanceOf[BufferedItemWriter[Tuple]]
+ val writerThread = new OutputPortResultWriterThread(bufferedItemWriter)
+ this.outputPortResultWriterThreads(portId) = writerThread
+ writerThread.start()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/WorkerPort.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/WorkerPort.scala
new file mode 100644
index 00000000000..504aaf93704
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/WorkerPort.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.tuple.Schema
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+
+import scala.collection.mutable
+case class WorkerPort(
+ schema: Schema,
+ channels: mutable.Set[ChannelIdentity] = mutable.Set(),
+ var completed: Boolean = false
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/WorkerTimerService.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/WorkerTimerService.scala
new file mode 100644
index 00000000000..3bb87febd93
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/WorkerTimerService.scala
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.pekko.actor.Cancellable
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.engine.architecture.common.AkkaActorService
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_FLUSH_NETWORK_BUFFER
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
+import org.apache.texera.amber.engine.common.virtualidentity.util.SELF
+
+import scala.concurrent.duration.{DurationInt, FiniteDuration, MILLISECONDS}
+
+class WorkerTimerService(actorService: AkkaActorService) {
+
+ private val enabledAdaptiveBatching = ApplicationConfig.enableAdaptiveNetworkBuffering
+ private val adaptiveBatchInterval = ApplicationConfig.adaptiveBufferingTimeoutMs
+
+ var adaptiveBatchingHandle: Option[Cancellable] = None
+ var isPaused = false
+
+ def startAdaptiveBatching(): Unit = {
+ if (!enabledAdaptiveBatching) {
+ return
+ }
+ if (this.adaptiveBatchingHandle.nonEmpty) {
+ return
+ }
+ this.adaptiveBatchingHandle = Some(
+ actorService.sendToSelfWithFixedDelay(
+ 0.milliseconds,
+ FiniteDuration.apply(adaptiveBatchInterval, MILLISECONDS),
+ ControlInvocation(
+ METHOD_FLUSH_NETWORK_BUFFER, // uses method descriptor instead of method name string
+ EmptyRequest(),
+ AsyncRPCContext(SELF, SELF),
+ AsyncRPCClient.IgnoreReplyAndDoNotLog
+ )
+ )
+ )
+ }
+
+ def stopAdaptiveBatching(): Unit = {
+ if (adaptiveBatchingHandle.nonEmpty) {
+ adaptiveBatchingHandle.get.cancel()
+ }
+ isPaused = false
+ }
+
+ def pauseAdaptiveBatching(): Unit = {
+ stopAdaptiveBatching()
+ isPaused = true
+ }
+
+ def resumeAdaptiveBatching(): Unit = {
+ if (isPaused) {
+ startAdaptiveBatching()
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala
new file mode 100644
index 00000000000..6618e857b1d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala
@@ -0,0 +1,256 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.pythonworker
+
+import com.twitter.util.{Await, Promise}
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.core.tuple.{Schema, Tuple}
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.pythonworker.WorkerBatchInternalQueue.{
+ ActorCommandElement,
+ ControlElement,
+ DataElement,
+ EmbeddedControlMessageElement
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ ControlInvocation,
+ EmbeddedControlMessage
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.ReturnInvocation
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.actormessage.{ActorCommand, PythonActorMessage}
+import org.apache.texera.amber.engine.common.ambermessage._
+import org.apache.texera.amber.util.ArrowUtils
+import org.apache.arrow.flight._
+import org.apache.arrow.memory.{ArrowBuf, BufferAllocator, RootAllocator}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema => ArrowSchema}
+import org.apache.arrow.vector.{VarBinaryVector, VectorSchemaRoot}
+
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicLong
+import scala.collection.compat.immutable.ArraySeq
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+class PythonProxyClient(portNumberPromise: Promise[Int], val actorId: ActorVirtualIdentity)
+ extends Runnable
+ with AmberLogging
+ with AutoCloseable
+ with WorkerBatchInternalQueue {
+
+ val allocator: BufferAllocator =
+ new RootAllocator().newChildAllocator("flight-client", 0, Long.MaxValue)
+ val location: Location = (() => {
+ // Read port number from promise until it's available
+ val portNumber = Await.result(portNumberPromise)
+ Location.forGrpcInsecure("localhost", portNumber)
+ })()
+
+ private val MAX_TRY_COUNT: Int = 2
+ private val UNIT_WAIT_TIME_MS = 200
+ private var flightClient: FlightClient = _
+ private var running: Boolean = true
+
+ private val pythonQueueInMemSize: AtomicLong = new AtomicLong(0)
+
+ def getQueuedCredit: Long = {
+ pythonQueueInMemSize.get()
+ }
+
+ override def run(): Unit = {
+ establishConnection()
+ mainLoop()
+ }
+
+ private def establishConnection(): Unit = {
+ var connected = false
+ var tryCount = 0
+ while (!connected && tryCount <= MAX_TRY_COUNT) {
+ try {
+ flightClient = FlightClient.builder(allocator, location).build()
+ connected = new String(flightClient.doAction(new Action("heartbeat")).next.getBody) == "ack"
+ if (!connected)
+ throw new RuntimeException("heartbeat failed")
+ } catch {
+ case _: RuntimeException =>
+ logger.warn(
+ s"Failed to connect to Flight Server in this attempt, retrying after $UNIT_WAIT_TIME_MS ms... remaining attempts: ${MAX_TRY_COUNT - tryCount}"
+ )
+ if (flightClient != null) flightClient.close()
+ Thread.sleep(UNIT_WAIT_TIME_MS)
+ tryCount += 1
+ }
+ }
+ if (!connected) {
+ throw new WorkflowRuntimeException(
+ s"Failed to connect to Flight Server after $MAX_TRY_COUNT attempts. Abort!"
+ )
+ }
+ }
+
+ private def mainLoop(): Unit = {
+ while (running) {
+ getElement match {
+ case DataElement(dataPayload, channel) =>
+ sendData(dataPayload, channel)
+ case ControlElement(cmd, channel) =>
+ sendControl(channel, cmd)
+ case EmbeddedControlMessageElement(cmd, channel) =>
+ sendECM(cmd, channel)
+ case ActorCommandElement(cmd) =>
+ sendActorCommand(cmd)
+ }
+ }
+ }
+
+ private def sendData(dataPayload: DataPayload, from: ChannelIdentity): Unit = {
+ dataPayload match {
+ case DataFrame(frame) =>
+ writeArrowStream(mutable.Queue(ArraySeq.unsafeWrapArray(frame): _*), from, "Data")
+ case StateFrame(state) =>
+ writeArrowStream(mutable.Queue(state.toTuple), from, "State")
+ }
+ }
+
+ private def sendECM(
+ ecm: EmbeddedControlMessage,
+ from: ChannelIdentity
+ ): Unit = {
+ val descriptor = FlightDescriptor.command(PythonDataHeader(from, "ECM").toByteArray)
+ val flightListener = new SyncPutListener
+
+ val field = new Field("payload", FieldType.nullable(new ArrowType.Binary), null)
+ val schema = new ArrowSchema(List(field).asJava)
+ val schemaRoot = VectorSchemaRoot.create(schema, allocator)
+
+ val writer = flightClient.startPut(descriptor, schemaRoot, flightListener)
+ schemaRoot.allocateNew()
+
+ val vector = schemaRoot.getVector("payload").asInstanceOf[VarBinaryVector]
+ vector.setSafe(0, ecm.toByteArray)
+ vector.setValueCount(1)
+ schemaRoot.setRowCount(1)
+
+ writer.putNext()
+ schemaRoot.clear()
+ writer.completed()
+
+ // for calculating sender credits - get back number of batches in Python worker queue
+ val ackMsgBuf: ArrowBuf = flightListener.poll(5, TimeUnit.SECONDS).getApplicationMetadata
+ pythonQueueInMemSize.set(ackMsgBuf.getLong(0))
+ logger.debug(s"data channel updated queue size $pythonQueueInMemSize")
+ ackMsgBuf.close()
+
+ flightListener.close()
+ }
+
+ private def sendControl(
+ from: ChannelIdentity,
+ payload: DirectControlMessagePayload
+ ): Result = {
+ var payloadV2 = DirectControlMessagePayloadV2.defaultInstance
+ payloadV2 = payload match {
+ case c: ControlInvocation =>
+ payloadV2.withControlInvocation(c)
+ case r: ReturnInvocation =>
+ payloadV2.withReturnInvocation(r)
+ case _ => ???
+ }
+ val controlMessage = PythonControlMessage(from, payloadV2)
+ val action: Action = new Action("control", controlMessage.toByteArray)
+ sendCreditedAction(action)
+ }
+
+ private def sendActorCommand(
+ command: ActorCommand
+ ): Result = {
+ val action: Action = new Action("actor", PythonActorMessage(command).toByteArray)
+ sendCreditedAction(action)
+ }
+
+ private def sendCreditedAction(action: Action) = {
+ logger.debug(s"sending ${action.getType} message")
+ // Arrow allows multiple results from the Action call return as a stream (interator).
+ // In Arrow 11, it alerts if the results are not consumed fully.
+ val results = flightClient.doAction(action)
+ // As we do our own Async RPC management, we are currently not using results from Action call.
+ // In the future, this results can include credits for flow control purpose.
+ val result = results.next()
+
+ // extract info needed to calculate sender credits from ack
+ // ackResult contains number of batches inside Python worker internal queue
+ pythonQueueInMemSize.set(new String(result.getBody).toLong)
+ logger.debug(s"action ${action.getType} updated queue size $pythonQueueInMemSize")
+ // However, we will only expect exactly one result for now.
+ assert(!results.hasNext)
+
+ result
+ }
+
+ private def writeArrowStream(
+ tuples: mutable.Queue[Tuple],
+ from: ChannelIdentity,
+ payloadType: String
+ ): Unit = {
+
+ val schema = if (tuples.isEmpty) new Schema() else tuples.front.getSchema
+ val descriptor = FlightDescriptor.command(PythonDataHeader(from, payloadType).toByteArray)
+ logger.debug(
+ s"sending data with descriptor ${PythonDataHeader(from, payloadType)}, schema $schema, size of batch ${tuples.size}"
+ )
+ val flightListener = new SyncPutListener
+ val schemaRoot = VectorSchemaRoot.create(ArrowUtils.fromTexeraSchema(schema), allocator)
+ val writer = flightClient.startPut(descriptor, schemaRoot, flightListener)
+ schemaRoot.allocateNew()
+ while (tuples.nonEmpty) {
+ ArrowUtils.appendTexeraTuple(tuples.dequeue(), schemaRoot)
+ }
+ writer.putNext()
+ schemaRoot.clear()
+ writer.completed()
+
+ // for calculating sender credits - get back number of batches in Python worker queue
+ val ackMsgBuf: ArrowBuf = flightListener.poll(5, TimeUnit.SECONDS).getApplicationMetadata
+ pythonQueueInMemSize.set(ackMsgBuf.getLong(0))
+ logger.debug(s"data channel updated queue size $pythonQueueInMemSize")
+ ackMsgBuf.close()
+
+ flightListener.close()
+
+ }
+
+ override def close(): Unit = {
+ val action: Action = new Action("shutdown")
+ try {
+ flightClient.doAction(action) // do not expect reply
+
+ flightClient.close()
+ } catch {
+ case _: NullPointerException =>
+ running = false
+ logger.warn(
+ s"Unable to close the flight client because it is null"
+ )
+ }
+ // stop the main loop
+ running = false
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServer.scala
new file mode 100644
index 00000000000..2ff866365bb
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServer.scala
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.pythonworker
+
+import com.google.common.primitives.Longs
+import com.twitter.util.Promise
+import org.apache.texera.amber.core.state.State
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.messaginglayer.NetworkOutputGateway
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessage
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.ambermessage.DirectControlMessagePayloadV2.Value.{
+ ControlInvocation => ControlInvocationV2,
+ ReturnInvocation => ReturnInvocationV2
+}
+import org.apache.texera.amber.engine.common.ambermessage._
+import org.apache.texera.amber.util.ArrowUtils
+import org.apache.arrow.flight._
+import org.apache.arrow.memory.{ArrowBuf, BufferAllocator, RootAllocator}
+import org.apache.arrow.util.AutoCloseables
+import org.apache.arrow.vector.VarBinaryVector
+
+import java.io.IOException
+import java.net.ServerSocket
+import java.nio.charset.Charset
+import java.nio.{ByteBuffer, ByteOrder}
+import java.util.concurrent.atomic.AtomicInteger
+import scala.collection.mutable
+
+private class AmberProducer(
+ actorId: ActorVirtualIdentity,
+ outputPort: NetworkOutputGateway,
+ promise: Promise[Int]
+) extends NoOpFlightProducer {
+ var _portNumber: AtomicInteger = new AtomicInteger(0)
+
+ def portNumber: AtomicInteger = _portNumber
+
+ override def doAction(
+ context: FlightProducer.CallContext,
+ action: Action,
+ listener: FlightProducer.StreamListener[Result]
+ ): Unit = {
+ action.getType match {
+ case "control" =>
+ val pythonControlMessage = PythonControlMessage.parseFrom(action.getBody)
+ pythonControlMessage.payload.value match {
+ case r: ReturnInvocationV2 =>
+ outputPort.sendTo(
+ to = pythonControlMessage.tag.toWorkerId,
+ payload = r.value
+ )
+
+ case c: ControlInvocationV2 =>
+ outputPort.sendTo(
+ to = pythonControlMessage.tag.toWorkerId,
+ payload = c.value
+ )
+ case payload =>
+ throw new RuntimeException(s"not supported payload $payload")
+ }
+
+ // get little-endian representation of credits
+ var creditVal: Long = 30L // TODO : replace with actual credit value
+ val creditByteArr: Array[Byte] =
+ ByteBuffer.allocate(Longs.BYTES).order(ByteOrder.LITTLE_ENDIAN).putLong(creditVal).array
+
+ listener.onNext(
+ new Result(creditByteArr)
+ )
+ listener.onCompleted()
+ case "handshake" =>
+ val strPortNumber: String = new String(action.getBody, Charset.forName("UTF-8"))
+ // Receive the port number from Python and put it into promise
+ promise.setValue(strPortNumber.toInt)
+ listener.onNext(new Result("ok".getBytes))
+ listener.onCompleted()
+ case _ => throw new NotImplementedError()
+ }
+
+ }
+
+ override def acceptPut(
+ context: FlightProducer.CallContext,
+ flightStream: FlightStream,
+ ackStream: FlightProducer.StreamListener[PutResult]
+ ): Runnable = { () =>
+ val dataHeader: PythonDataHeader = PythonDataHeader
+ .parseFrom(flightStream.getDescriptor.getCommand)
+ val to: ChannelIdentity = dataHeader.tag
+ val root = flightStream.getRoot
+
+ // send back ack with credits on ackStream
+ val bufferAllocator = new RootAllocator(8 * 1024)
+ try {
+ val arrowBuf: ArrowBuf = bufferAllocator.buffer(Longs.BYTES + 4)
+ arrowBuf.writeLong(
+ 31L
+ ) // TODO : replace with actual credit value
+ ackStream.onNext(PutResult.metadata(arrowBuf))
+ arrowBuf.close()
+ } finally if (bufferAllocator != null) bufferAllocator.close()
+
+ // consume all data in the stream, it will store on the root vectors.
+ while (flightStream.next) {}
+
+ // closing the stream will release the dictionaries
+ flightStream.takeDictionaryOwnership
+
+ dataHeader.payloadType match {
+ case "State" =>
+ assert(root.getRowCount == 1)
+ outputPort.sendTo(to, StateFrame(State.fromTuple(ArrowUtils.getTexeraTuple(0, root))))
+ case "ECM" =>
+ assert(root.getRowCount == 1)
+ outputPort.sendTo(
+ to,
+ EmbeddedControlMessage.parseFrom(
+ root.getVector("payload").asInstanceOf[VarBinaryVector].get(0)
+ )
+ )
+ case _ => // normal data batches
+ val queue = mutable.Queue[Tuple]()
+ for (i <- 0 until root.getRowCount)
+ queue.enqueue(ArrowUtils.getTexeraTuple(i, root))
+ outputPort.sendTo(to, DataFrame(queue.toArray))
+ }
+ }
+}
+
+class PythonProxyServer(
+ outputPort: NetworkOutputGateway,
+ val actorId: ActorVirtualIdentity,
+ promise: Promise[Int]
+) extends Runnable
+ with AutoCloseable
+ with AmberLogging {
+ private lazy val portNumber: AtomicInteger = new AtomicInteger(getFreeLocalPort)
+
+ def getPortNumber: AtomicInteger = portNumber
+
+ val allocator: BufferAllocator =
+ new RootAllocator().newChildAllocator("flight-server", 0, Long.MaxValue)
+
+ val producer: FlightProducer = new AmberProducer(actorId, outputPort, promise)
+
+ val location: Location = (() => {
+ Location.forGrpcInsecure("localhost", portNumber.intValue())
+ })()
+
+ val server: FlightServer = FlightServer.builder(allocator, location, producer).build()
+
+ override def run(): Unit = {
+ server.start()
+ }
+
+ @throws[Exception]
+ override def close(): Unit = {
+ AutoCloseables.close(server, allocator)
+ }
+
+ /**
+ * Get a random free port.
+ *
+ * @return The port number.
+ * @throws IOException , might happen when getting a free port.
+ */
+ @throws[IOException]
+ private def getFreeLocalPort: Int = {
+ var s: ServerSocket = null
+ try {
+ // ServerSocket(0) results in availability of a free random port
+ s = new ServerSocket(0)
+ s.getLocalPort
+ } catch {
+ case e: Exception =>
+ throw new RuntimeException(e)
+ } finally {
+ assert(s != null)
+ s.close()
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorker.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorker.scala
new file mode 100644
index 00000000000..4ff5ff15ae3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorker.scala
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.pythonworker
+
+import org.apache.pekko.actor.Props
+import com.twitter.util.Promise
+import org.apache.texera.amber.config.{StorageConfig, UdfConfig}
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkAck
+import org.apache.texera.amber.engine.architecture.messaginglayer.{
+ NetworkInputGateway,
+ NetworkOutputGateway
+}
+import org.apache.texera.amber.engine.architecture.pythonworker.WorkerBatchInternalQueue.{
+ DataElement,
+ EmbeddedControlMessageElement
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessage
+import org.apache.texera.amber.engine.architecture.scheduling.config.WorkerConfig
+import org.apache.texera.amber.engine.common.actormessage.{Backpressure, CreditUpdate}
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
+import org.apache.texera.amber.engine.common.ambermessage._
+import org.apache.texera.amber.engine.common.{CheckpointState, Utils}
+import org.apache.texera.amber.config.PythonUtils
+
+import java.nio.file.Path
+import java.util.concurrent.{ExecutorService, Executors}
+import scala.sys.process.{BasicIO, Process}
+
+object PythonWorkflowWorker {
+ def props(workerConfig: WorkerConfig): Props = Props(new PythonWorkflowWorker(workerConfig))
+}
+
+class PythonWorkflowWorker(
+ workerConfig: WorkerConfig
+) extends WorkflowActor(replayLogConfOpt = None, actorId = workerConfig.workerId) {
+
+ // For receiving the Python server port number that will be available later
+ private lazy val portNumberPromise = Promise[Int]()
+ // Proxy Server and Client
+ private lazy val serverThreadExecutor: ExecutorService = Executors.newSingleThreadExecutor
+ private lazy val clientThreadExecutor: ExecutorService = Executors.newSingleThreadExecutor
+ private var pythonProxyServer: PythonProxyServer = _
+ private lazy val pythonProxyClient: PythonProxyClient =
+ new PythonProxyClient(portNumberPromise, workerConfig.workerId)
+
+ val pythonSrcDirectory: Path = Utils.amberHomePath
+ .resolve("src")
+ .resolve("main")
+ .resolve("python")
+ val RENVPath: String = UdfConfig.rPath.trim
+
+ // Python process
+ private var pythonServerProcess: Process = _
+
+ private val networkInputGateway = new NetworkInputGateway(workerConfig.workerId)
+ private val networkOutputGateway = new NetworkOutputGateway(
+ workerConfig.workerId,
+ // handler for output messages
+ msg => {
+ logManager.sendCommitted(Right(msg))
+ }
+ )
+
+ override def handleInputMessage(messageId: Long, workflowMsg: WorkflowFIFOMessage): Unit = {
+ val channel = networkInputGateway.getChannel(workflowMsg.channelId)
+ channel.acceptMessage(workflowMsg)
+ while (channel.isEnabled && channel.hasMessage) {
+ val msg = channel.take
+ msg.payload match {
+ case payload: DirectControlMessagePayload =>
+ pythonProxyClient.enqueueCommand(payload, workflowMsg.channelId)
+ case payload: DataPayload =>
+ pythonProxyClient.enqueueData(DataElement(payload, workflowMsg.channelId))
+ case ecm: EmbeddedControlMessage =>
+ pythonProxyClient.enqueueData(EmbeddedControlMessageElement(ecm, workflowMsg.channelId))
+ case p => logger.error(s"unhandled control payload: $p")
+ }
+ }
+ sender() ! NetworkAck(
+ messageId,
+ getInMemSize(workflowMsg),
+ getQueuedCredit(workflowMsg.channelId)
+ )
+ }
+
+ override def receiveCreditMessages: Receive = {
+ case WorkflowActor.CreditRequest(channel) =>
+ pythonProxyClient.enqueueActorCommand(CreditUpdate())
+ sender() ! WorkflowActor.CreditResponse(channel, getQueuedCredit(channel))
+ case WorkflowActor.CreditResponse(channel, credit) =>
+ transferService.updateChannelCreditFromReceiver(channel, credit)
+ }
+
+ /** flow-control */
+ override def getQueuedCredit(channelId: ChannelIdentity): Long = {
+ pythonProxyClient.getQueuedCredit(channelId) + pythonProxyClient.getQueuedCredit
+ }
+
+ override def handleBackpressure(enableBackpressure: Boolean): Unit = {
+ pythonProxyClient.enqueueActorCommand(Backpressure(enableBackpressure))
+ }
+
+ override def postStop(): Unit = {
+ super.postStop()
+ try {
+ // try to send shutdown command so that it can gracefully shutdown
+ pythonProxyClient.close()
+
+ clientThreadExecutor.shutdown()
+
+ serverThreadExecutor.shutdown()
+
+ // destroy python process
+ pythonServerProcess.destroy()
+ } catch {
+ case e: Exception =>
+ logger.error(s"$e - happened during shutdown")
+ }
+ }
+
+ override def initState(): Unit = {
+ startProxyServer()
+ startPythonProcess()
+ startProxyClient()
+ }
+
+ private def startProxyServer(): Unit = {
+ // Try to start the server until it succeeds
+ var serverStart = false
+ while (!serverStart) {
+ pythonProxyServer =
+ new PythonProxyServer(networkOutputGateway, workerConfig.workerId, portNumberPromise)
+ val future = serverThreadExecutor.submit(pythonProxyServer)
+ try {
+ future.get()
+ serverStart = true
+ } catch {
+ case e: Exception =>
+ future.cancel(true)
+ logger.info("Failed to start the server: " + e.getMessage + ", will try again")
+ }
+ }
+ }
+
+ private def startProxyClient(): Unit = {
+ clientThreadExecutor.submit(pythonProxyClient)
+ }
+
+ private def startPythonProcess(): Unit = {
+ val udfEntryScriptPath: String =
+ pythonSrcDirectory.resolve("texera_run_python_worker.py").toString
+ // Set the Iceberg related arguments based on the catalog type.
+ val isPostgres = StorageConfig.icebergCatalogType == "postgres"
+ val isRest = StorageConfig.icebergCatalogType == "rest"
+ pythonServerProcess = Process(
+ Seq(
+ PythonUtils.getPythonExecutable,
+ "-u",
+ udfEntryScriptPath,
+ workerConfig.workerId.name,
+ Integer.toString(pythonProxyServer.getPortNumber.get()),
+ UdfConfig.pythonLogStreamHandlerLevel,
+ RENVPath,
+ StorageConfig.icebergCatalogType,
+ if (isPostgres) StorageConfig.icebergPostgresCatalogUriWithoutScheme else "",
+ if (isPostgres) StorageConfig.icebergPostgresCatalogUsername else "",
+ if (isPostgres) StorageConfig.icebergPostgresCatalogPassword else "",
+ if (isRest) StorageConfig.icebergRESTCatalogUri else "",
+ if (isRest) StorageConfig.icebergRESTCatalogWarehouseName else "",
+ StorageConfig.icebergTableResultNamespace,
+ StorageConfig.fileStorageDirectoryPath.toString,
+ StorageConfig.icebergTableCommitBatchSize.toString,
+ StorageConfig.s3Endpoint,
+ StorageConfig.s3Region,
+ StorageConfig.s3Username,
+ StorageConfig.s3Password
+ )
+ ).run(BasicIO.standard(false))
+ }
+
+ override def loadFromCheckpoint(chkpt: CheckpointState): Unit = ???
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/WorkerBatchInternalQueue.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/WorkerBatchInternalQueue.scala
new file mode 100644
index 00000000000..11f9b6b802b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/WorkerBatchInternalQueue.scala
@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.pythonworker
+
+import lbmq.LinkedBlockingMultiQueue
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+import org.apache.texera.amber.engine.architecture.pythonworker.WorkerBatchInternalQueue._
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessage
+import org.apache.texera.amber.engine.common.actormessage.ActorCommand
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataFrame,
+ DataPayload,
+ DirectControlMessagePayload
+}
+
+import scala.collection.mutable
+
+object WorkerBatchInternalQueue {
+ final val DATA_QUEUE = 1
+ final val CONTROL_QUEUE = 0
+
+ // 4 kinds of elements can be accepted by internal queue
+ sealed trait InternalQueueElement
+
+ case class DataElement(dataPayload: DataPayload, from: ChannelIdentity)
+ extends InternalQueueElement
+
+ case class ControlElement(cmd: DirectControlMessagePayload, from: ChannelIdentity)
+ extends InternalQueueElement
+ case class EmbeddedControlMessageElement(cmd: EmbeddedControlMessage, from: ChannelIdentity)
+ extends InternalQueueElement
+ case class ActorCommandElement(cmd: ActorCommand) extends InternalQueueElement
+}
+
+/** Inspired by the mailbox-ed thread, the internal queue should
+ * be a part of DP thread.
+ */
+trait WorkerBatchInternalQueue {
+
+ private val lbmq = new LinkedBlockingMultiQueue[Int, InternalQueueElement]()
+
+ lbmq.addSubQueue(DATA_QUEUE, DATA_QUEUE)
+ lbmq.addSubQueue(CONTROL_QUEUE, CONTROL_QUEUE)
+
+ private val dataQueue = lbmq.getSubQueue(DATA_QUEUE)
+
+ private val controlQueue = lbmq.getSubQueue(CONTROL_QUEUE)
+
+ // the values in below maps are in batches
+ private val inQueueSizeMapping =
+ new mutable.HashMap[ChannelIdentity, Long]() // read and written by main thread
+ @volatile private var outQueueSizeMapping =
+ new mutable.HashMap[ChannelIdentity, Long]() // written by DP thread, read by main thread
+
+ def enqueueData(elem: InternalQueueElement): Unit = {
+ dataQueue.add(elem)
+ elem match {
+ case DataElement(dataPayload, from) =>
+ dataPayload match {
+ case frame: DataFrame =>
+ inQueueSizeMapping(from) =
+ inQueueSizeMapping.getOrElseUpdate(from, 0L) + frame.inMemSize
+ case _ =>
+ // do nothing
+ }
+ case _ =>
+ // do nothing
+ }
+ }
+
+ def enqueueCommand(cmd: DirectControlMessagePayload, from: ChannelIdentity): Unit = {
+ controlQueue.add(ControlElement(cmd, from))
+ }
+
+ def enqueueActorCommand(command: ActorCommand): Unit = {
+ controlQueue.add(ActorCommandElement(command))
+ }
+
+ def getElement: InternalQueueElement = {
+ val elem = lbmq.take()
+ elem match {
+ case DataElement(dataPayload, from) =>
+ dataPayload match {
+ case frame: DataFrame =>
+ outQueueSizeMapping(from) =
+ outQueueSizeMapping.getOrElseUpdate(from, 0L) + frame.inMemSize
+ case _ =>
+ // do nothing
+ }
+ case _ =>
+ // do nothing
+ }
+ elem
+ }
+
+ def disableDataQueue(): Unit = dataQueue.enable(false)
+
+ def enableDataQueue(): Unit = dataQueue.enable(true)
+
+ def getDataQueueLength: Int = dataQueue.size()
+
+ def getControlQueueLength: Int = controlQueue.size()
+
+ def isControlQueueEmpty: Boolean = controlQueue.isEmpty
+
+ def getQueuedCredit(sender: ChannelIdentity): Long = {
+ val inBytes = inQueueSizeMapping.getOrElseUpdate(sender, 0L)
+ val outBytes = outQueueSizeMapping.getOrElseUpdate(sender, 0L)
+ inBytes - outBytes
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/CostBasedScheduleGenerator.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/CostBasedScheduleGenerator.scala
new file mode 100644
index 00000000000..401ccddc0a4
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/CostBasedScheduleGenerator.scala
@@ -0,0 +1,663 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.storage.VFSURIFactory.createResultURI
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, PhysicalOpIdentity}
+import org.apache.texera.amber.core.workflow._
+import org.apache.texera.amber.engine.architecture.scheduling.SchedulingUtils.replaceVertex
+import org.apache.texera.amber.engine.architecture.scheduling.config.{
+ IntermediateInputPortConfig,
+ OutputPortConfig,
+ ResourceConfig
+}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.jgrapht.Graph
+import org.jgrapht.alg.connectivity.BiconnectivityInspector
+import org.jgrapht.graph.{DirectedAcyclicGraph, DirectedPseudograph}
+
+import java.net.URI
+import java.util.concurrent.TimeoutException
+import scala.collection.mutable
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration.DurationInt
+import scala.concurrent.{Await, Future}
+import scala.jdk.CollectionConverters._
+import scala.util.control.Breaks.{break, breakable}
+import scala.util.{Failure, Success, Try}
+
+class CostBasedScheduleGenerator(
+ workflowContext: WorkflowContext,
+ initialPhysicalPlan: PhysicalPlan,
+ val actorId: ActorVirtualIdentity
+) extends ScheduleGenerator(
+ workflowContext,
+ initialPhysicalPlan
+ )
+ with AmberLogging {
+
+ case class SearchResult(
+ state: Set[PhysicalLink],
+ regionDAG: DirectedAcyclicGraph[Region, RegionLink],
+ cost: Double,
+ searchTimeNanoSeconds: Long = 0,
+ numStatesExplored: Int = 0
+ )
+
+ private val costEstimator =
+ new DefaultCostEstimator(
+ workflowContext = workflowContext,
+ resourceAllocator = resourceAllocator,
+ actorId = actorId
+ )
+
+ private case class CostEstimatorMemoKey(
+ physicalOpIds: Set[PhysicalOpIdentity],
+ physicalLinkIds: Set[PhysicalLink],
+ portIds: Set[GlobalPortIdentity],
+ resourceConfig: Option[ResourceConfig]
+ )
+
+ private val costEstimatorMemoization
+ : mutable.Map[CostEstimatorMemoKey, (ResourceConfig, Double)] =
+ new mutable.HashMap()
+
+ def generate(): (Schedule, PhysicalPlan) = {
+ val startTime = System.nanoTime()
+ val regionDAG = createRegionDAG()
+ val totalRPGTime = System.nanoTime() - startTime
+ val regionPlan = RegionPlan(
+ regions = regionDAG.iterator().asScala.toSet,
+ regionLinks = regionDAG.edgeSet().asScala.toSet
+ )
+ val schedule = generateScheduleFromRegionPlan(regionPlan)
+ logger.info(
+ s"WID: ${workflowContext.workflowId.id}, EID: ${workflowContext.executionId.id}, total RPG time: " +
+ s"${totalRPGTime / 1e6} ms."
+ )
+ (
+ schedule,
+ physicalPlan
+ )
+ }
+
+ /**
+ * Partitions a physical plan into Regions and assigns storage URIs in two passes.
+ *
+ * Overview
+ *
+ * Region construction:
+ * Remove all materialized edges from the DAG and compute undirected connected
+ * components. The resulting “Region Graph” may contain directed cycles.
+ * Pass 1 – Output URIs:
+ * For each Region, allocate storage URIs on every output port of materialized edges.
+ * Pass 2 – Input URIs:
+ * Re-traverse the same Regions and attach reader URIs on input ports using
+ * the URIs created in Pass 1.
+ *
+ *
+ * Why two passes?
+ *
+ * Potential directed cycles in the Region Graph makes a topological
+ * traversal of regions inpossible.
+ * To ensure every output URI exists before its corresponding reader is assigned,
+ * and avoiding “reader before writer” holes, two passes are required.
+ *
+ *
+ * @param physicalPlan the original physical plan (without materializations)
+ * @param matEdges edges to be materialized (including blocking edges)
+ * @return a set of `Region`s whose `ResourceConfig` contains only `URI`s for `PortConfig`s
+ * (`Partitioning` to be assigned later in `ResourceAllocator`; see `IntermediateInputPortConfig`.)
+ */
+ private def createRegions(
+ physicalPlan: PhysicalPlan,
+ matEdges: Set[PhysicalLink]
+ ): Set[Region] = {
+
+ // Pass 0 – remove materialized edges and create connected components
+
+ val matEdgesRemovedDAG: PhysicalPlan = matEdges.foldLeft(physicalPlan)(_.removeLink(_))
+
+ val connectedComponents: Set[Graph[PhysicalOpIdentity, PhysicalLink]] =
+ new BiconnectivityInspector[PhysicalOpIdentity, PhysicalLink](
+ matEdgesRemovedDAG.dag
+ ).getConnectedComponents.asScala.toSet
+
+ // Pass 1 – build Regions only output-port storage URIs
+
+ val regionsWithOnlyOutputPortURIs: Set[Region] = connectedComponents.zipWithIndex.map {
+ case (connectedSubDAG, idx) =>
+ // Operators and intra‑region pipelined links
+
+ val operators: Set[PhysicalOpIdentity] = connectedSubDAG.vertexSet().asScala.toSet
+
+ val links: Set[PhysicalLink] = operators
+ .flatMap { opId =>
+ physicalPlan.getUpstreamPhysicalLinks(opId) ++
+ physicalPlan.getDownstreamPhysicalLinks(opId)
+ }
+ .filter(link => operators.contains(link.fromOpId))
+ .diff(matEdges) // keep only pipelined edges
+
+ val physicalOps: Set[PhysicalOp] = operators.map(physicalPlan.getOperator)
+
+ // Frontend-specified ports that need to be materailized (output ports of "eye-icon" physicalOps)
+ val outputPortIdsToViewResult: Set[GlobalPortIdentity] =
+ workflowContext.workflowSettings.outputPortsNeedingStorage
+ .filter(pid => operators.contains(pid.opId))
+
+ // Contains both frontend-specified and scheduler-decided ports that require materailizations.
+ val outputPortIdsNeedingStorage: Set[GlobalPortIdentity] =
+ matEdges
+ .filter(e => operators.contains(e.fromOpId))
+ .map(e => GlobalPortIdentity(e.fromOpId, e.fromPortId)) ++
+ outputPortIdsToViewResult
+
+ // Allocate an URI for each of these output ports
+ val outputPortConfigs: Map[GlobalPortIdentity, OutputPortConfig] =
+ outputPortIdsNeedingStorage.map { gpid =>
+ val outputWriterURI = createResultURI(
+ workflowId = workflowContext.workflowId,
+ executionId = workflowContext.executionId,
+ globalPortId = gpid
+ )
+ gpid -> OutputPortConfig(outputWriterURI)
+ }.toMap
+
+ val resourceConfig = ResourceConfig(portConfigs = outputPortConfigs)
+
+ // Enumerate all ports belonging to the Region
+ val ports: Set[GlobalPortIdentity] = physicalOps.flatMap { op =>
+ op.inputPorts.keys
+ .map(inputPortId => GlobalPortIdentity(op.id, inputPortId, input = true))
+ .toSet ++ op.outputPorts.keys
+ .map(outputPortId => GlobalPortIdentity(op.id, outputPortId))
+ .toSet
+ }
+
+ // Build the Region skeleton (no input‑port URIs yet)
+ Region(
+ id = RegionIdentity(idx),
+ physicalOps = physicalOps,
+ physicalLinks = links,
+ ports = ports,
+ resourceConfig = Some(resourceConfig)
+ )
+ }
+
+ // Collect writer‑side configs so we can look them up in Pass 2
+ val allOutputPortConfigs: Map[GlobalPortIdentity, OutputPortConfig] =
+ regionsWithOnlyOutputPortURIs
+ .flatMap(_.resourceConfig) // Seq[ResourceConfig]
+ .flatMap(_.portConfigs.collect { // PortConfig → OutputPortConfig
+ case (id, cfg: OutputPortConfig) => id -> cfg
+ })
+ .toMap
+
+ // Pass 2 – add input‑port storage configs (reader URIs)
+
+ regionsWithOnlyOutputPortURIs.map { existingRegion =>
+ // MatEdges that originally connected to the input ports of this region.
+ val relevantMatEdges: Set[PhysicalLink] = matEdges.filter { matEdge =>
+ existingRegion.getOperators.exists(_.id == matEdge.toOpId)
+ }
+
+ // Assign storage URIs to input ports of each materialized edge (each input port could have more than one URI)
+ val inputPortConfigs: Map[GlobalPortIdentity, IntermediateInputPortConfig] =
+ relevantMatEdges
+ .foldLeft(Map.empty[GlobalPortIdentity, List[URI]]) { (acc, link) =>
+ val globalOutputPortId = GlobalPortIdentity(link.fromOpId, link.fromPortId)
+ val globalInputPortId = GlobalPortIdentity(link.toOpId, link.toPortId, input = true)
+
+ // Writer‑side URI that must already exist thanks to Pass 1
+ val inputReaderURI = allOutputPortConfigs
+ .getOrElse(
+ globalOutputPortId,
+ throw new IllegalStateException(
+ s"Materialization edge $link: attempting to assign a materialization " +
+ s"reader URI for input port $globalInputPortId when " +
+ s"the outout port $globalOutputPortId has not been assigned a URI yet."
+ )
+ )
+ .storageURI
+
+ // Group all available URIs of this input port together
+ acc.updated(
+ globalInputPortId,
+ acc.getOrElse(globalInputPortId, List.empty[URI]) :+ inputReaderURI
+ )
+ }
+ .map {
+ case (inputPortId, uris) =>
+ inputPortId -> IntermediateInputPortConfig(uris)
+ }
+
+ val newResourceConfig: Option[ResourceConfig] = existingRegion.resourceConfig match {
+ case Some(existingConfig) =>
+ Some(ResourceConfig(portConfigs = existingConfig.portConfigs ++ inputPortConfigs))
+ case None =>
+ if (inputPortConfigs.nonEmpty) Some(ResourceConfig(portConfigs = inputPortConfigs))
+ else None
+ }
+
+ existingRegion.copy(resourceConfig = newResourceConfig)
+ }
+ }
+
+ /**
+ * Checks a plan for schedulability, and returns a region DAG if the plan is schedulable.
+ *
+ * @param matEdges Set of edges to materialize (including the original blocking edges).
+ * @return If the plan is schedulable, a region DAG will be returned. Otherwise a DirectedPseudograph (with directed
+ * cycles) will be returned to indicate that the plan is unschedulable.
+ */
+ private def tryConnectRegionDAG(
+ matEdges: Set[PhysicalLink]
+ ): Either[DirectedAcyclicGraph[Region, RegionLink], DirectedPseudograph[Region, RegionLink]] = {
+ val regionDAG = new DirectedAcyclicGraph[Region, RegionLink](classOf[RegionLink])
+ val regionGraph = new DirectedPseudograph[Region, RegionLink](classOf[RegionLink])
+ val opToRegionMap = new mutable.HashMap[PhysicalOpIdentity, Region]
+ createRegions(physicalPlan, matEdges).foreach(region => {
+ region.getOperators.foreach(op => opToRegionMap(op.id) = region)
+ regionGraph.addVertex(region)
+ regionDAG.addVertex(region)
+ })
+ var isAcyclic = true
+ matEdges.foreach(matEdge => {
+ val fromRegion = opToRegionMap(matEdge.fromOpId)
+ val toRegion = opToRegionMap(matEdge.toOpId)
+ regionGraph.addEdge(fromRegion, toRegion, RegionLink(fromRegion.id, toRegion.id))
+ try {
+ regionDAG.addEdge(fromRegion, toRegion, RegionLink(fromRegion.id, toRegion.id))
+ } catch {
+ case _: IllegalArgumentException =>
+ isAcyclic = false
+ }
+ })
+ if (isAcyclic) Left(regionDAG)
+ else Right(regionGraph)
+ }
+
+ /**
+ * Performs a search to generate a region DAG.
+ * Materializations are added only after the plan is determined to be schedulable.
+ *
+ * @return A region DAG.
+ */
+ private def createRegionDAG(): DirectedAcyclicGraph[Region, RegionLink] = {
+ val searchResultFuture: Future[SearchResult] = Future {
+ workflowContext.workflowSettings.executionMode match {
+ case ExecutionMode.MATERIALIZED =>
+ getFullyMaterializedSearchState
+ case ExecutionMode.PIPELINED =>
+ if (ApplicationConfig.useTopDownSearch)
+ topDownSearch(globalSearch = ApplicationConfig.useGlobalSearch)
+ else
+ bottomUpSearch(globalSearch = ApplicationConfig.useGlobalSearch)
+ }
+ }
+ val searchResult = Try(
+ Await.result(searchResultFuture, ApplicationConfig.searchTimeoutMilliseconds.milliseconds)
+ ) match {
+ case Failure(exception) =>
+ exception match {
+ case _: TimeoutException =>
+ logger.warn(
+ s"WID: ${workflowContext.workflowId.id}, EID: ${workflowContext.executionId.id}, search for region plan " +
+ s"timed out, falling back to bottom-up greedy search.",
+ exception
+ )
+ bottomUpSearch()
+ case _ => throw new RuntimeException(exception)
+ }
+
+ case Success(result) =>
+ result
+ }
+ logger.info(
+ s"WID: ${workflowContext.workflowId.id}, EID: ${workflowContext.executionId.id}, search time: " +
+ s"${searchResult.searchTimeNanoSeconds / 1e6} ms."
+ )
+
+ val regionDAG = searchResult.regionDAG
+ regionDAG
+ }
+
+ /**
+ * The core of the search algorithm. If the input physical plan is already schedulable, no search will be executed.
+ * Otherwise, depending on the configuration, either a global search or a greedy search will be performed to find
+ * an optimal plan. The search starts from a plan where all non-blocking edges are pipelined, and leads to a low-cost
+ * schedulable plan by changing pipelined non-blocking edges to materialized. By default all pruning techniques
+ * are enabled (chains, clean edges, and early stopping on schedulable states).
+ *
+ * @return A SearchResult containing the plan, the region DAG (without materializations added yet), the cost, the
+ * time to finish search, and the number of states explored.
+ */
+ def bottomUpSearch(
+ globalSearch: Boolean = false,
+ oChains: Boolean = true,
+ oCleanEdges: Boolean = true,
+ oEarlyStop: Boolean = true
+ ): SearchResult = {
+ val startTime = System.nanoTime()
+ val originalNonBlockingEdges =
+ if (oCleanEdges) {
+ physicalPlan.getNonBridgeNonBlockingLinks
+ } else {
+ physicalPlan.links.diff(
+ physicalPlan.getBlockingAndDependeeLinks
+ )
+ }
+ // Queue to hold states to be explored, starting with the empty set
+ val queue: mutable.Queue[Set[PhysicalLink]] = mutable.Queue(Set.empty[PhysicalLink])
+ // Keep track of visited states to avoid revisiting
+ val visited: mutable.Set[Set[PhysicalLink]] = mutable.Set.empty[Set[PhysicalLink]]
+ // Used for the Early Stop optimization technique
+ val schedulableStates: mutable.Set[Set[PhysicalLink]] = mutable.Set.empty[Set[PhysicalLink]]
+ // Initialize the bestResult with an impossible high cost for comparison
+ var bestResult: SearchResult = SearchResult(
+ state = Set.empty,
+ regionDAG = new DirectedAcyclicGraph[Region, RegionLink](classOf[RegionLink]),
+ cost = Double.PositiveInfinity
+ )
+
+ while (queue.nonEmpty) {
+ // A state is represented as a set of materialized non-blocking edges.
+ val currentState = queue.dequeue()
+ breakable {
+ if (
+ oEarlyStop && schedulableStates
+ .exists(ancestorState => ancestorState.subsetOf(currentState))
+ ) {
+ // Early stop: stopping exploring states beyond a schedulable state since the cost will only increase.
+ // A state X is a descendant of an ancestor state Y in the bottom-up search process if Y's set of materialized
+ // edges is a subset of that of X's (since X is reachable from Y by adding more materialized edges.)
+ break()
+ }
+ visited.add(currentState)
+ tryConnectRegionDAG(
+ physicalPlan.getBlockingAndDependeeLinks ++ currentState
+ ) match {
+ case Left(regionDAG) =>
+ updateOptimumIfApplicable(regionDAG)
+ addNeighborStatesToFrontier()
+ case Right(_) =>
+ addNeighborStatesToFrontier()
+ }
+ }
+
+ /**
+ * An internal method of bottom-up search that updates the current optimum if the examined state is schedulable
+ * and has a lower cost.
+ */
+ def updateOptimumIfApplicable(regionDAG: DirectedAcyclicGraph[Region, RegionLink]): Unit = {
+ if (oEarlyStop) schedulableStates.add(currentState)
+ // Calculate the current state's cost and update the bestResult if it's lower
+ val cost = allocateResourcesAndEvaluateCost(regionDAG)
+ if (cost < bestResult.cost) {
+ bestResult = SearchResult(currentState, regionDAG, cost)
+ }
+ }
+
+ /**
+ * An internal method of bottom-up search that performs state transitions (changing an pipelined edge to
+ * materialized) to include the unvisited neighbor(s) of the current state in the frontier (i.e., the queue).
+ * If using global search, all unvisited neighbors will be included. Otherwise in a greedy search, only the
+ * neighbor with the lowest cost will be included.
+ */
+ def addNeighborStatesToFrontier(): Unit = {
+ val allCurrentMaterializedEdges =
+ currentState ++ physicalPlan.getBlockingAndDependeeLinks
+ // Generate and enqueue all neighbour states that haven't been visited
+ var candidateEdges = originalNonBlockingEdges
+ .diff(currentState)
+ if (oChains) {
+ val edgesInChainWithMaterializedEdges = physicalPlan.maxChains
+ .filter(chain => chain.intersect(allCurrentMaterializedEdges).nonEmpty)
+ .flatten
+ candidateEdges = candidateEdges.diff(
+ edgesInChainWithMaterializedEdges
+ ) // Edges in chain with blocking edges should not be materialized
+ }
+
+ val unvisitedNeighborStates = candidateEdges
+ .map(edge => currentState + edge)
+ .filter(neighborState =>
+ !visited.contains(neighborState) && !queue.contains(neighborState)
+ )
+
+ val filteredNeighborStates = if (oEarlyStop) {
+ // Any descendant state of a schedulable state is not worth exploring.
+ unvisitedNeighborStates.filter(neighborState =>
+ !schedulableStates.exists(ancestorState => ancestorState.subsetOf(neighborState))
+ )
+ } else {
+ unvisitedNeighborStates
+ }
+
+ if (globalSearch) {
+ // include all unvisited neighbors
+ filteredNeighborStates.foreach(neighborState => queue.enqueue(neighborState))
+ } else {
+ // greedy search, only include an unvisited neighbor with the lowest cost
+ if (filteredNeighborStates.nonEmpty) {
+ val minCostNeighborState = filteredNeighborStates.minBy(neighborState =>
+ tryConnectRegionDAG(
+ physicalPlan.getBlockingAndDependeeLinks ++ neighborState
+ ) match {
+ case Left(regionDAG) =>
+ allocateResourcesAndEvaluateCost(regionDAG)
+ case Right(_) =>
+ Double.MaxValue
+ }
+ )
+ queue.enqueue(minCostNeighborState)
+ }
+ }
+ }
+ }
+
+ val searchTime = System.nanoTime() - startTime
+ bestResult.copy(
+ searchTimeNanoSeconds = searchTime,
+ numStatesExplored = visited.size
+ )
+ }
+
+ /** Constructs a baseline fully materialized region plan (one operator per region) and evaluates its cost. */
+ def getFullyMaterializedSearchState: SearchResult = {
+ val startTime = System.nanoTime()
+
+ val (regionDAG, cost) =
+ tryConnectRegionDAG(physicalPlan.links) match {
+ case Left(dag) => (dag, allocateResourcesAndEvaluateCost(dag))
+ case Right(_) =>
+ (
+ new DirectedAcyclicGraph[Region, RegionLink](classOf[RegionLink]),
+ Double.PositiveInfinity
+ )
+ }
+
+ SearchResult(
+ state = Set.empty,
+ regionDAG = regionDAG,
+ cost = cost,
+ searchTimeNanoSeconds = System.nanoTime() - startTime,
+ numStatesExplored = 1
+ )
+ }
+
+ /**
+ * Another direction to perform the search. Depending on the configuration, either a global search or a greedy search
+ * will be performed to find an optimal plan. The search starts from a plan where all edges are materialized, and
+ * leads to a low-cost schedulable plan by changing materialized non-blocking edges to pipelined.
+ * By default, all pruning techniques are enabled (chains, clean edges).
+ *
+ * @return A SearchResult containing the plan, the region DAG (without materializations added yet), the cost, the
+ * time to finish search, and the number of states explored.
+ */
+ def topDownSearch(
+ globalSearch: Boolean = false,
+ oChains: Boolean = true,
+ oCleanEdges: Boolean = true
+ ): SearchResult = {
+ val startTime = System.nanoTime()
+ // Starting from a state where all non-blocking edges are materialized
+ val originalSeedState = physicalPlan.links.diff(
+ physicalPlan.getBlockingAndDependeeLinks
+ )
+
+ // Chain optimization: an edge in the same chain as a blocking edge should not be materialized
+ val seedStateOptimizedByChainsIfApplicable = if (oChains) {
+ val edgesInChainWithBlockingEdge = physicalPlan.maxChains
+ .filter(chain => chain.intersect(physicalPlan.getBlockingAndDependeeLinks).nonEmpty)
+ .flatten
+ originalSeedState.diff(edgesInChainWithBlockingEdge)
+ } else {
+ originalSeedState
+ }
+
+ // Clean edge optimization: a clean edge should not be materialized
+ val finalSeedState = if (oCleanEdges) {
+ seedStateOptimizedByChainsIfApplicable.intersect(physicalPlan.getNonBridgeNonBlockingLinks)
+ } else {
+ seedStateOptimizedByChainsIfApplicable
+ }
+
+ // Queue to hold states to be explored, starting with the seed state
+ val queue: mutable.Queue[Set[PhysicalLink]] = mutable.Queue(finalSeedState)
+ // Keep track of visited states to avoid revisiting
+ val visited: mutable.Set[Set[PhysicalLink]] = mutable.Set.empty[Set[PhysicalLink]]
+ // Initialize the bestResult with an impossible high cost for comparison
+ var bestResult: SearchResult = SearchResult(
+ state = Set.empty,
+ regionDAG = new DirectedAcyclicGraph[Region, RegionLink](classOf[RegionLink]),
+ cost = Double.PositiveInfinity
+ )
+
+ while (queue.nonEmpty) {
+ val currentState = queue.dequeue()
+ visited.add(currentState)
+ tryConnectRegionDAG(
+ physicalPlan.getBlockingAndDependeeLinks ++ currentState
+ ) match {
+ case Left(regionDAG) =>
+ updateOptimumIfApplicable(regionDAG)
+ addNeighborStatesToFrontier()
+ // No need to explore further
+ case Right(_) =>
+ addNeighborStatesToFrontier()
+ }
+
+ /**
+ * An internal method of top-down search that updates the current optimum if the examined state is schedulable
+ * and has a lower cost.
+ */
+ def updateOptimumIfApplicable(regionDAG: DirectedAcyclicGraph[Region, RegionLink]): Unit = {
+ // Calculate the current state's cost and update the bestResult if it's lower
+ val cost = allocateResourcesAndEvaluateCost(regionDAG)
+ if (cost < bestResult.cost) {
+ bestResult = SearchResult(currentState, regionDAG, cost)
+ }
+ }
+
+ /**
+ * An internal method of top-down search that performs state transitions (changing an materialized edge to
+ * pipelined) to include the unvisited neighbor(s) of the current state in the frontier (i.e., the queue).
+ * If using global search, all unvisited neighbors will be included. Otherwise in a greedy search, only the
+ * neighbor with the lowest cost will be included.
+ */
+ def addNeighborStatesToFrontier(): Unit = {
+ val unvisitedNeighborStates = currentState
+ .map(edge => currentState - edge)
+ .filter(neighborState =>
+ !visited.contains(neighborState) && !queue.contains(neighborState)
+ )
+
+ if (globalSearch) {
+ // include all unvisited neighbors
+ unvisitedNeighborStates.foreach(neighborState => queue.enqueue(neighborState))
+ } else {
+ // greedy search, only include an unvisited neighbor with the lowest cost
+ if (unvisitedNeighborStates.nonEmpty) {
+ val minCostNeighborState = unvisitedNeighborStates.minBy(neighborState =>
+ tryConnectRegionDAG(
+ physicalPlan.getBlockingAndDependeeLinks ++ neighborState
+ ) match {
+ case Left(regionDAG) =>
+ allocateResourcesAndEvaluateCost(regionDAG)
+ case Right(_) =>
+ Double.MaxValue
+ }
+ )
+ queue.enqueue(minCostNeighborState)
+ }
+ }
+ }
+ }
+
+ val searchTime = System.nanoTime() - startTime
+ bestResult.copy(
+ searchTimeNanoSeconds = searchTime,
+ numStatesExplored = visited.size
+ )
+ }
+
+ /**
+ * Takes a region DAG, generates one or more (to be done in the future) schedules based on the region DAG, allocates
+ * resources to each region in the region DAG, and calculates the cost of the schedule(s) using Cost Estimator. Uses
+ * the cost of the best schedule (currently only considers one schedule) as the cost of the region DAG.
+ *
+ * @return A cost determined by the cost estimator.
+ */
+ private def allocateResourcesAndEvaluateCost(
+ regionDAG: DirectedAcyclicGraph[Region, RegionLink]
+ ): Double = {
+ val regionPlan =
+ RegionPlan(regionDAG.vertexSet().asScala.toSet, regionDAG.edgeSet().asScala.toSet)
+ val schedule = generateScheduleFromRegionPlan(regionPlan)
+ // In the future we may allow multiple regions in a level and split the resources.
+ schedule
+ .map(level =>
+ level
+ .map(region => {
+ val costEstimatorMemoKey = CostEstimatorMemoKey(
+ physicalOpIds = region.physicalOps.map(_.id),
+ physicalLinkIds = region.physicalLinks,
+ portIds = region.ports,
+ resourceConfig = region.resourceConfig
+ )
+ val (resourceConfig, regionCost) = costEstimatorMemoization
+ .getOrElseUpdate(
+ costEstimatorMemoKey,
+ costEstimator.allocateResourcesAndEstimateCost(region, 1)
+ )
+ // Update the region in the regionDAG to be the new region with resources allocated.
+ val regionWithResourceConfig = region.copy(resourceConfig = Some(resourceConfig))
+ replaceVertex(regionDAG, region, regionWithResourceConfig)
+ regionCost
+ })
+ .sum
+ )
+ .sum
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/CostEstimator.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/CostEstimator.scala
new file mode 100644
index 00000000000..d86101a1f06
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/CostEstimator.scala
@@ -0,0 +1,165 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflow.WorkflowContext
+import org.apache.texera.amber.engine.architecture.scheduling.DefaultCostEstimator.DEFAULT_OPERATOR_COST
+import org.apache.texera.amber.engine.architecture.scheduling.config.ResourceConfig
+import org.apache.texera.amber.engine.architecture.scheduling.resourcePolicies.ResourceAllocator
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.SqlServer.withTransaction
+import org.apache.texera.dao.jooq.generated.Tables.{WORKFLOW_EXECUTIONS, WORKFLOW_VERSION}
+
+import java.net.URI
+import scala.util.{Failure, Success, Try}
+
+/**
+ * A cost estimator should estimate a cost of running a region under the given resource constraints as units.
+ */
+trait CostEstimator {
+
+ /**
+ * Uses the given resource units to allocate resources to the region, and determine a cost based on the allocation.
+ *
+ * Note currently the ResourceAllocator is not cost-based and thus we use a cost model that does not rely on the
+ * allocator, i.e., the cost estimation process is external to the ResourceAllocator.
+ * @return A ResourceConfig for the region and an estimated cost of this region.
+ */
+ def allocateResourcesAndEstimateCost(region: Region, resourceUnits: Int): (ResourceConfig, Double)
+}
+
+object DefaultCostEstimator {
+ val DEFAULT_OPERATOR_COST: Double = 1.0
+}
+
+/**
+ * A default cost estimator using past statistics. If past statistics of a workflow are available, the cost of a region
+ * is the execution time of its longest-running operator. Otherwise the cost is the number of materialized ports in the
+ * region.
+ */
+class DefaultCostEstimator(
+ workflowContext: WorkflowContext,
+ val resourceAllocator: ResourceAllocator,
+ val actorId: ActorVirtualIdentity
+) extends CostEstimator
+ with AmberLogging {
+
+ // Requires mysql database to retrieve execution statistics, otherwise use number of materialized ports as a default.
+ private val operatorEstimatedTimeOption = Try(
+ this.getOperatorExecutionTimeInSeconds(
+ this.workflowContext.workflowId.id
+ )
+ ) match {
+ case Failure(_) => None
+ case Success(result) => result
+ }
+
+ operatorEstimatedTimeOption match {
+ case None =>
+ logger.info(
+ s"WID: ${workflowContext.workflowId.id}, EID: ${workflowContext.executionId.id}, " +
+ s"no past execution statistics available. Using number of materialized output ports as the cost. "
+ )
+ case Some(_) =>
+ }
+
+ override def allocateResourcesAndEstimateCost(
+ region: Region,
+ resourceUnits: Int
+ ): (ResourceConfig, Double) = {
+ // Currently the dummy cost from resourceAllocator is discarded.
+ val (resourceConfig, _) = resourceAllocator.allocate(region)
+ // We use a cost model that does not rely on the resource allocation.
+ // TODO: Once the ResourceAllocator actually calculates a cost, we can use its calculated cost.
+ val cost = this.operatorEstimatedTimeOption match {
+ case Some(operatorEstimatedTime) =>
+ // Use past statistics (wall-clock runtime). We use the execution time of the longest-running
+ // operator in each region to represent the region's execution time, and use the sum of all the regions'
+ // execution time as the wall-clock runtime of the workflow.
+ // This assumes a schedule is a total-order of the regions.
+ val opExecutionTimes = region.getOperators.map(op => {
+ operatorEstimatedTime.getOrElse(op.id.logicalOpId.id, DEFAULT_OPERATOR_COST)
+ })
+ val longestRunningOpExecutionTime = opExecutionTimes.max
+ longestRunningOpExecutionTime
+ case None =>
+ // Without past statistics (e.g., first execution), we use number of ports needing storage as the cost.
+ // Each port needing storage has a portConfig.
+ // This is independent of the schedule / resource allocator.
+ resourceConfig.portConfigs.size
+ }
+ (resourceConfig, cost)
+ }
+
+ /**
+ * Retrieve the latest successful execution to get statistics to calculate costs in DefaultCostEstimator.
+ * Using the total control processing time plus data processing time of an operator as its cost.
+ * If no past statistics are available (e.g., first execution), return None.
+ */
+ private def getOperatorExecutionTimeInSeconds(
+ wid: Long
+ ): Option[Map[String, Double]] = {
+
+ val uriString: String = withTransaction(
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ ) { context =>
+ context
+ .select(WORKFLOW_EXECUTIONS.RUNTIME_STATS_URI)
+ .from(WORKFLOW_EXECUTIONS)
+ .join(WORKFLOW_VERSION)
+ .on(WORKFLOW_VERSION.VID.eq(WORKFLOW_EXECUTIONS.VID))
+ .where(
+ WORKFLOW_VERSION.WID
+ .eq(wid.toInt)
+ .and(WORKFLOW_EXECUTIONS.STATUS.eq(3.toByte))
+ )
+ .orderBy(WORKFLOW_EXECUTIONS.STARTING_TIME.desc())
+ .limit(1)
+ .fetchOneInto(classOf[String])
+ }
+
+ if (uriString == null || uriString.isEmpty) {
+ None
+ } else {
+ val uri: URI = new URI(uriString)
+ val document = DocumentFactory.openDocument(uri)
+
+ val maxStats = document._1
+ .get()
+ .foldLeft(Map.empty[String, Double]) { (acc, tuple) =>
+ val record = tuple.asInstanceOf[Tuple]
+ val operatorId = record.getField(0).asInstanceOf[String]
+ val dataProcessingTime = record.getField(6).asInstanceOf[Long]
+ val controlProcessingTime = record.getField(7).asInstanceOf[Long]
+ val currentMaxTime = acc.getOrElse(operatorId, 0.0)
+ val newTime = (dataProcessingTime + controlProcessingTime) / 1e9
+ acc + (operatorId -> Math.max(currentMaxTime, newTime))
+ }
+
+ if (maxStats.isEmpty) None else Some(maxStats)
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/ExpansionGreedyScheduleGenerator.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/ExpansionGreedyScheduleGenerator.scala
new file mode 100644
index 00000000000..4bb89338967
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/ExpansionGreedyScheduleGenerator.scala
@@ -0,0 +1,495 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.core.storage.VFSURIFactory.createResultURI
+import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
+import org.apache.texera.amber.core.workflow.{
+ GlobalPortIdentity,
+ PhysicalLink,
+ PhysicalPlan,
+ WorkflowContext
+}
+import org.apache.texera.amber.engine.architecture.scheduling.SchedulingUtils.replaceVertex
+import org.apache.texera.amber.engine.architecture.scheduling.config.{
+ IntermediateInputPortConfig,
+ OutputPortConfig,
+ ResourceConfig
+}
+import org.jgrapht.alg.connectivity.BiconnectivityInspector
+import org.jgrapht.graph.DirectedAcyclicGraph
+import org.jgrapht.traverse.TopologicalOrderIterator
+
+import java.net.URI
+import scala.annotation.tailrec
+import scala.collection.mutable
+import scala.jdk.CollectionConverters.{CollectionHasAsScala, IteratorHasAsScala}
+
+@deprecated(
+ "This greedy schedule generator will be removed in the future. Use CostBasedScheduleGenerator instead."
+)
+class ExpansionGreedyScheduleGenerator(
+ workflowContext: WorkflowContext,
+ initialPhysicalPlan: PhysicalPlan
+) extends ScheduleGenerator(workflowContext, initialPhysicalPlan)
+ with LazyLogging {
+ def generate(): (Schedule, PhysicalPlan) = {
+
+ val regionDAG = createRegionDAG()
+ val regionPlan = RegionPlan(
+ regions = regionDAG.vertexSet().asScala.toSet,
+ regionLinks = regionDAG.edgeSet().asScala.toSet
+ )
+ val schedule = generateScheduleFromRegionPlan(regionPlan)
+
+ (
+ schedule,
+ physicalPlan
+ )
+ }
+
+ /**
+ * Takes in a pair of operatorIds, `upstreamOpId` and `downstreamOpId`, finds all regions they each
+ * belong to, and creates the order relationships between the Regions of upstreamOpId, with the Regions
+ * of downstreamOpId. The relation ship can be N to M.
+ *
+ * This method does not consider ports.
+ *
+ * Returns pairs of (upstreamRegion, downstreamRegion) indicating the order from
+ * upstreamRegion to downstreamRegion.
+ */
+ private def toRegionOrderPairs(
+ upstreamOpId: PhysicalOpIdentity,
+ downstreamOpId: PhysicalOpIdentity,
+ regionDAG: DirectedAcyclicGraph[Region, RegionLink]
+ ): Set[(Region, Region)] = {
+
+ val upstreamRegions = getRegions(upstreamOpId, regionDAG)
+ val downstreamRegions = getRegions(downstreamOpId, regionDAG)
+
+ upstreamRegions.flatMap { upstreamRegion =>
+ downstreamRegions
+ .filterNot(regionDAG.getDescendants(upstreamRegion).contains(_))
+ .map(downstreamRegion => (upstreamRegion, downstreamRegion))
+ }
+ }
+
+ /**
+ * Create Regions based on the PhysicalPlan. The Region are to be added to regionDAG separately.
+ */
+ private def createRegions(physicalPlan: PhysicalPlan): Set[Region] = {
+ val dependeeLinksRemovedDAG = physicalPlan.getDependeeLinksRemovedDAG
+ val connectedComponents = new BiconnectivityInspector[PhysicalOpIdentity, PhysicalLink](
+ dependeeLinksRemovedDAG.dag
+ ).getConnectedComponents.asScala.toSet
+ connectedComponents.zipWithIndex.map {
+ case (connectedSubDAG, idx) =>
+ val operatorIds = connectedSubDAG.vertexSet().asScala.toSet
+ val links = operatorIds
+ .flatMap(operatorId => {
+ physicalPlan.getUpstreamPhysicalLinks(operatorId) ++ physicalPlan
+ .getDownstreamPhysicalLinks(operatorId)
+ })
+ .filter(link => operatorIds.contains(link.fromOpId))
+ .diff(physicalPlan.getDependeeLinks) // dependee links should not belong to a region.
+ val operators = operatorIds.map(operatorId => physicalPlan.getOperator(operatorId))
+ val ports = operators.flatMap(op =>
+ op.inputPorts.keys
+ .map(inputPortId => GlobalPortIdentity(op.id, inputPortId, input = true))
+ .toSet ++ op.outputPorts.keys
+ .map(outputPortId => GlobalPortIdentity(op.id, outputPortId))
+ .toSet
+ )
+ Region(
+ id = RegionIdentity(idx),
+ physicalOps = operators,
+ physicalLinks = links,
+ ports = ports
+ )
+ }
+ }
+
+ /**
+ * Try connect the regions in the DAG while respecting the dependencies of PhysicalLinks (e.g., HashJoin).
+ * This function returns either a successful connected region DAG, or a list of PhysicalLinks that should be
+ * replaced for materialization.
+ *
+ * This function builds a region DAG from scratch. It first adds all the regions into the DAG. Then it starts adding
+ * edges on the DAG. To do so, it examines each PhysicalOp and checks its input links. The links will be problematic
+ * if the link's toOp (this PhysicalOp) has another link that has higher priority to run than this link (i.e., it has
+ * a dependency). If such links are found, the function will terminate after this PhysicalOp and return the set of
+ * links.
+ *
+ * If the function finds no such links for all PhysicalOps, it will return the connected Region DAG.
+ *
+ * @return Either a partially connected region DAG, or a set of PhysicalLinks for materialization replacement.
+ */
+ private def tryConnectRegionDAG()
+ : Either[DirectedAcyclicGraph[Region, RegionLink], Set[PhysicalLink]] = {
+
+ // creates an empty regionDAG
+ val regionDAG = new DirectedAcyclicGraph[Region, RegionLink](classOf[RegionLink])
+
+ // add Regions as vertices
+ createRegions(physicalPlan).foreach(region => regionDAG.addVertex(region))
+
+ // add regionLinks as edges, if failed, return the problematic PhysicalLinks.
+ physicalPlan
+ .topologicalIterator()
+ .foreach(physicalOpId => {
+ handleInputPortDependencies(physicalOpId, regionDAG)
+ .map(links => return Right(links))
+ })
+
+ // if success, a partially connected region DAG without edges between materialization operators is returned.
+ // The edges between materialization are to be added later.
+ Left(regionDAG)
+ }
+
+ /**
+ * A dependee input port is one that is depended on by another input port of the same operator.
+ * The incoming edge of a dependee input port is called a dependee edge.
+ * Similarly, the other port of this dependency relationship is called a depender input port and connects
+ * to a depender edge.
+ *
+ * Core design: a dependee edge needs to be materialized, and is mapped to a region edge in the region DAG.
+ * Note: currently we assume there CANNOT be dependencies between two dependee input ports.
+ * This method reasons about the input port dependencies of a given operator during the greedy expansion-based
+ * construction of a region DAG.
+ *
+ * This method first reasons about the dependencies of the input ports of the given operator to find
+ * pairs of dependency relationships, and then enforces the dependency of each pair:
+ * All the incoming edges of a dependee port will be added to the partial region DAG as a region edge.
+ * If adding a dependee edge results in a cycle that breaks the region DAG, we use a heurestic which is to
+ * return the other depender edge and indicate that this depender edge needs to be materialized. This will
+ * break the cycle and maintain the acyclicity of the region DAG.
+ *
+ * Previously we relied purely on edges and cache read operators for implementing materializations for
+ * materialized edges and find regions in this method.
+ *
+ * After introducing materailizations on output and input ports, materializing an
+ * edge could result in an operator that does not have any edges connected to one or more of its input
+ * ports (i.e., it becomes a "starter" operator in a region). For such input ports, we can only use port to
+ * find regions.
+ *
+ * @param physicalOpId The id of the input physical operator on which we need to handle input port dependies.
+ * @param regionDAG The partial region DAG that is always acyclic.
+ * @return Optionally a set of [[PhysicalLink]]s to do materialization-replacements on.
+ */
+ private def handleInputPortDependencies(
+ physicalOpId: PhysicalOpIdentity,
+ regionDAG: DirectedAcyclicGraph[Region, RegionLink]
+ ): Option[Set[PhysicalLink]] = {
+ // For operators like HashJoin's Probe that have dependencies between their input ports
+ physicalPlan
+ .getOperator(physicalOpId)
+ .getInputPortDependencyPairs
+ .sliding(2, 1)
+ .foreach {
+ case List(dependeePort, dependerPort) =>
+ // Create edges between regions
+ val dependeeEdges =
+ physicalPlan
+ .getUpstreamPhysicalLinks(physicalOpId)
+ .filter(l => l.toPortId == dependeePort)
+ val dependerEdges =
+ physicalPlan
+ .getUpstreamPhysicalLinks(physicalOpId)
+ .filter(l => l.toPortId == dependerPort)
+
+ if (dependerEdges.nonEmpty) {
+ // The depender port is connected to some edges of this same region
+ val regionOrderPairs =
+ toRegionOrderPairs(
+ dependeeEdges.head.fromOpId,
+ dependerEdges.head.fromOpId,
+ regionDAG
+ )
+ // Attempt to add these depender edges to regionDAG
+ try {
+ regionOrderPairs.foreach {
+ case (dependeeRegion, dependerRegion) =>
+ regionDAG.addEdge(
+ dependeeRegion,
+ dependerRegion,
+ RegionLink(dependeeRegion.id, dependerRegion.id)
+ )
+ }
+ } catch {
+ case _: IllegalArgumentException =>
+ // Adding the depender edge causes cycle. return the edge for materialization replacement
+ return Some(Set(dependerEdges.head))
+ }
+ } else {
+ // The depender port is not connected to any edges (due to materializations)
+ try {
+ // Any region that the dependee port belongs to needs to run first.
+ val dependeeRegions = getRegions(dependeeEdges.head.fromOpId, regionDAG)
+ // Any region that this depender port belongs to need to run after those dependee regions.
+ val dependerRegion = getRegions(physicalOpId, regionDAG)
+ .filter(region =>
+ region.getPorts.contains(
+ GlobalPortIdentity(
+ opId = physicalOpId,
+ portId = dependerPort,
+ input = true
+ )
+ )
+ )
+ .head
+ // We can safely add region edges created from this dependency relationship and it should
+ // never cause cycles (since the edges of this depender port are already "cut").
+ dependeeRegions.foreach(fromRegion =>
+ regionDAG
+ .addEdge(fromRegion, dependerRegion, RegionLink(fromRegion.id, dependerRegion.id))
+ )
+ } catch {
+ case _: IllegalArgumentException =>
+ // A cycle is detected. This logic should never be reached.
+ throw new WorkflowRuntimeException(
+ "Cyclic dependency when trying to handle input port dependencies in building a region plan"
+ )
+ }
+ }
+ case _ =>
+ }
+ None
+ }
+
+ /**
+ * Create `PortConfig`s containing only `URI`s for both input and output ports. For the greedy scheduler, this step
+ * after a region DAG is created.
+ */
+ private def assignPortConfigs(
+ matReaderWriterPairs: Set[(GlobalPortIdentity, GlobalPortIdentity)],
+ regionDAG: DirectedAcyclicGraph[Region, RegionLink]
+ ): Unit = {
+
+ val outputPortsToMaterialize = matReaderWriterPairs.map(_._1)
+
+ (outputPortsToMaterialize ++ workflowContext.workflowSettings.outputPortsNeedingStorage)
+ .foreach(outputPortId => {
+ getRegions(outputPortId.opId, regionDAG).foreach(fromRegion => {
+ val portConfigToAdd = outputPortId -> {
+ val uriToAdd = getStorageURIFromGlobalOutputPortId(outputPortId)
+ OutputPortConfig(uriToAdd)
+ }
+ val newResourceConfig = fromRegion.resourceConfig match {
+ case Some(existingConfig) =>
+ existingConfig.copy(portConfigs = existingConfig.portConfigs + portConfigToAdd)
+ case None => ResourceConfig(portConfigs = Map(portConfigToAdd))
+ }
+ val newFromRegion = fromRegion.copy(resourceConfig = Some(newResourceConfig))
+ replaceVertex(regionDAG, fromRegion, newFromRegion)
+ })
+ })
+
+ matReaderWriterPairs
+ // Group all pairs by the input port (_2)
+ .groupBy { case (_, inputPort) => inputPort }
+ // For each input port, build its PortConfig based on all its upstream output ports
+ .foreach {
+ case (inputPort, pairsForThisInput) =>
+ // Extract all the output ports paired with this input
+ val urisToAdd: List[URI] = pairsForThisInput.map {
+ case (outputPort, _) => getStorageURIFromGlobalOutputPortId(outputPort)
+ }.toList
+
+ val portConfigToAdd =
+ inputPort -> IntermediateInputPortConfig(urisToAdd)
+
+ getRegions(inputPort.opId, regionDAG).foreach(toRegion => {
+ val newResourceConfig = toRegion.resourceConfig match {
+ case Some(existingConfig) =>
+ existingConfig.copy(portConfigs = existingConfig.portConfigs + portConfigToAdd)
+ case None => ResourceConfig(portConfigs = Map(portConfigToAdd))
+ }
+ val newToRegion = toRegion.copy(resourceConfig = Some(newResourceConfig))
+ replaceVertex(regionDAG, toRegion, newToRegion)
+ })
+ }
+ }
+
+ private def getStorageURIFromGlobalOutputPortId(outputPortId: GlobalPortIdentity) = {
+ assert(!outputPortId.input)
+ createResultURI(
+ workflowId = workflowContext.workflowId,
+ executionId = workflowContext.executionId,
+ globalPortId = outputPortId
+ )
+ }
+
+ private def replaceLinkWithMaterialization(
+ physicalLink: PhysicalLink,
+ writerReaderPairs: mutable.Set[(GlobalPortIdentity, GlobalPortIdentity)]
+ ): PhysicalPlan = {
+ val outputGlobalPortId = GlobalPortIdentity(
+ physicalLink.fromOpId,
+ physicalLink.fromPortId
+ )
+
+ val inputGlobalPortId = GlobalPortIdentity(
+ physicalLink.toOpId,
+ physicalLink.toPortId,
+ input = true
+ )
+
+ val pair = (outputGlobalPortId, inputGlobalPortId)
+
+ writerReaderPairs += pair
+
+ val newPhysicalPlan = physicalPlan
+ .removeLink(physicalLink)
+ newPhysicalPlan
+ }
+
+ private def allocateResource(
+ regionDAG: DirectedAcyclicGraph[Region, RegionLink]
+ ): Unit = {
+ // generate the resource configs
+ new TopologicalOrderIterator(regionDAG).asScala
+ .foreach(region => {
+ val (resourceConfig, _) = resourceAllocator.allocate(region)
+ val regionWithResourceConfig = region.copy(resourceConfig = Some(resourceConfig))
+ replaceVertex(regionDAG, region, regionWithResourceConfig)
+ })
+ }
+
+ private def getRegions(
+ physicalOpId: PhysicalOpIdentity,
+ regionDAG: DirectedAcyclicGraph[Region, RegionLink]
+ ): Set[Region] = {
+ regionDAG
+ .vertexSet()
+ .asScala
+ .filter(region => region.getOperators.map(_.id).contains(physicalOpId))
+ .toSet
+ }
+
+ /**
+ * For a dependee input link, although it connects two regions A->B, we include this link and its toOp in region A
+ * so that the dependee link will be completed first.
+ */
+ private def populateDependeeLinks(
+ regionDAG: DirectedAcyclicGraph[Region, RegionLink]
+ ): Unit = {
+
+ val dependeeLinks = physicalPlan
+ .topologicalIterator()
+ .flatMap { physicalOpId =>
+ val upstreamPhysicalOpIds = physicalPlan.getUpstreamPhysicalOpIds(physicalOpId)
+ upstreamPhysicalOpIds.flatMap { upstreamPhysicalOpId =>
+ physicalPlan
+ .getLinksBetween(upstreamPhysicalOpId, physicalOpId)
+ .filter(link =>
+ physicalPlan
+ .getOperator(physicalOpId)
+ .isInputLinkDependee(link)
+ )
+ }
+ }
+ .toSet
+
+ dependeeLinks
+ .flatMap { link => getRegions(link.fromOpId, regionDAG).map(region => region -> link) }
+ .groupBy(_._1)
+ .view
+ .mapValues(_.map(_._2))
+ .foreach {
+ case (region, links) =>
+ val newRegion = region.copy(
+ physicalLinks = region.physicalLinks ++ links,
+ physicalOps =
+ region.getOperators ++ links.map(_.toOpId).map(id => physicalPlan.getOperator(id)),
+ ports = region.getPorts ++ links.map(dependeeLink =>
+ GlobalPortIdentity(dependeeLink.toOpId, dependeeLink.toPortId, input = true)
+ )
+ )
+ replaceVertex(regionDAG, region, newRegion)
+ }
+ }
+
+ /**
+ * This function creates and connects a region DAG while conducting materialization replacement.
+ * It keeps attempting to create a region DAG from the given PhysicalPlan. When failed, a list
+ * of PhysicalLinks that causes the failure will be given to conduct materialization replacement,
+ * which changes the PhysicalPlan. It keeps attempting with the updated PhysicalPLan until a
+ * region DAG is built after connecting materialized pairs.
+ *
+ * @return a fully connected region DAG.
+ */
+ private def createRegionDAG(): DirectedAcyclicGraph[Region, RegionLink] = {
+
+ val materializedOutputInputPortPairs =
+ new mutable.HashSet[(GlobalPortIdentity, GlobalPortIdentity)]()
+
+ @tailrec
+ def recConnectRegionDAG(): DirectedAcyclicGraph[Region, RegionLink] = {
+ tryConnectRegionDAG() match {
+ case Left(dag) => dag
+ case Right(links) =>
+ links.foreach { link =>
+ physicalPlan = replaceLinkWithMaterialization(
+ link,
+ materializedOutputInputPortPairs
+ )
+ }
+ recConnectRegionDAG()
+ }
+ }
+
+ // the region is partially connected successfully.
+ val regionDAG: DirectedAcyclicGraph[Region, RegionLink] = recConnectRegionDAG()
+
+ // also need to materialize all the dependee links.
+ physicalPlan.getDependeeLinks.foreach(link => {
+ physicalPlan = replaceLinkWithMaterialization(link, materializedOutputInputPortPairs)
+ })
+
+ // try to add dependencies between materialization writer and reader regions
+ try {
+ materializedOutputInputPortPairs.foreach {
+ case (upstreamOutputPort, downstreamInputPort) =>
+ toRegionOrderPairs(upstreamOutputPort.opId, downstreamInputPort.opId, regionDAG).foreach {
+ case (fromRegion, toRegion) =>
+ regionDAG.addEdge(fromRegion, toRegion, RegionLink(fromRegion.id, toRegion.id))
+ }
+ }
+ } catch {
+ case _: IllegalArgumentException =>
+ // a cycle is detected. it should not reach here.
+ throw new WorkflowRuntimeException(
+ "Cyclic dependency between regions detected"
+ )
+ }
+
+ assignPortConfigs(materializedOutputInputPortPairs.toSet, regionDAG)
+
+ // mark links that go to downstream regions
+ populateDependeeLinks(regionDAG)
+
+ // allocate resources on regions
+ allocateResource(regionDAG)
+
+ regionDAG
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/Region.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/Region.scala
new file mode 100644
index 00000000000..248a3ece3c1
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/Region.scala
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
+import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PhysicalLink, PhysicalOp}
+import org.apache.texera.amber.engine.architecture.scheduling.config.ResourceConfig
+import org.jgrapht.graph.{DefaultEdge, DirectedAcyclicGraph}
+import org.jgrapht.traverse.TopologicalOrderIterator
+
+import scala.jdk.CollectionConverters.IteratorHasAsScala
+
+case class RegionLink(fromRegionId: RegionIdentity, toRegionId: RegionIdentity)
+
+case class RegionIdentity(id: Long)
+case class Region(
+ id: RegionIdentity,
+ physicalOps: Set[PhysicalOp],
+ physicalLinks: Set[PhysicalLink],
+ ports: Set[GlobalPortIdentity] = Set.empty,
+ resourceConfig: Option[ResourceConfig] = None
+) {
+
+ private val operators: Map[PhysicalOpIdentity, PhysicalOp] =
+ getOperators.map(op => op.id -> op).toMap
+
+ @transient lazy val dag: DirectedAcyclicGraph[PhysicalOpIdentity, DefaultEdge] = {
+ val jgraphtDag = new DirectedAcyclicGraph[PhysicalOpIdentity, DefaultEdge](classOf[DefaultEdge])
+ getOperators.foreach(op => jgraphtDag.addVertex(op.id))
+ getLinks.foreach(link => jgraphtDag.addEdge(link.fromOpId, link.toOpId))
+ jgraphtDag
+ }
+
+ def topologicalIterator(): Iterator[PhysicalOpIdentity] = {
+ new TopologicalOrderIterator(dag).asScala
+ }
+
+ def getOperators: Set[PhysicalOp] = physicalOps
+
+ def getLinks: Set[PhysicalLink] = physicalLinks
+
+ /**
+ * Ideally ports should be derived from operators. However, as we are including an operator with a dependee input
+ * link in the previous region, such operator's other ports should not belong to the previous region. As a result
+ * ports of a regioin are saved separately.
+ * TODO: Improve this design once we have clean separation of regions.
+ */
+ def getPorts: Set[GlobalPortIdentity] = ports
+
+ def getOperator(physicalOpId: PhysicalOpIdentity): PhysicalOp = {
+ operators(physicalOpId)
+ }
+
+ /**
+ * Effective source operators in a region.
+ * The effective source contains operators that have 0 input links in this region.
+ */
+ def getSourceOperators: Set[PhysicalOp] = {
+ getOperators
+ .filter(physicalOp =>
+ physicalOp
+ .getInputLinks()
+ .map(link => link.fromOpId)
+ .forall(upstreamOpId => !getOperators.map(_.id).contains(upstreamOpId))
+ )
+
+ }
+
+ /**
+ * Operators that should be started first. An operator need to start first either because it is a source operator,
+ * or because it has an input port that needs to read from materialization.
+ */
+ def getStarterOperators: Set[PhysicalOp] = {
+ val opsReadingFromMaterialization = resourceConfig match {
+ case Some(config) =>
+ config.portConfigs
+ .filter {
+ case (globalPortId, config) =>
+ globalPortId.input && config.storageURIs.nonEmpty
+ }
+ .map {
+ case (globalPortId, _) => globalPortId.opId
+ }
+ .toSet
+ .map(opId => getOperator(opId))
+ case None => Set.empty
+ }
+ opsReadingFromMaterialization ++ getSourceOperators
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinator.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinator.scala
new file mode 100644
index 00000000000..254c16bf34b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinator.scala
@@ -0,0 +1,596 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.pekko.pattern.gracefulStop
+import com.twitter.util.{Duration => TwitterDuration, Future, JavaTimer, Return, Throw, Timer}
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.storage.VFSURIFactory.decodeURI
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PhysicalLink, PhysicalOp}
+import org.apache.texera.amber.engine.architecture.common.{
+ AkkaActorRefMappingService,
+ AkkaActorService,
+ ExecutorDeployment
+}
+import org.apache.texera.amber.engine.architecture.controller.execution.{
+ OperatorExecution,
+ RegionExecution,
+ WorkflowExecution
+}
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerConfig,
+ ExecutionStatsUpdate,
+ RuntimeStatisticsPersist,
+ WorkerAssignmentUpdate
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.scheduling.config.{
+ InputPortConfig,
+ OperatorConfig,
+ OutputPortConfig,
+ ResourceConfig
+}
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.Partitioning
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.FutureBijection._
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.web.SessionState
+import org.apache.texera.web.model.websocket.event.RegionStateEvent
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicReference
+import scala.concurrent.duration.{Duration => ScalaDuration}
+
+/**
+ * The executor of a region.
+ *
+ * We currently use a two-phase execution scheme to handle input-port dependency relationships. This is based on these
+ * assumptions:
+ *
+ * - We only allow input port dependencies where the input ports of a region can be grouped as two layers, with one
+ * layer of “dependee” ports and another layer of “depender” ports. We do not allow the case where an input port
+ * can both be a dependee and a depender.
+ * - We only allow depender ports to send data to output ports. Depenee input ports cannot send data to output ports.
+ * - All the physical operators must have output ports so that we can use the existence of output ports to decide
+ * whether to `FinalizeExecutor()` for a worker. (See `OutputManager.finalizeOutput()`)
+ *
+ * Under these assumptions, we can `syncStatusAndTransitionRegionExecutionPhase` for a region in this sequence:
+ *
+ * 0. `Unexecuted`
+ *
+ * 1. `ExecutingDependeePortsPhase`: All the dependee input ports are executed first until they complete.
+ * The corresponding workers of those input ports are also started in this phase. No output ports are allowed. If no
+ * dependee ports exist in a region, this first phase will be skipped.
+ *
+ * 2. `ExecutingNonDependeePortsPhase`: All other ports (non-dependee input ports, output ports) and
+ * their workers are executed. Region completion is indicated by the completion of all the ports when in this phase.
+ *
+ * 3. `Completed`
+ */
+class RegionExecutionCoordinator(
+ region: Region,
+ isRestart: Boolean,
+ workflowExecution: WorkflowExecution,
+ asyncRPCClient: AsyncRPCClient,
+ controllerConfig: ControllerConfig,
+ actorService: AkkaActorService,
+ actorRefService: AkkaActorRefMappingService
+) extends AmberLogging {
+
+ initRegionExecution()
+
+ private sealed trait RegionExecutionPhase
+ private case object Unexecuted extends RegionExecutionPhase
+ private case object ExecutingDependeePortsPhase extends RegionExecutionPhase
+ private case object ExecutingNonDependeePortsPhase extends RegionExecutionPhase
+ private case object Completed extends RegionExecutionPhase
+
+ private val currentPhaseRef: AtomicReference[RegionExecutionPhase] = new AtomicReference(
+ Unexecuted
+ )
+ private val terminationFutureRef: AtomicReference[Future[Unit]] = new AtomicReference(null)
+ private val killRetryTimer: Timer = new JavaTimer(true)
+ private val killRetryDelay: TwitterDuration = TwitterDuration.fromMilliseconds(200)
+
+ /**
+ * Sync the status of `RegionExecution` and transition this coordinator's phase to `Completed` only when the
+ * coordinator is currently in `ExecutingNonDependeePortsPhase`, all the ports of this region are completed, and
+ * all workers in this region are terminated.
+ *
+ * Additionally, this method will also terminate all the workers of this region:
+ *
+ * 1. An `EndWorker` control message is first sent to all the workers. This will be the last message each worker
+ * receives. We wait for all workers have replied to indicate they have finished processing all control messages.
+ *
+ * 2. Only after all workers have processed all control messages do we send a `gracefulStop` (pekko message) to each
+ * worker. JVM workers will be terminated by `gracefulStop`. Python proxy workes will also be terminated by
+ * `gracefulStop`, whose termination logic will also kill the PVMs.
+ */
+ private def tryCompleteRegionExecution(): Future[Unit] = {
+ // Only `ExecutingNonDependeePortsPhase` can transition to `Completed`
+ if (currentPhaseRef.get != ExecutingNonDependeePortsPhase) {
+ return Future.Unit
+ }
+
+ // Sync the status with RegionExecution
+ val regionExecution = workflowExecution.getRegionExecution(region.id)
+ if (!regionExecution.isCompleted) {
+ return Future.Unit
+ }
+
+ val existingTerminationFuture = terminationFutureRef.get
+ if (existingTerminationFuture != null) {
+ existingTerminationFuture
+ } else {
+ val terminationFuture = terminateWorkersWithRetry(regionExecution).flatMap { _ =>
+ // Set this coordinator's status to be completed so that subsequent regions can be started by
+ // WorkflowExecutionCoordinator.
+ setPhase(Completed)
+ Future.Unit
+ }
+ if (terminationFutureRef.compareAndSet(null, terminationFuture)) {
+ terminationFuture
+ } else {
+ terminationFutureRef.get
+ }
+ }
+ }
+
+ private def terminateWorkers(regionExecution: RegionExecution) = {
+ // 1. Send EndWorkers to every worker
+ val endWorkerRequests =
+ regionExecution.getAllOperatorExecutions.flatMap {
+ case (_, opExec) =>
+ opExec.getWorkerIds.map { workerId =>
+ asyncRPCClient.workerInterface
+ .endWorker(EmptyRequest(), asyncRPCClient.mkContext(workerId))
+ }
+ }.toSeq
+
+ val endWorkerFuture: Future[Unit] =
+ Future.collect(endWorkerRequests).unit
+
+ // 2. Send GracefulStops only after 1 has finished
+ val gracefulStopRequests: Future[Unit] =
+ endWorkerFuture.flatMap { _ =>
+ val gracefulStops =
+ regionExecution.getAllOperatorExecutions.flatMap {
+ case (_, opExec) =>
+ opExec.getWorkerIds.map { workerId =>
+ val actorRef = actorRefService.getActorRef(workerId)
+ // Remove the actorRef so that no other actors can find the worker and send messages.
+ actorRefService.removeActorRef(workerId)
+ // Restarted regions reuse actorId. Remove stale control channels so the
+ // controller does not reuse old control-message sequence numbers for new workers.
+ asyncRPCClient.inputGateway.removeControlChannel(workerId)
+ asyncRPCClient.outputGateway.removeControlChannel(workerId)
+ gracefulStop(actorRef, ScalaDuration(5, TimeUnit.SECONDS)).asTwitter()
+ }
+ }.toSeq
+
+ Future.collect(gracefulStops).unit
+ }
+
+ // 3. Log whether the kills were successful
+ gracefulStopRequests.transform {
+ case Return(_) =>
+ logger.info(s"Region ${region.id.id} successfully terminated.")
+ regionExecution.getAllOperatorExecutions.foreach {
+ case (_, opExec) =>
+ opExec.getWorkerIds.foreach { workerId =>
+ opExec.getWorkerExecution(workerId).update(System.nanoTime(), WorkerState.TERMINATED)
+ }
+ }
+ Future.Unit // propagate success
+ case Throw(err) =>
+ logger.warn(s"Error when terminating region ${region.id}.")
+ Future.exception(err) // propagate failure
+ }
+ }
+
+ private def terminateWorkersWithRetry(
+ regionExecution: RegionExecution,
+ attempt: Int = 1
+ ): Future[Unit] = {
+ terminateWorkers(regionExecution).rescue {
+ case err =>
+ logger.warn(
+ s"Failed to terminate region ${region.id.id} on attempt $attempt. Retrying in ${killRetryDelay.inMilliseconds} ms.",
+ err
+ )
+ Future
+ .sleep(killRetryDelay)(killRetryTimer)
+ .flatMap(_ => terminateWorkersWithRetry(regionExecution, attempt + 1))
+ }
+ }
+
+ def isCompleted: Boolean = currentPhaseRef.get == Completed
+
+ /**
+ * Returns the region termination future if termination has been initiated.
+ * This is only set by `tryCompleteRegionExecution()`.
+ */
+ def getTerminationFutureOpt: Option[Future[Unit]] = Option(terminationFutureRef.get)
+
+ /**
+ * This will sync and transition the region execution phase from one to another depending on its current phase:
+ *
+ * `Unexecuted` -> `ExecutingDependeePortsPhase` -> `ExecutingNonDependeePortsPhase` -> `Completed`
+ */
+ def syncStatusAndTransitionRegionExecutionPhase(): Future[Unit] =
+ currentPhaseRef.get match {
+ case Unexecuted =>
+ executeDependeePortPhase()
+ case ExecutingDependeePortsPhase =>
+ val regionExecution = workflowExecution.getRegionExecution(region.id)
+ if (
+ region.getOperators.forall { op =>
+ val operatorExecution = regionExecution.getOperatorExecution(op.id)
+ op.dependeeInputs.forall { dependeePortId =>
+ operatorExecution.isInputPortCompleted(dependeePortId)
+ }
+ }
+ ) {
+ // All dependee ports are completed. Can proceed with the next phase.
+ executeNonDependeePortPhase()
+ } else {
+ // Some dependee ports are still executing. Continue with this phase.
+ Future.Unit
+ }
+ case ExecutingNonDependeePortsPhase =>
+ tryCompleteRegionExecution()
+ case Completed =>
+ // Already completed, no further action needed.
+ Future.Unit
+ }
+
+ private def executeDependeePortPhase(): Future[Unit] = {
+ setPhase(ExecutingDependeePortsPhase)
+ if (!region.getOperators.exists(_.dependeeInputs.nonEmpty)) {
+ // Skip to the next phase when there are no dependee input ports
+ return syncStatusAndTransitionRegionExecutionPhase()
+ }
+ val ops = region.getOperators.filter(_.dependeeInputs.nonEmpty)
+
+ launchPhaseExecutionInternal(
+ ops,
+ () => assignPorts(region, isDependeePhase = true),
+ () => Future.value(Seq.empty),
+ () => sendStarts(region, isDependeePhase = true)
+ )
+ }
+
+ private def executeNonDependeePortPhase(): Future[Unit] = {
+ setPhase(ExecutingNonDependeePortsPhase)
+ // Allocate output port storage objects
+ region.resourceConfig.get.portConfigs
+ .collect {
+ case (id, cfg: OutputPortConfig) => id -> cfg
+ }
+ .foreach {
+ case (pid, cfg) =>
+ createOutputPortStorageObjects(Map(pid -> cfg))
+ }
+
+ val ops = region.getOperators.filter(_.dependeeInputs.isEmpty)
+
+ launchPhaseExecutionInternal(
+ ops,
+ () => assignPorts(region, isDependeePhase = false),
+ () => connectChannels(region.getLinks),
+ () => sendStarts(region, isDependeePhase = false)
+ )
+ }
+
+ /**
+ * Unified logic for launching either of the two phases asynchronously.
+ */
+ private def launchPhaseExecutionInternal(
+ operatorsToRun: Set[PhysicalOp],
+ assignPortsLogic: () => Future[Seq[EmptyReturn]],
+ connectChannelsLogic: () => Future[Seq[EmptyReturn]],
+ startWorkersLogic: () => Future[Seq[Unit]]
+ ): Future[Unit] = {
+
+ val resourceConfig = region.resourceConfig.get
+ val regionExecution = workflowExecution.getRegionExecution(region.id)
+
+ val stats = workflowExecution.getAllRegionExecutionsStats
+ asyncRPCClient.sendToClient(ExecutionStatsUpdate(stats))
+ asyncRPCClient.sendToClient(RuntimeStatisticsPersist(stats))
+ asyncRPCClient.sendToClient(
+ WorkerAssignmentUpdate(
+ operatorsToRun
+ .map(_.id)
+ .map { pid =>
+ pid.logicalOpId.id -> regionExecution
+ .getOperatorExecution(pid)
+ .getWorkerIds
+ .map(_.name)
+ .toList
+ }
+ .toMap
+ )
+ )
+ Future(())
+ .flatMap(_ => initExecutors(operatorsToRun, resourceConfig))
+ .flatMap(_ => assignPortsLogic())
+ .flatMap(_ => connectChannelsLogic())
+ .flatMap(_ => openOperators(operatorsToRun))
+ .flatMap(_ => startWorkersLogic())
+ .unit
+ }
+
+ /**
+ * Initialize the execution states of all the operators in the region, and also create workers for each operator.
+ */
+ private def initRegionExecution(): Unit = {
+ val resourceConfig = region.resourceConfig.get
+ val regionExecution = workflowExecution.getRegionExecution(region.id)
+
+ region.getOperators.foreach { physicalOp =>
+ val existOpExecution =
+ workflowExecution.getAllRegionExecutions.exists(_.hasOperatorExecution(physicalOp.id))
+
+ val operatorExecution = regionExecution.initOperatorExecution(
+ physicalOp.id,
+ if (existOpExecution)
+ Some(workflowExecution.getLatestOperatorExecution(physicalOp.id))
+ else
+ None
+ )
+
+ if (!existOpExecution) {
+ buildOperator(
+ actorService,
+ physicalOp,
+ resourceConfig.operatorConfigs(physicalOp.id),
+ operatorExecution
+ )
+ }
+ }
+ }
+
+ private def buildOperator(
+ actorService: AkkaActorService,
+ physicalOp: PhysicalOp,
+ operatorConfig: OperatorConfig,
+ operatorExecution: OperatorExecution
+ ): Unit = {
+ ExecutorDeployment.createWorkers(
+ physicalOp,
+ actorService,
+ operatorExecution,
+ operatorConfig,
+ controllerConfig.stateRestoreConfOpt,
+ controllerConfig.faultToleranceConfOpt
+ )
+ }
+
+ private def initExecutors(
+ operators: Set[PhysicalOp],
+ resourceConfig: ResourceConfig
+ ): Future[Seq[EmptyReturn]] = {
+ Future
+ .collect(
+ operators
+ .flatMap(physicalOp => {
+ val workerConfigs = resourceConfig.operatorConfigs(physicalOp.id).workerConfigs
+ workerConfigs.map(_.workerId).map { workerId =>
+ asyncRPCClient.workerInterface.initializeExecutor(
+ InitializeExecutorRequest(
+ workerConfigs.length,
+ physicalOp.opExecInitInfo,
+ physicalOp.isSourceOperator
+ ),
+ asyncRPCClient.mkContext(workerId)
+ )
+ }
+ })
+ .toSeq
+ )
+ }
+
+ private def assignPorts(
+ region: Region,
+ isDependeePhase: Boolean
+ ): Future[Seq[EmptyReturn]] = {
+ val resourceConfig = region.resourceConfig.get
+ Future.collect(
+ region.getOperators
+ .flatMap { physicalOp: PhysicalOp =>
+ // assign input ports
+ val inputPortMapping = physicalOp.inputPorts
+ .filter {
+ case (portId, _) =>
+ // keep only the ports that belong to the requested phase
+ isDependeePhase == physicalOp.dependeeInputs.contains(portId)
+ }
+ .flatMap {
+ case (inputPortId, (_, _, Right(schema))) =>
+ val globalInputPortId = GlobalPortIdentity(physicalOp.id, inputPortId, input = true)
+ val (storageURIs, partitionings) =
+ resourceConfig.portConfigs.get(globalInputPortId) match {
+ case Some(cfg: InputPortConfig) =>
+ (cfg.storagePairs.map(_._1.toString), cfg.storagePairs.map(_._2))
+ case _ => (List.empty[String], List.empty[Partitioning])
+ }
+ Some(globalInputPortId -> (storageURIs, partitionings, schema))
+ case _ => None
+ }
+
+ // Currently an output port uses the same AssignPortRequest as an Input port.
+ // However, an output port does not need a list of URIs or partitionings.
+ // TODO: Separate AssignPortRequest for Input and Output Ports
+
+ // assign output ports (only for non-dependee phase)
+ val outputPortMapping =
+ if (isDependeePhase) {
+ Iterable.empty
+ } else {
+ physicalOp.outputPorts
+ .filter {
+ case (outputPortId, _) =>
+ val globalInputPortId = GlobalPortIdentity(physicalOp.id, outputPortId)
+ region.getPorts.contains(globalInputPortId)
+ }
+ .flatMap {
+ case (outputPortId, (_, _, Right(schema))) =>
+ val storageURI = resourceConfig.portConfigs
+ .collectFirst {
+ case (gid, cfg: OutputPortConfig)
+ if gid == GlobalPortIdentity(
+ opId = physicalOp.id,
+ portId = outputPortId
+ ) =>
+ cfg.storageURI.toString
+ }
+ .getOrElse("")
+ Some(
+ GlobalPortIdentity(physicalOp.id, outputPortId) -> (List(
+ storageURI
+ ), List.empty, schema)
+ )
+ case _ => None
+ }
+ }
+
+ inputPortMapping ++ outputPortMapping
+ }
+ // Issue AssignPort control messages to each worker.
+ .flatMap {
+ case (globalPortId, (storageUris, partitionings, schema)) =>
+ resourceConfig.operatorConfigs(globalPortId.opId).workerConfigs.map(_.workerId).map {
+ workerId =>
+ asyncRPCClient.workerInterface.assignPort(
+ AssignPortRequest(
+ globalPortId.portId,
+ globalPortId.input,
+ schema.toRawSchema,
+ storageUris,
+ partitionings
+ ),
+ asyncRPCClient.mkContext(workerId)
+ )
+ }
+ }
+ .toSeq
+ )
+ }
+
+ private def connectChannels(links: Set[PhysicalLink]): Future[Seq[EmptyReturn]] = {
+ Future.collect(
+ links.map { link: PhysicalLink =>
+ asyncRPCClient.controllerInterface.linkWorkers(
+ LinkWorkersRequest(link),
+ asyncRPCClient.mkContext(CONTROLLER)
+ )
+ }.toSeq
+ )
+ }
+
+ private def openOperators(operators: Set[PhysicalOp]): Future[Seq[EmptyReturn]] = {
+ Future
+ .collect(
+ operators
+ .map(_.id)
+ .flatMap(opId =>
+ workflowExecution.getRegionExecution(region.id).getOperatorExecution(opId).getWorkerIds
+ )
+ .map { workerId =>
+ asyncRPCClient.workerInterface
+ .openExecutor(EmptyRequest(), asyncRPCClient.mkContext(workerId))
+ }
+ .toSeq
+ )
+ }
+
+ private def sendStarts(
+ region: Region,
+ isDependeePhase: Boolean
+ ): Future[Seq[Unit]] = {
+ val stats = workflowExecution.getAllRegionExecutionsStats
+ asyncRPCClient.sendToClient(ExecutionStatsUpdate(stats))
+ asyncRPCClient.sendToClient(RuntimeStatisticsPersist(stats))
+ val allStarterOperators = region.getStarterOperators
+ val starterOpsForThisPhase =
+ if (isDependeePhase) allStarterOperators.filter(_.dependeeInputs.nonEmpty)
+ else allStarterOperators
+ Future.collect(
+ starterOpsForThisPhase
+ .map(_.id)
+ .flatMap { opId =>
+ workflowExecution
+ .getRegionExecution(region.id)
+ .getOperatorExecution(opId)
+ .getWorkerIds
+ .map { workerId =>
+ asyncRPCClient.workerInterface
+ .startWorker(EmptyRequest(), asyncRPCClient.mkContext(workerId))
+ .map(resp =>
+ // update worker state
+ workflowExecution
+ .getRegionExecution(region.id)
+ .getOperatorExecution(opId)
+ .getWorkerExecution(workerId)
+ .update(System.nanoTime(), resp.state)
+ )
+ }
+ }
+ .toSeq
+ )
+ }
+
+ private def createOutputPortStorageObjects(
+ portConfigs: Map[GlobalPortIdentity, OutputPortConfig]
+ ): Unit = {
+ portConfigs.foreach {
+ case (outputPortId, portConfig) =>
+ val storageUriToAdd = portConfig.storageURI
+ val (_, eid, _, _) = decodeURI(storageUriToAdd)
+ val schemaOptional =
+ region.getOperator(outputPortId.opId).outputPorts(outputPortId.portId)._3
+ val schema =
+ schemaOptional.getOrElse(throw new IllegalStateException("Schema is missing"))
+ DocumentFactory.createDocument(storageUriToAdd, schema)
+ if (!isRestart) {
+ WorkflowExecutionsResource.insertOperatorPortResultUri(
+ eid = eid,
+ globalPortId = outputPortId,
+ uri = storageUriToAdd
+ )
+ }
+ }
+ }
+
+ private def setPhase(phase: RegionExecutionPhase): Unit = {
+ currentPhaseRef.set(phase)
+ SessionState.getAllSessionStates.foreach { state =>
+ state.send(RegionStateEvent(region.id.id, phase.toString))
+ }
+ }
+
+ override def actorId: ActorVirtualIdentity = CONTROLLER
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionPlan.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionPlan.scala
new file mode 100644
index 00000000000..95937517de2
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionPlan.scala
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PhysicalLink}
+import org.jgrapht.graph.DirectedAcyclicGraph
+import org.jgrapht.traverse.TopologicalOrderIterator
+
+import scala.jdk.CollectionConverters.IteratorHasAsScala
+
+case class RegionPlan(
+ regions: Set[Region],
+ regionLinks: Set[RegionLink]
+) {
+
+ @transient lazy val dag: DirectedAcyclicGraph[RegionIdentity, RegionLink] = {
+ val jgraphtDag = new DirectedAcyclicGraph[RegionIdentity, RegionLink](classOf[RegionLink])
+ regionMapping.keys.foreach(regionId => jgraphtDag.addVertex(regionId))
+ regionLinks.foreach(l => jgraphtDag.addEdge(l.fromRegionId, l.toRegionId, l))
+ jgraphtDag
+ }
+ @transient private lazy val regionMapping: Map[RegionIdentity, Region] =
+ regions.map(region => region.id -> region).toMap
+
+ def getRegionOfLink(link: PhysicalLink): Region = {
+ regions.find(region => region.getLinks.contains(link)).get
+ }
+
+ def getRegionOfPortId(portId: GlobalPortIdentity): Option[Region] = {
+ regions.find(region => region.getPorts.contains(portId))
+ }
+
+ def topologicalIterator(): Iterator[RegionIdentity] = {
+ new TopologicalOrderIterator(dag).asScala
+ }
+
+ def getRegion(regionId: RegionIdentity): Region = {
+ regionMapping(regionId)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/Schedule.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/Schedule.scala
new file mode 100644
index 00000000000..6f34c9ed1e5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/Schedule.scala
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+case class Schedule(private val levelSets: Map[Int, Set[Region]]) extends Iterator[Set[Region]] {
+ private var currentLevel = levelSets.keys.minOption.getOrElse(0)
+
+ def getRegions: List[Region] = levelSets.values.flatten.toList
+
+ override def hasNext: Boolean = levelSets.isDefinedAt(currentLevel)
+
+ override def next(): Set[Region] = {
+ val regions = levelSets(currentLevel)
+ currentLevel += 1
+ regions
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/ScheduleGenerator.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/ScheduleGenerator.scala
new file mode 100644
index 00000000000..a5748fc71b1
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/ScheduleGenerator.scala
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.workflow._
+import org.apache.texera.amber.engine.architecture.scheduling.resourcePolicies.{
+ DefaultResourceAllocator,
+ ExecutionClusterInfo
+}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+abstract class ScheduleGenerator(
+ workflowContext: WorkflowContext,
+ var physicalPlan: PhysicalPlan
+) {
+ private val executionClusterInfo = new ExecutionClusterInfo()
+ val resourceAllocator =
+ new DefaultResourceAllocator(
+ physicalPlan,
+ executionClusterInfo,
+ workflowContext.workflowSettings
+ )
+
+ def generate(): (Schedule, PhysicalPlan)
+
+ /**
+ * A schedule is a ranking on the regions of a region plan.
+ * Regions are dispatched in batches of up to AmberConfig.maxConcurrentRegions, respecting the DAG dependencies.
+ * When maxConcurrentRegions == 1, this is equivalent to a total order (fully sequential execution).
+ */
+ def generateScheduleFromRegionPlan(regionPlan: RegionPlan): Schedule = {
+ val inDegree = mutable.Map.empty[RegionIdentity, Int]
+ regionPlan.topologicalIterator().foreach { rid =>
+ inDegree(rid) = regionPlan.dag.incomingEdgesOf(rid).asScala.size
+ }
+
+ val readyRegionsQueue = mutable.Queue(
+ inDegree.collect { case (rid, 0) => rid }.toSeq: _*
+ )
+
+ val tmpLevelSets = mutable.Map.empty[Int, Set[RegionIdentity]]
+ var level = 0
+
+ // While there are ready regions:
+ // 1. Dequeue up to maxConcurrentRegions regions to form the current batch.
+ // 2. Record this batch under the current level.
+ // 3. For each region in the batch:
+ // a. For each successor region, decrement its in-degree.
+ // b. If a successor's in-degree reaches zero, enqueue it to the readyRegionsQueue.
+ // 4. Increment level and repeat.
+ while (readyRegionsQueue.nonEmpty) {
+ val batchIds = (1 to ApplicationConfig.maxConcurrentRegions).flatMap { _ =>
+ if (readyRegionsQueue.nonEmpty) Some(readyRegionsQueue.dequeue()) else None
+ }.toSet
+
+ tmpLevelSets(level) = batchIds
+ batchIds.foreach { rid =>
+ regionPlan.dag
+ .outgoingEdgesOf(rid)
+ .asScala
+ .map(edge => regionPlan.dag.getEdgeTarget(edge))
+ .foreach { succ =>
+ inDegree(succ) -= 1
+ if (inDegree(succ) == 0) readyRegionsQueue.enqueue(succ)
+ }
+ }
+ level += 1
+ }
+ val levelSets: Map[Int, Set[Region]] = tmpLevelSets.view.map {
+ case (lvl, idSet) =>
+ lvl -> idSet.iterator.map(regionPlan.getRegion).toSet
+ }.toMap
+ Schedule(levelSets)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/SchedulingUtils.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/SchedulingUtils.scala
new file mode 100644
index 00000000000..4a287d6c6fd
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/SchedulingUtils.scala
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.jgrapht.graph.DirectedAcyclicGraph
+
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+object SchedulingUtils {
+ // TODO: remove this function
+ def replaceVertex(
+ graph: DirectedAcyclicGraph[Region, RegionLink],
+ oldVertex: Region,
+ newVertex: Region
+ ): Unit = {
+ if (oldVertex.equals(newVertex)) {
+ return
+ }
+ graph.addVertex(newVertex)
+ graph
+ .outgoingEdgesOf(oldVertex)
+ .asScala
+ .toList
+ .foreach(oldEdge => {
+ val dest = graph.getEdgeTarget(oldEdge)
+ graph.removeEdge(oldEdge)
+ graph.addEdge(newVertex, dest, RegionLink(newVertex.id, dest.id))
+ })
+ graph
+ .incomingEdgesOf(oldVertex)
+ .asScala
+ .toList
+ .foreach(oldEdge => {
+ val source = graph.getEdgeSource(oldEdge)
+ graph.removeEdge(oldEdge)
+ graph.addEdge(source, newVertex, RegionLink(source.id, newVertex.id))
+ })
+ graph.removeVertex(oldVertex)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinator.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinator.scala
new file mode 100644
index 00000000000..4b639fc241e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinator.scala
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import com.twitter.util.Future
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PhysicalLink}
+import org.apache.texera.amber.engine.architecture.common.{
+ AkkaActorRefMappingService,
+ AkkaActorService
+}
+import org.apache.texera.amber.engine.architecture.controller.ControllerConfig
+import org.apache.texera.amber.engine.architecture.controller.ExecutionStateUpdate
+import org.apache.texera.amber.engine.architecture.controller.execution.WorkflowExecution
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient
+
+import java.util.concurrent.atomic.AtomicBoolean
+import scala.collection.mutable
+
+class WorkflowExecutionCoordinator(
+ getNextRegions: () => Set[Region],
+ workflowExecution: WorkflowExecution,
+ controllerConfig: ControllerConfig,
+ asyncRPCClient: AsyncRPCClient
+) extends LazyLogging {
+
+ private val executedRegions: mutable.ListBuffer[Set[Region]] = mutable.ListBuffer()
+
+ private val regionExecutionCoordinators
+ : mutable.HashMap[RegionIdentity, RegionExecutionCoordinator] =
+ mutable.HashMap()
+ private val completionNotified: AtomicBoolean = new AtomicBoolean(false)
+
+ @transient var actorRefService: AkkaActorRefMappingService = _
+
+ def setupActorRefService(actorRefService: AkkaActorRefMappingService): Unit = {
+ this.actorRefService = actorRefService
+ }
+
+ /**
+ * Each invocation first syncs the internal statuses of each exisiting `RegionExecutionCoordintor`, after which each
+ * of the `RegionExecutionCoordintor`s will launch the corresponding next phase of whenever needed until it is
+ * in `Completed` status (phase).
+ *
+ * After the syncs, if there are no running region(s), it will start new regions (if available).
+ */
+ def coordinateRegionExecutors(actorService: AkkaActorService): Future[Unit] = {
+ val unfinishedRegionCoordinators =
+ regionExecutionCoordinators.values.filter(!_.isCompleted).toSeq
+
+ // Trigger sync for each unfinished region.
+ unfinishedRegionCoordinators.foreach(_.syncStatusAndTransitionRegionExecutionPhase())
+
+ // Wait only for region termination futures (kill path), then re-run coordination.
+ val terminationFutures = unfinishedRegionCoordinators.flatMap(_.getTerminationFutureOpt)
+ if (terminationFutures.nonEmpty) {
+ return Future
+ .collect(terminationFutures)
+ .unit
+ .flatMap(_ => coordinateRegionExecutors(actorService))
+ }
+
+ if (regionExecutionCoordinators.values.exists(!_.isCompleted)) {
+ // Some regions are still not completed yet. Cannot start the new regions.
+ return Future.Unit
+ }
+
+ // All existing regions are completed. Start the next region (if any).
+ val nextRegions = getNextRegions()
+ if (nextRegions.isEmpty) {
+ if (workflowExecution.isCompleted && completionNotified.compareAndSet(false, true)) {
+ asyncRPCClient.sendToClient(ExecutionStateUpdate(workflowExecution.getState))
+ }
+ return Future.Unit
+ }
+
+ executedRegions.append(nextRegions)
+ Future
+ .collect(
+ nextRegions
+ .map(region => {
+ val isRestart = workflowExecution.hasRegionExecution(region.id)
+ if (isRestart) {
+ workflowExecution.restartRegionExecution(region)
+ } else {
+ workflowExecution.initRegionExecution(region)
+ }
+ regionExecutionCoordinators(region.id) = new RegionExecutionCoordinator(
+ region,
+ isRestart,
+ workflowExecution,
+ asyncRPCClient,
+ controllerConfig,
+ actorService,
+ actorRefService
+ )
+ regionExecutionCoordinators(region.id)
+ })
+ .map(_.syncStatusAndTransitionRegionExecutionPhase())
+ .toSeq
+ )
+ .unit
+ }
+
+ def getRegionOfLink(link: PhysicalLink): Region = {
+ getExecutingRegions.find(region => region.getLinks.contains(link)).get
+ }
+
+ def getRegionOfPortId(portId: GlobalPortIdentity): Option[Region] = {
+ getExecutingRegions.find(region => region.getPorts.contains(portId))
+ }
+
+ def getExecutingRegions: Set[Region] = {
+ executedRegions.flatten
+ .filterNot(region => workflowExecution.getRegionExecution(region.id).isCompleted)
+ .toSet
+ }
+
+ def hasUnfinishedRegionCoordinators: Boolean = {
+ regionExecutionCoordinators.values.exists(!_.isCompleted)
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/ChannelConfig.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/ChannelConfig.scala
new file mode 100644
index 00000000000..83d4ed985a3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/ChannelConfig.scala
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling.config
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow._
+
+case object ChannelConfig {
+ def generateChannelConfigs(
+ fromWorkerIds: List[ActorVirtualIdentity],
+ toWorkerIds: List[ActorVirtualIdentity],
+ toPortId: PortIdentity,
+ partitionInfo: PartitionInfo
+ ): List[ChannelConfig] = {
+ partitionInfo match {
+ case HashPartition(_) | RangePartition(_, _, _) | BroadcastPartition() | UnknownPartition() =>
+ fromWorkerIds.flatMap(fromWorkerId =>
+ toWorkerIds.map(toWorkerId =>
+ ChannelConfig(ChannelIdentity(fromWorkerId, toWorkerId, isControl = false), toPortId)
+ )
+ )
+
+ case SinglePartition() =>
+ assert(toWorkerIds.size == 1)
+ val toWorkerId = toWorkerIds.head
+ fromWorkerIds.map(fromWorkerId =>
+ ChannelConfig(ChannelIdentity(fromWorkerId, toWorkerId, isControl = false), toPortId)
+ )
+ case OneToOnePartition() =>
+ fromWorkerIds.zip(toWorkerIds).map {
+ case (fromWorkerId, toWorkerId) =>
+ ChannelConfig(ChannelIdentity(fromWorkerId, toWorkerId, isControl = false), toPortId)
+ }
+ case _ =>
+ List()
+
+ }
+ }
+}
+
+case class ChannelConfig(
+ channelId: ChannelIdentity,
+ toPortId: PortIdentity
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/LinkConfig.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/LinkConfig.scala
new file mode 100644
index 00000000000..ef0117834cc
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/LinkConfig.scala
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling.config
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow._
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings._
+
+case object LinkConfig {
+ def toPartitioning(
+ fromWorkerIds: List[ActorVirtualIdentity],
+ toWorkerIds: List[ActorVirtualIdentity],
+ partitionInfo: PartitionInfo,
+ dataTransferBatchSize: Int
+ ): Partitioning = {
+ partitionInfo match {
+ case HashPartition(hashAttributeNames) =>
+ HashBasedShufflePartitioning(
+ dataTransferBatchSize,
+ fromWorkerIds.flatMap(from =>
+ toWorkerIds.map(to => ChannelIdentity(from, to, isControl = false))
+ ),
+ hashAttributeNames
+ )
+
+ case RangePartition(rangeAttributeNames, rangeMin, rangeMax) =>
+ RangeBasedShufflePartitioning(
+ dataTransferBatchSize,
+ fromWorkerIds.flatMap(fromId =>
+ toWorkerIds.map(toId => ChannelIdentity(fromId, toId, isControl = false))
+ ),
+ rangeAttributeNames,
+ rangeMin,
+ rangeMax
+ )
+
+ case SinglePartition() =>
+ assert(toWorkerIds.size == 1)
+ OneToOnePartitioning(
+ dataTransferBatchSize,
+ fromWorkerIds.map(fromWorkerId =>
+ ChannelIdentity(fromWorkerId, toWorkerIds.head, isControl = false)
+ )
+ )
+
+ case OneToOnePartition() =>
+ OneToOnePartitioning(
+ dataTransferBatchSize,
+ fromWorkerIds.zip(toWorkerIds).map {
+ case (fromWorkerId, toWorkerId) =>
+ ChannelIdentity(fromWorkerId, toWorkerId, isControl = false)
+ }
+ )
+
+ case BroadcastPartition() =>
+ BroadcastPartitioning(
+ dataTransferBatchSize,
+ fromWorkerIds.zip(toWorkerIds).map {
+ case (fromWorkerId, toWorkerId) =>
+ ChannelIdentity(fromWorkerId, toWorkerId, isControl = false)
+ }
+ )
+
+ case UnknownPartition() =>
+ RoundRobinPartitioning(
+ dataTransferBatchSize,
+ fromWorkerIds.flatMap(from =>
+ toWorkerIds.map(to => ChannelIdentity(from, to, isControl = false))
+ )
+ )
+
+ case _ =>
+ throw new UnsupportedOperationException()
+
+ }
+ }
+}
+
+case class LinkConfig(
+ channelConfigs: List[ChannelConfig],
+ partitioning: Partitioning
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/OperatorConfig.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/OperatorConfig.scala
new file mode 100644
index 00000000000..883a09e4bf0
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/OperatorConfig.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling.config
+
+case object OperatorConfig {
+ def empty: OperatorConfig = {
+ OperatorConfig(workerConfigs = List())
+ }
+}
+
+case class OperatorConfig(
+ workerConfigs: List[WorkerConfig]
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/PortConfig.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/PortConfig.scala
new file mode 100644
index 00000000000..b4a1e058b44
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/PortConfig.scala
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling.config
+
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.Partitioning
+
+import java.net.URI
+
+/**
+ * Super-type for any per-port config.
+ * After resource allocation, only OutputPortConfig or InputPortConfig remain.
+ */
+sealed trait PortConfig {
+ def storageURIs: List[URI]
+}
+
+/** An output port requires exactly one materialization URI. */
+final case class OutputPortConfig(storageURI: URI) extends PortConfig {
+ override val storageURIs: List[URI] = List(storageURI)
+}
+
+/**
+ * This class is needed as we fill the ResouceConfig of a region in two passes (before and after the ResouceAllocator
+ * is invoked.) In the first pass, the ScheduleGenerator builds a schedule and assigns materialization URIs to
+ * input/output ports. The URI allocation happens in this pass as the ScheduleGenerator can assign URIs as it creates
+ * each region, utilizing its global information about materializations across regions. After a Region DAG is
+ * finalized by the ScheduleGenerator, the ScheduleGenerator invokes ResouceAllocator, which allocates workers. As
+ * Partitioning can only be created after worker allocation, IntermediateInputPortConfig serves as the intermediate result
+ * before the ResourceAllocator is invoked. After ResourceAllocator finishes allocating resources, it will be
+ * upgraded to an InputPortConfig.
+ */
+final case class IntermediateInputPortConfig(storageURIs: List[URI]) extends PortConfig
+
+/**
+ * Final form after ResourceAllocator is invoked by the ScheduleGenerator.
+ * Each URI is associated with its Partitioning.
+ */
+final case class InputPortConfig(
+ storagePairs: List[(URI, Partitioning)]
+) extends PortConfig {
+ override val storageURIs: List[URI] = storagePairs.map(_._1)
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/ResourceConfig.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/ResourceConfig.scala
new file mode 100644
index 00000000000..b89ff9a5a91
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/ResourceConfig.scala
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling.config
+
+import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
+import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PhysicalLink}
+
+case class ResourceConfig(
+ operatorConfigs: Map[PhysicalOpIdentity, OperatorConfig] = Map.empty,
+ linkConfigs: Map[PhysicalLink, LinkConfig] = Map.empty,
+ portConfigs: Map[GlobalPortIdentity, PortConfig] = Map.empty
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/WorkerConfig.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/WorkerConfig.scala
new file mode 100644
index 00000000000..60a55151850
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/WorkerConfig.scala
@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling.config
+
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflow.PhysicalOp
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+case object WorkerConfig {
+ def generateWorkerConfigs(physicalOp: PhysicalOp): List[WorkerConfig] = {
+ val workerCount = if (physicalOp.parallelizable) {
+ physicalOp.suggestedWorkerNum match {
+ // Keep suggested number of workers
+ case Some(num) => num
+ // If no suggested number, use default value
+ case None => ApplicationConfig.numWorkerPerOperatorByDefault
+ }
+ } else {
+ // Non parallelizable operator has only 1 worker
+ 1
+ }
+
+ (0 until workerCount).toList.map(idx =>
+ WorkerConfig(
+ VirtualIdentityUtils.createWorkerIdentity(physicalOp.workflowId, physicalOp.id, idx)
+ )
+ )
+ }
+}
+
+case class WorkerConfig(
+ workerId: ActorVirtualIdentity
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/resourcePolicies/ExecutionClusterInfo.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/resourcePolicies/ExecutionClusterInfo.scala
new file mode 100644
index 00000000000..c99121b735e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/resourcePolicies/ExecutionClusterInfo.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling.resourcePolicies
+
+class ExecutionClusterInfo() {}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/resourcePolicies/ResourceAllocator.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/resourcePolicies/ResourceAllocator.scala
new file mode 100644
index 00000000000..2953ece0b34
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/resourcePolicies/ResourceAllocator.scala
@@ -0,0 +1,208 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling.resourcePolicies
+
+import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
+import org.apache.texera.amber.core.workflow._
+import org.apache.texera.amber.engine.architecture.scheduling.Region
+import org.apache.texera.amber.engine.architecture.scheduling.config.ChannelConfig.generateChannelConfigs
+import org.apache.texera.amber.engine.architecture.scheduling.config.LinkConfig.toPartitioning
+import org.apache.texera.amber.engine.architecture.scheduling.config.WorkerConfig.generateWorkerConfigs
+import org.apache.texera.amber.engine.architecture.scheduling.config._
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.Partitioning
+import org.apache.texera.amber.util.VirtualIdentityUtils.getFromActorIdForInputPortStorage
+
+import java.net.URI
+import scala.collection.mutable
+
+trait ResourceAllocator {
+ def allocate(region: Region): (ResourceConfig, Double)
+}
+
+class DefaultResourceAllocator(
+ physicalPlan: PhysicalPlan,
+ executionClusterInfo: ExecutionClusterInfo,
+ workflowSettings: WorkflowSettings
+) extends ResourceAllocator {
+
+ // a map of a physical link to the partition info of the upstream/downstream of this link
+ private val linkPartitionInfos = new mutable.HashMap[PhysicalLink, PartitionInfo]()
+
+ private val operatorConfigs = new mutable.HashMap[PhysicalOpIdentity, OperatorConfig]()
+ private val linkConfigs = new mutable.HashMap[PhysicalLink, LinkConfig]()
+
+ /**
+ * Allocates resources for a given region and its operators.
+ *
+ * This method calculates and assigns worker configurations for each operator
+ * in the region. For the operators that are parallelizable, it respects the
+ * suggested worker number if provided. Otherwise, it falls back to a default
+ * value. Non-parallelizable operators are assigned a single worker.
+ *
+ * @param region The region for which to allocate resources.
+ * @return A tuple containing:
+ * 1) A resource configuration.
+ * 2) An estimated cost of the workflow with the resource configuration,
+ * represented as a Double value (currently set to 0, but will be
+ * updated in the future).
+ */
+ def allocate(
+ region: Region
+ ): (ResourceConfig, Double) = {
+
+ val opToOperatorConfigMapping = region.getOperators
+ .map(physicalOp => physicalOp.id -> OperatorConfig(generateWorkerConfigs(physicalOp)))
+ .toMap
+
+ operatorConfigs ++= opToOperatorConfigMapping
+
+ propagatePartitionRequirement(region)
+
+ val linkToLinkConfigMapping = region.getLinks.map { physicalLink =>
+ physicalLink -> LinkConfig(
+ generateChannelConfigs(
+ operatorConfigs(physicalLink.fromOpId).workerConfigs.map(_.workerId),
+ operatorConfigs(physicalLink.toOpId).workerConfigs.map(_.workerId),
+ toPortId = physicalLink.toPortId,
+ linkPartitionInfos(physicalLink)
+ ),
+ toPartitioning(
+ operatorConfigs(physicalLink.fromOpId).workerConfigs.map(_.workerId),
+ operatorConfigs(physicalLink.toOpId).workerConfigs.map(_.workerId),
+ linkPartitionInfos(physicalLink),
+ workflowSettings.dataTransferBatchSize
+ )
+ )
+ }.toMap
+
+ linkConfigs ++= linkToLinkConfigMapping
+
+ val portConfigs: Map[GlobalPortIdentity, PortConfig] = region.resourceConfig match {
+ case Some(existing) =>
+ val upgradedInputPortConfigs: Map[GlobalPortIdentity, InputPortConfig] =
+ existing.portConfigs.collect {
+ case (globalPortId, rawInConfig: IntermediateInputPortConfig) if globalPortId.input =>
+ val uris: List[URI] = rawInConfig.storageURIs
+ // derive partitionings for each upstream materialization
+ val portPartitionings: List[Partitioning] = uris.map { inputMatUri =>
+ val toWorkerActorIds =
+ operatorConfigs(globalPortId.opId).workerConfigs.map(_.workerId)
+ val fromVirtualThreadActorIds = toWorkerActorIds.map(toWorkerActorId =>
+ getFromActorIdForInputPortStorage(inputMatUri.toString, toWorkerActorId)
+ )
+ // Extract the input port partitionInfo defined in the physicalOp, defaulting to UnknownPartition.
+ val inputPortPartitionInfo = region
+ .getOperator(globalPortId.opId)
+ .partitionRequirement
+ .applyOrElse(globalPortId.portId.id, (_: Int) => None)
+ .getOrElse(UnknownPartition())
+
+ toPartitioning(
+ fromVirtualThreadActorIds,
+ toWorkerActorIds,
+ inputPortPartitionInfo,
+ workflowSettings.dataTransferBatchSize
+ )
+ }
+ // new InputPortConfig that carries both URIs and per-URI partitionings
+ globalPortId -> InputPortConfig(uris.zip(portPartitionings))
+ }
+
+ existing.portConfigs ++ upgradedInputPortConfigs
+
+ case None =>
+ Map.empty[GlobalPortIdentity, PortConfig]
+ }
+
+ val resourceConfig = ResourceConfig(
+ opToOperatorConfigMapping,
+ linkToLinkConfigMapping,
+ portConfigs
+ )
+
+ (resourceConfig, 0)
+ }
+
+ /**
+ * This method propagates partitioning requirements in the PhysicalPlan DAG.
+ *
+ * This method is invoked once for each region, and only propagate partitioning requirements within
+ * the region. For example, suppose we have the following physical Plan:
+ *
+ * A ->
+ * HJ
+ * B ->
+ * The link A->HJ will be propagated in the first region. The link B->HJ will be propagated in the second region.
+ * The output partition info of HJ will be derived after both links are propagated, which is in the second region.
+ */
+ private def propagatePartitionRequirement(region: Region): Unit = {
+ region
+ .topologicalIterator()
+ .foreach(physicalOpId => {
+ val physicalOp = region.getOperator(physicalOpId)
+ val outputPartitionInfo = if (physicalPlan.getSourceOperatorIds.contains(physicalOpId)) {
+ Some(physicalOp.partitionRequirement.headOption.flatten.getOrElse(UnknownPartition()))
+ } else {
+ val inputPartitionInfos = physicalOp.inputPorts.keys
+ .flatMap((portId: PortIdentity) => {
+ physicalOp
+ .getInputLinks(Some(portId))
+ .filter(link => region.getLinks.contains(link))
+ .map(link => {
+ val previousLinkPartitionInfo =
+ linkPartitionInfos.getOrElse(link, UnknownPartition())
+ val updatedLinkPartitionInfo = physicalPlan.getOutputPartitionInfo(
+ link,
+ previousLinkPartitionInfo,
+ operatorConfigs.map {
+ case (opId, operatorConfig) => opId -> operatorConfig.workerConfigs.length
+ }.toMap
+ )
+ linkPartitionInfos.put(link, updatedLinkPartitionInfo)
+ (link.toPortId, updatedLinkPartitionInfo)
+ })
+ })
+ // group upstream partition infos by input port of this physicalOp
+ .groupBy(_._1)
+ .values
+ .toList
+ // if there are multiple partition infos on an input port, reduce them to once
+ .map(_.map(_._2).reduce((p1, p2) => p1.merge(p2)))
+
+ if (inputPartitionInfos.length == physicalOp.inputPorts.size) {
+ // derive the output partition info with all the input partition infos
+ Some(physicalOp.derivePartition(inputPartitionInfos))
+ } else {
+ None
+ }
+
+ }
+
+ if (outputPartitionInfo.isDefined) {
+ physicalOp.outputPorts.keys
+ .flatMap(physicalOp.getOutputLinks)
+ .foreach(link =>
+ // by default, a link's partition info comes from its input, unless updated to match its output.
+ linkPartitionInfos.put(link, outputPartitionInfo.get)
+ )
+ }
+ })
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/BroadcastPartitioner.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/BroadcastPartitioner.scala
new file mode 100644
index 00000000000..71de92ced60
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/BroadcastPartitioner.scala
@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.sendsemantics.partitioners
+
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.BroadcastPartitioning
+
+case class BroadcastPartitioner(partitioning: BroadcastPartitioning) extends Partitioner {
+
+ private val receivers = partitioning.channels.map(_.toWorkerId).distinct
+
+ override def getBucketIndex(tuple: Tuple): Iterator[Int] = {
+ receivers.indices.iterator
+ }
+
+ override def allReceivers: Seq[ActorVirtualIdentity] = receivers
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/HashBasedShufflePartitioner.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/HashBasedShufflePartitioner.scala
new file mode 100644
index 00000000000..7025a4bd2cf
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/HashBasedShufflePartitioner.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.sendsemantics.partitioners
+
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.HashBasedShufflePartitioning
+
+case class HashBasedShufflePartitioner(partitioning: HashBasedShufflePartitioning)
+ extends Partitioner {
+
+ private val receivers = partitioning.channels.map(_.toWorkerId).distinct
+
+ override def getBucketIndex(tuple: Tuple): Iterator[Int] = {
+ val numBuckets = receivers.length
+ val partialTuple =
+ if (partitioning.hashAttributeNames.isEmpty) tuple
+ else tuple.getPartialTuple(partitioning.hashAttributeNames.toList)
+ val index = Math.floorMod(partialTuple.hashCode(), numBuckets)
+ Iterator(index)
+ }
+
+ override def allReceivers: Seq[ActorVirtualIdentity] = receivers
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/OneToOnePartitioner.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/OneToOnePartitioner.scala
new file mode 100644
index 00000000000..f1360f8be4a
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/OneToOnePartitioner.scala
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.sendsemantics.partitioners
+
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.OneToOnePartitioning
+
+case class OneToOnePartitioner(partitioning: OneToOnePartitioning, actorId: ActorVirtualIdentity)
+ extends Partitioner {
+
+ override def getBucketIndex(tuple: Tuple): Iterator[Int] = Iterator(0)
+
+ override def allReceivers: Seq[ActorVirtualIdentity] =
+ Seq(partitioning.channels.filter(_.fromWorkerId == actorId).head.toWorkerId)
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/Partitioner.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/Partitioner.scala
new file mode 100644
index 00000000000..eac77bce365
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/Partitioner.scala
@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.sendsemantics.partitioners
+
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.state.State
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.messaginglayer.NetworkOutputGateway
+import org.apache.texera.amber.engine.common.ambermessage.{DataFrame, StateFrame}
+
+import scala.collection.mutable.ArrayBuffer
+
+trait Partitioner extends Serializable {
+ def getBucketIndex(tuple: Tuple): Iterator[Int]
+
+ def allReceivers: Seq[ActorVirtualIdentity]
+}
+
+class NetworkOutputBuffer(
+ val to: ActorVirtualIdentity,
+ val dataOutputPort: NetworkOutputGateway,
+ val batchSize: Int = ApplicationConfig.defaultDataTransferBatchSize
+) {
+
+ var buffer = new ArrayBuffer[Tuple]()
+
+ def addTuple(tuple: Tuple): Unit = {
+ buffer.append(tuple)
+ if (buffer.size >= batchSize) {
+ flush()
+ }
+ }
+
+ def sendState(state: State): Unit = {
+ flush()
+ dataOutputPort.sendTo(to, StateFrame(state))
+ flush()
+ }
+
+ def flush(): Unit = {
+ if (buffer.nonEmpty) {
+ dataOutputPort.sendTo(to, DataFrame(buffer.toArray))
+ buffer = new ArrayBuffer[Tuple]()
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/RangeBasedShufflePartitioner.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/RangeBasedShufflePartitioner.scala
new file mode 100644
index 00000000000..9b8aae263f1
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/RangeBasedShufflePartitioner.scala
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.sendsemantics.partitioners
+
+import org.apache.texera.amber.core.tuple.{AttributeType, Tuple}
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.RangeBasedShufflePartitioning
+
+case class RangeBasedShufflePartitioner(partitioning: RangeBasedShufflePartitioning)
+ extends Partitioner {
+
+ private val receivers = partitioning.channels.map(_.toWorkerId).distinct
+ private val keysPerReceiver =
+ ((partitioning.rangeMax - partitioning.rangeMin) / receivers.length) + 1
+
+ override def getBucketIndex(tuple: Tuple): Iterator[Int] = {
+ // Do range partitioning only on the first attribute in `rangeAttributeNames`.
+ val attribute = tuple.getSchema.getAttribute(partitioning.rangeAttributeNames.head)
+ var fieldVal: Long = -1
+ attribute.getType match {
+ case AttributeType.LONG =>
+ fieldVal = tuple.getField[Long](attribute)
+ case AttributeType.INTEGER =>
+ fieldVal = tuple.getField[Int](attribute)
+ case AttributeType.DOUBLE =>
+ fieldVal = tuple.getField[Double](attribute).toLong
+ case _ =>
+ throw new RuntimeException(s"unsupported attribute type: ${attribute.getType}")
+ }
+
+ if (fieldVal < partitioning.rangeMin) {
+ return Iterator(0)
+ }
+ if (fieldVal > partitioning.rangeMax) {
+ return Iterator(receivers.length - 1)
+ }
+ Iterator(((fieldVal - partitioning.rangeMin) / keysPerReceiver).toInt)
+ }
+
+ override def allReceivers: Seq[ActorVirtualIdentity] = receivers
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/RoundRobinPartitioner.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/RoundRobinPartitioner.scala
new file mode 100644
index 00000000000..185f2e5e3ac
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/RoundRobinPartitioner.scala
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.sendsemantics.partitioners
+
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.RoundRobinPartitioning
+
+case class RoundRobinPartitioner(partitioning: RoundRobinPartitioning) extends Partitioner {
+ private var roundRobinIndex = 0
+ private val receivers = partitioning.channels.map(_.toWorkerId).distinct
+
+ override def getBucketIndex(tuple: Tuple): Iterator[Int] = {
+ roundRobinIndex = (roundRobinIndex + 1) % receivers.length
+ Iterator(roundRobinIndex)
+ }
+
+ override def allReceivers: Seq[ActorVirtualIdentity] = receivers
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DPThread.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DPThread.scala
new file mode 100644
index 00000000000..0c136d613a5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DPThread.scala
@@ -0,0 +1,222 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.logreplay.ReplayLogManager
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessage
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ DPInputQueueElement,
+ MainThreadDelegateMessage
+}
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.{
+ READY,
+ UNINITIALIZED
+}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.actormessage.{ActorCommand, Backpressure}
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataPayload,
+ DirectControlMessagePayload,
+ WorkflowFIFOMessage
+}
+import org.apache.texera.amber.engine.common.virtualidentity.util.SELF
+import org.apache.texera.amber.error.ErrorUtils.safely
+
+import java.util.concurrent._
+
+class DPThread(
+ val actorId: ActorVirtualIdentity,
+ dp: DataProcessor,
+ logManager: ReplayLogManager,
+ internalQueue: LinkedBlockingQueue[DPInputQueueElement]
+) extends AmberLogging {
+
+ // initialize dp thread upon construction
+ @transient
+ var dpThreadExecutor: ExecutorService = _
+ @transient
+ var dpThread: Future[_] = _
+
+ var backpressureStatus = false
+
+ def getThreadName: String = "DP-thread"
+
+ private val endFuture = new CompletableFuture[Unit]()
+
+ def stop(): Unit = {
+ if (dpThread != null) {
+ dpThread.cancel(true) // interrupt
+ stopped = true
+ endFuture.get()
+ }
+ if (dpThreadExecutor != null) {
+ dpThreadExecutor.shutdownNow() // destroy thread
+ }
+ }
+
+ @volatile
+ private var stopped = false
+
+ def start(): Unit = {
+ if (dpThreadExecutor != null) {
+ logger.info("DP Thread is already running")
+ return
+ }
+ dpThreadExecutor = Executors.newSingleThreadExecutor
+ if (dp.stateManager.getCurrentState == UNINITIALIZED) {
+ dp.stateManager.transitTo(READY)
+ }
+ if (dpThread == null) {
+ // TODO: setup context
+ // operator.context = new OperatorContext(new TimeService(logManager))
+ val startFuture = new CompletableFuture[Unit]()
+ dpThread = dpThreadExecutor.submit(new Runnable() {
+ def run(): Unit = {
+ Thread.currentThread().setName(getThreadName)
+ logger.info("DP thread started")
+ startFuture.complete(())
+ dp.statisticsManager.initializeWorkerStartTime(System.nanoTime())
+ try {
+ runDPThreadMainLogic()
+ } catch safely {
+ case _: InterruptedException =>
+ // dp thread will stop here
+ logger.info("DP Thread exits")
+ case err: Throwable =>
+ logger.error("DP Thread exists unexpectedly", err)
+ dp.outputHandler(Left(MainThreadDelegateMessage((worker) => {
+ // notify main thread
+ throw err
+ })))
+ }
+ dp.statisticsManager.updateTotalExecutionTime(System.nanoTime())
+ endFuture.complete(())
+ }
+ })
+ startFuture.get()
+ }
+ }
+
+ def handleActorCommand(cmd: ActorCommand): Unit = {
+ cmd match {
+ case Backpressure(enabled) =>
+ backpressureStatus = enabled
+ case _ => // no op
+ }
+ }
+
+ @throws[Exception]
+ private[this] def runDPThreadMainLogic(): Unit = {
+ //
+ // Main loop step 1: receive messages from actor and apply FIFO
+ //
+ var waitingForInput = false
+ while (!stopped) {
+ while (internalQueue.size > 0 || waitingForInput) {
+ val elem = internalQueue.take
+ waitingForInput = false
+ elem match {
+ case WorkflowWorker.FIFOMessageElement(msg) =>
+ val channel = dp.inputGateway.getChannel(msg.channelId)
+ channel.acceptMessage(msg)
+ case WorkflowWorker.TimerBasedControlElement(control) =>
+ // establish order according to receiving order.
+ // Note: this will not guarantee fifo & exactly-once
+ // Please make sure the control here is IDEMPOTENT and ORDER-INDEPENDENT.
+ val controlChannelId = ChannelIdentity(SELF, SELF, isControl = true)
+ val channel = dp.inputGateway.getChannel(controlChannelId)
+ channel.acceptMessage(
+ WorkflowFIFOMessage(controlChannelId, channel.getCurrentSeq, control)
+ )
+ case WorkflowWorker.ActorCommandElement(msg) =>
+ handleActorCommand(msg)
+ }
+ }
+
+ //
+ // Main loop step 2: do input selection
+ //
+ var channelId: ChannelIdentity = null
+ var msgOpt: Option[WorkflowFIFOMessage] = None
+ if (
+ dp.inputManager.hasUnfinishedInput || dp.outputManager.hasUnfinishedOutput || dp.pauseManager.isPaused
+ ) {
+ dp.inputGateway.tryPickControlChannel match {
+ case Some(channel) =>
+ channelId = channel.channelId
+ msgOpt = Some(channel.take)
+ case None =>
+ // continue processing
+ if (!dp.pauseManager.isPaused && !backpressureStatus) {
+ channelId = dp.inputManager.currentChannelId
+ } else {
+ waitingForInput = true
+ }
+ }
+ } else {
+ // take from input port
+ if (backpressureStatus) {
+ dp.inputGateway.tryPickControlChannel
+ } else {
+ dp.inputGateway.tryPickChannel
+ } match {
+ case Some(channel) =>
+ channelId = channel.channelId
+ msgOpt = Some(channel.take)
+ case None => waitingForInput = true
+ }
+ }
+
+ //
+ // Main loop step 3: process selected message payload
+ //
+ if (channelId != null) {
+ // for logging, skip large data frames.
+ val msgToLog = msgOpt.filter(_.payload.isInstanceOf[DirectControlMessagePayload])
+ logManager.withFaultTolerant(channelId, msgToLog) {
+ msgOpt match {
+ case None =>
+ dp.continueDataProcessing()
+ case Some(msg) =>
+ msg.payload match {
+ case payload: DirectControlMessagePayload =>
+ dp.processDCM(msg.channelId, payload)
+ case payload: DataPayload =>
+ dp.processDataPayload(msg.channelId, payload)
+ case ecm: EmbeddedControlMessage =>
+ dp.processECM(msg.channelId, ecm, logManager)
+ }
+ }
+ }
+ }
+ // As the computation is chopped into steps, the checkpoint
+ // serialization must happen after/before a step. Otherwise
+ // DP state will be restored in the middle of a step, which
+ // is often not what we want. Thus, we have this one-time
+ // additional serializationCall assigned inside the checkpoint
+ // handler.
+ dp.serializationManager.applySerialization()
+
+ dp.statisticsManager.updateTotalExecutionTime(System.nanoTime())
+ // End of Main loop
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessor.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessor.scala
new file mode 100644
index 00000000000..84f1e8ec659
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessor.scala
@@ -0,0 +1,311 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import com.softwaremill.macwire.wire
+import io.grpc.MethodDescriptor
+import org.apache.texera.amber.core.executor.OperatorExecutor
+import org.apache.texera.amber.core.state.State
+import org.apache.texera.amber.core.tuple._
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.common.AmberProcessor
+import org.apache.texera.amber.engine.architecture.logreplay.ReplayLogManager
+import org.apache.texera.amber.engine.architecture.messaginglayer.{
+ InputManager,
+ OutputManager,
+ WorkerTimerService
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.{
+ NO_ALIGNMENT,
+ PORT_ALIGNMENT
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_END_CHANNEL
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ DPInputQueueElement,
+ MainThreadDelegateMessage
+}
+import org.apache.texera.amber.engine.architecture.worker.managers.SerializationManager
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.{
+ COMPLETED,
+ READY,
+ RUNNING
+}
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerStatistics
+import org.apache.texera.amber.engine.common.ambermessage._
+import org.apache.texera.amber.engine.common.statetransition.WorkerStateManager
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.error.ErrorUtils.{mkConsoleMessage, safely}
+
+import java.util.concurrent.LinkedBlockingQueue
+
+class DataProcessor(
+ actorId: ActorVirtualIdentity,
+ outputHandler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit,
+ inputMessageQueue: LinkedBlockingQueue[DPInputQueueElement]
+) extends AmberProcessor(actorId, outputHandler)
+ with Serializable {
+
+ @transient var executor: OperatorExecutor = _
+
+ def initTimerService(adaptiveBatchingMonitor: WorkerTimerService): Unit = {
+ this.adaptiveBatchingMonitor = adaptiveBatchingMonitor
+ }
+
+ @transient var adaptiveBatchingMonitor: WorkerTimerService = _
+
+ // inner dependencies
+ private val initializer = new DataProcessorRPCHandlerInitializer(this)
+ val pauseManager: PauseManager = wire[PauseManager]
+ val stateManager: WorkerStateManager = new WorkerStateManager(actorId)
+ val inputManager: InputManager = new InputManager(actorId, inputMessageQueue)
+ val outputManager: OutputManager = new OutputManager(actorId, outputGateway)
+ val ecmManager: EmbeddedControlMessageManager =
+ new EmbeddedControlMessageManager(actorId, inputGateway, inputManager)
+ val serializationManager: SerializationManager = new SerializationManager(actorId)
+
+ def getQueuedCredit(channelId: ChannelIdentity): Long = {
+ inputGateway.getChannel(channelId).getQueuedCredit
+ }
+
+ /**
+ * provide API for actor to get stats of this operator
+ */
+ def collectStatistics(): WorkerStatistics =
+ statisticsManager.getStatistics(executor)
+
+ /**
+ * process currentInputTuple through executor logic.
+ * this function is only called by the DP thread.
+ */
+ private[this] def processInputTuple(tuple: Tuple): Unit = {
+ try {
+ val portIdentity: PortIdentity =
+ this.inputGateway.getChannel(inputManager.currentChannelId).getPortId
+ outputManager.outputIterator.setTupleOutput(
+ executor.processTupleMultiPort(
+ tuple,
+ portIdentity.id
+ )
+ )
+
+ statisticsManager.increaseInputStatistics(portIdentity, tuple.inMemSize)
+
+ } catch safely {
+ case e =>
+ // forward input tuple to the user and pause DP thread
+ handleExecutorException(e)
+ }
+ }
+
+ private[this] def processInputState(state: State, port: Int): Unit = {
+ try {
+ val outputState = executor.processState(state, port)
+ if (outputState.isDefined) {
+ outputManager.emitState(outputState.get)
+ }
+ } catch safely {
+ case e =>
+ handleExecutorException(e)
+ }
+ }
+
+ /** transfer one tuple from iterator to downstream.
+ * this function is only called by the DP thread
+ */
+ private[this] def outputOneTuple(): Unit = {
+ adaptiveBatchingMonitor.startAdaptiveBatching()
+ var out: (TupleLike, Option[PortIdentity]) = null
+ try {
+ out = outputManager.outputIterator.next()
+ } catch safely {
+ case e =>
+ // invalidate current output tuple
+ out = null
+ // also invalidate outputIterator
+ outputManager.outputIterator.setTupleOutput(Iterator.empty)
+ // forward input tuple to the user and pause DP thread
+ handleExecutorException(e)
+ }
+ if (out == null) return
+
+ val (outputTuple, outputPortOpt) = out
+
+ if (outputTuple == null) return
+ outputTuple match {
+ case FinalizeExecutor() =>
+ sendECMToDataChannels(METHOD_END_CHANNEL, PORT_ALIGNMENT)
+ // Send Completed signal to worker actor.
+ executor.close()
+ adaptiveBatchingMonitor.stopAdaptiveBatching()
+ stateManager.transitTo(COMPLETED)
+ logger.info(
+ s"$executor completed, # of input ports = ${inputManager.getAllPorts.size}, " +
+ s"input tuple count = ${statisticsManager.getInputTupleCount}, " +
+ s"output tuple count = ${statisticsManager.getOutputTupleCount}"
+ )
+ asyncRPCClient.controllerInterface.workerExecutionCompleted(
+ EmptyRequest(),
+ asyncRPCClient.mkContext(CONTROLLER)
+ )
+ case FinalizePort(portId, input) =>
+ if (!input) {
+ outputManager.closeOutputStorageWriterIfNeeded(portId)
+ }
+ asyncRPCClient.controllerInterface.portCompleted(
+ PortCompletedRequest(portId, input),
+ asyncRPCClient.mkContext(CONTROLLER)
+ )
+ case schemaEnforceable: SchemaEnforceable =>
+ val portIdentity = outputPortOpt.getOrElse(outputManager.getSingleOutputPortIdentity)
+ val tuple = schemaEnforceable.enforceSchema(outputManager.getPort(portIdentity).schema)
+ statisticsManager.increaseOutputStatistics(portIdentity, tuple.inMemSize)
+ outputManager.passTupleToDownstream(tuple, outputPortOpt)
+ outputManager.saveTupleToStorageIfNeeded(tuple, outputPortOpt)
+
+ case other => // skip for now
+ }
+ }
+
+ def continueDataProcessing(): Unit = {
+ val dataProcessingStartTime = System.nanoTime()
+ if (outputManager.hasUnfinishedOutput) {
+ outputOneTuple()
+ } else {
+ processInputTuple(inputManager.getNextTuple)
+ }
+ statisticsManager.increaseDataProcessingTime(System.nanoTime() - dataProcessingStartTime)
+ }
+
+ def processDataPayload(
+ channelId: ChannelIdentity,
+ dataPayload: DataPayload
+ ): Unit = {
+ val dataProcessingStartTime = System.nanoTime()
+ val portId = this.inputGateway.getChannel(channelId).getPortId
+ dataPayload match {
+ case DataFrame(tuples) =>
+ stateManager.conditionalTransitTo(
+ READY,
+ RUNNING,
+ () => {
+ asyncRPCClient.controllerInterface.workerStateUpdated(
+ WorkerStateUpdatedRequest(stateManager.getCurrentState),
+ asyncRPCClient.mkContext(CONTROLLER)
+ )
+ }
+ )
+ inputManager.initBatch(channelId, tuples)
+ processInputTuple(inputManager.getNextTuple)
+ case StateFrame(state) =>
+ processInputState(state, portId.id)
+ }
+ statisticsManager.increaseDataProcessingTime(System.nanoTime() - dataProcessingStartTime)
+ }
+
+ def processECM(
+ channelId: ChannelIdentity,
+ ecm: EmbeddedControlMessage,
+ logManager: ReplayLogManager
+ ): Unit = {
+ inputManager.currentChannelId = channelId
+ val command = ecm.commandMapping.get(actorId.name)
+ logger.info(s"receive ECM from $channelId, id = ${ecm.id}, cmd = $command")
+ if (ecm.ecmType != NO_ALIGNMENT) {
+ pauseManager.pauseInputChannel(ECMPause(ecm.id), List(channelId))
+ }
+ if (ecmManager.isECMAligned(channelId, ecm)) {
+ logManager.markAsReplayDestination(ecm.id)
+ // invoke the control command carried with the ECM
+ logger.info(s"process ECM from $channelId, id = ${ecm.id}, cmd = $command")
+ if (command.isDefined) {
+ // The reply must go back to the actor that originated the invocation
+ // (recorded in command.context.sender), not to channelId.fromWorkerId.
+ // For ECM-embedded commands those differ: channelId is the data
+ // channel between two workers, while the originator is typically the
+ // controller. Fall back to the channel sender when the context is
+ // unset (e.g. unit-test inputs).
+ val ctx = command.get.context
+ val replyTo =
+ if (ctx.sender.name.nonEmpty) ctx.sender else channelId.fromWorkerId
+ asyncRPCServer.receive(command.get, replyTo)
+ }
+ // if this worker is not the final destination of the ECM, pass it downstream
+ val downstreamChannelsInScope = ecm.scope.filter(_.fromWorkerId == actorId).toSet
+ if (downstreamChannelsInScope.nonEmpty) {
+ outputManager.flush(Some(downstreamChannelsInScope))
+ outputGateway.getActiveChannels.foreach { activeChannelId =>
+ if (downstreamChannelsInScope.contains(activeChannelId)) {
+ logger.info(
+ s"send ECM to $activeChannelId, id = ${ecm.id}, cmd = $command"
+ )
+ outputGateway.sendTo(activeChannelId, ecm)
+ }
+ }
+ }
+ // unblock input channels
+ if (ecm.ecmType != NO_ALIGNMENT) {
+ pauseManager.resume(ECMPause(ecm.id))
+ }
+ }
+ }
+
+ def sendECMToDataChannels(
+ method: MethodDescriptor[EmptyRequest, EmptyReturn],
+ alignment: EmbeddedControlMessageType
+ ): Unit = {
+ outputManager.flush()
+ outputGateway.getActiveChannels
+ .filter(!_.isControl)
+ .foreach { activeChannelId =>
+ asyncRPCClient.sendECMToChannel(
+ EmbeddedControlMessageIdentity(method.getBareMethodName),
+ alignment,
+ Set(),
+ Map(
+ activeChannelId.toWorkerId.name ->
+ ControlInvocation(
+ method.getBareMethodName,
+ EmptyRequest(),
+ AsyncRPCContext(ActorVirtualIdentity(""), ActorVirtualIdentity("")),
+ -1
+ )
+ ),
+ activeChannelId
+ )
+ }
+ }
+
+ def handleExecutorException(e: Throwable): Unit = {
+ asyncRPCClient.controllerInterface.consoleMessageTriggered(
+ ConsoleMessageTriggeredRequest(mkConsoleMessage(actorId, e)),
+ asyncRPCClient.mkContext(CONTROLLER)
+ )
+ logger.warn(e.getLocalizedMessage + "\n" + e.getStackTrace.mkString("\n"))
+ // invoke a pause in-place
+ pauseManager.pause(OperatorLogicPause)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessorRPCHandlerInitializer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessorRPCHandlerInitializer.scala
new file mode 100644
index 00000000000..6b0c62ac3f2
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessorRPCHandlerInitializer.scala
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.executor.{
+ ExecFactory,
+ OpExecInitInfo,
+ OpExecSource,
+ OpExecWithClassName,
+ OpExecWithCode
+}
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ DebugCommandRequest,
+ EmptyRequest,
+ EvaluatePythonExpressionRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{EmptyReturn, EvaluatedValue}
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceFs2Grpc
+import org.apache.texera.amber.engine.architecture.worker.promisehandlers._
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCHandlerInitializer
+import org.apache.texera.amber.operator.source.cache.CacheSourceOpExec
+
+import java.net.URI
+
+class DataProcessorRPCHandlerInitializer(val dp: DataProcessor)
+ extends AsyncRPCHandlerInitializer(dp.asyncRPCClient, dp.asyncRPCServer)
+ with WorkerServiceFs2Grpc[Future, AsyncRPCContext]
+ with AmberLogging
+ with InitializeExecutorHandler
+ with OpenExecutorHandler
+ with PauseHandler
+ with AddPartitioningHandler
+ with QueryStatisticsHandler
+ with ResumeHandler
+ with StartHandler
+ with EndHandler
+ with StartChannelHandler
+ with EndChannelHandler
+ with AssignPortHandler
+ with AddInputChannelHandler
+ with FlushNetworkBufferHandler
+ with RetrieveStateHandler
+ with PrepareCheckpointHandler
+ with FinalizeCheckpointHandler
+ with UpdateExecutorHandler {
+ val actorId: ActorVirtualIdentity = dp.actorId
+
+ var cachedTotalWorkerCount = 0
+
+ override def debugCommand(
+ request: DebugCommandRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = ???
+
+ override def evaluatePythonExpression(
+ request: EvaluatePythonExpressionRequest,
+ ctx: AsyncRPCContext
+ ): Future[EvaluatedValue] = ???
+
+ override def retryCurrentTuple(request: EmptyRequest, ctx: AsyncRPCContext): Future[EmptyReturn] =
+ ???
+
+ override def noOperation(request: EmptyRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = ???
+
+ def setupExecutor(execInitInfo: OpExecInitInfo, workerIdx: Int, workerCount: Int): Unit = {
+ dp.executor = execInitInfo match {
+ case OpExecWithClassName(className, descString) =>
+ ExecFactory.newExecFromJavaClassName(className, descString, workerIdx, workerCount)
+ case OpExecWithCode(code, _) =>
+ ExecFactory.newExecFromJavaCode(code)
+ case OpExecSource(storageUri, _) =>
+ new CacheSourceOpExec(URI.create(storageUri))
+ case OpExecInitInfo.Empty =>
+ throw new IllegalArgumentException("Empty executor initialization info")
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/EmbeddedControlMessageManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/EmbeddedControlMessageManager.scala
new file mode 100644
index 00000000000..2837c11d80d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/EmbeddedControlMessageManager.scala
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.architecture.messaginglayer.{InputGateway, InputManager}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessage
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.{
+ ALL_ALIGNMENT,
+ NO_ALIGNMENT,
+ PORT_ALIGNMENT
+}
+import org.apache.texera.amber.engine.common.{AmberLogging, CheckpointState}
+
+import scala.collection.mutable
+
+class EmbeddedControlMessageManager(
+ val actorId: ActorVirtualIdentity,
+ inputGateway: InputGateway,
+ inputManager: InputManager
+) extends AmberLogging {
+
+ private val ecmReceived =
+ new mutable.HashMap[EmbeddedControlMessageIdentity, Set[ChannelIdentity]]()
+
+ val checkpoints = new mutable.HashMap[EmbeddedControlMessageIdentity, CheckpointState]()
+
+ /**
+ * Determines if an ECM is fully received from all relevant senders within its scope.
+ * This method checks if the ECM, based on its type, has been received from all necessary channels.
+ * For ECMs requiring alignment, it verifies receipt from all senders in the scope. For non-aligned ECMs,
+ * it checks if it's the first received ECM. Post verification, it cleans up the ECMs.
+ *
+ * @return Boolean indicating if the ECM is completely received from all senders
+ * within the scope. Returns true if the ECM is aligned, otherwise false.
+ */
+ def isECMAligned(
+ from: ChannelIdentity,
+ ecm: EmbeddedControlMessage
+ ): Boolean = {
+ val portId = inputGateway.getChannel(from).getPortId
+ if (!ecmReceived.contains(ecm.id)) {
+ ecmReceived(ecm.id) = Set()
+ }
+ ecmReceived.update(ecm.id, ecmReceived(ecm.id) + from)
+ val ecmReceivedFromAllChannels =
+ getChannelsWithinScope(ecm).subsetOf(ecmReceived(ecm.id))
+ // check if the ECM is completed
+ val ecmCompleted = ecm.ecmType match {
+ case ALL_ALIGNMENT =>
+ ecmReceivedFromAllChannels
+ case PORT_ALIGNMENT =>
+ inputManager.getPort(portId).channels.subsetOf(ecmReceived(ecm.id))
+ case NO_ALIGNMENT =>
+ ecmReceived(ecm.id).size == 1 // only the first ECM triggers
+ case _ =>
+ throw new IllegalArgumentException(
+ s"Unsupported ECM type: ${ecm.ecmType}"
+ )
+ }
+ if (ecmReceivedFromAllChannels) {
+ ecmReceived.remove(ecm.id) // clean up if all ECMs are received
+ }
+ ecmCompleted
+ }
+
+ private def getChannelsWithinScope(ecm: EmbeddedControlMessage): Set[ChannelIdentity] = {
+ if (ecm.scope.isEmpty) inputGateway.getAllDataChannels.map(_.channelId)
+ else {
+ val upstreams = ecm.scope.filter(_.toWorkerId == actorId)
+ inputGateway.getAllChannels
+ .map(_.channelId)
+ .filter { id =>
+ upstreams.contains(id)
+ }
+ }
+ }.toSet
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/PauseManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/PauseManager.scala
new file mode 100644
index 00000000000..80cf9c1415d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/PauseManager.scala
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.messaginglayer.InputGateway
+import org.apache.texera.amber.engine.common.AmberLogging
+
+import scala.collection.mutable
+
+class PauseManager(val actorId: ActorVirtualIdentity, inputGateway: InputGateway)
+ extends AmberLogging {
+
+ private val globalPauses = new mutable.HashSet[PauseType]()
+ private val specificInputPauses = mutable.MultiDict[PauseType, ChannelIdentity]()
+
+ def pause(pauseType: PauseType): Unit = {
+ globalPauses.add(pauseType)
+ // disable all data queues
+ inputGateway.getAllDataChannels.foreach(_.enable(false))
+ }
+
+ def pauseInputChannel(pauseType: PauseType, inputs: List[ChannelIdentity]): Unit = {
+ inputs.foreach(input => {
+ specificInputPauses.addOne((pauseType, input))
+ // disable specified data queues
+ inputGateway.getChannel(input).enable(false)
+ })
+ }
+
+ def resume(pauseType: PauseType): Unit = {
+ globalPauses.remove(pauseType)
+ specificInputPauses.removeKey(pauseType)
+
+ // still globally paused no action, don't need to resume anything
+ if (globalPauses.nonEmpty) {
+ return
+ }
+ // global pause is empty, specific input pause is also empty, resume all
+ if (specificInputPauses.isEmpty) {
+ inputGateway.getAllDataChannels.foreach(_.enable(true))
+ return
+ }
+ // need to resume specific input channels
+ val pausedChannels = specificInputPauses.values.toSet
+ inputGateway.getAllChannels.foreach(_.enable(true))
+ pausedChannels.foreach { ChannelIdentity =>
+ inputGateway.getChannel(ChannelIdentity).enable(false)
+ }
+ }
+
+ def isPaused: Boolean = {
+ globalPauses.nonEmpty
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/PauseType.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/PauseType.scala
new file mode 100644
index 00000000000..57fd5cee3ae
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/PauseType.scala
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import org.apache.texera.amber.core.virtualidentity.EmbeddedControlMessageIdentity
+
+sealed trait PauseType
+
+object UserPause extends PauseType
+
+object BackpressurePause extends PauseType
+
+object OperatorLogicPause extends PauseType
+
+case class ECMPause(id: EmbeddedControlMessageIdentity) extends PauseType
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/WorkflowWorker.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/WorkflowWorker.scala
new file mode 100644
index 00000000000..d1a0a300d93
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/WorkflowWorker.scala
@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import org.apache.pekko.actor.Props
+import org.apache.texera.amber.core.virtualidentity.{
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkAck
+import org.apache.texera.amber.engine.architecture.controller.ReplayStatusUpdate
+import org.apache.texera.amber.engine.architecture.messaginglayer.WorkerTimerService
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.ControlInvocation
+import org.apache.texera.amber.engine.architecture.scheduling.config.WorkerConfig
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker._
+import org.apache.texera.amber.engine.common.actormessage.{ActorCommand, Backpressure}
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
+import org.apache.texera.amber.engine.common.{CheckpointState, SerializedState}
+
+import java.net.URI
+import java.util.concurrent.LinkedBlockingQueue
+import scala.collection.mutable
+
+object WorkflowWorker {
+ def props(
+ workerConfig: WorkerConfig,
+ replayInitialization: WorkerReplayInitialization
+ ): Props =
+ Props(
+ new WorkflowWorker(
+ workerConfig,
+ replayInitialization
+ )
+ )
+
+ final case class TriggerSend(msg: WorkflowFIFOMessage)
+
+ final case class MainThreadDelegateMessage(closure: WorkflowWorker => Unit)
+
+ sealed trait DPInputQueueElement
+
+ final case class FIFOMessageElement(msg: WorkflowFIFOMessage) extends DPInputQueueElement
+
+ final case class TimerBasedControlElement(control: ControlInvocation) extends DPInputQueueElement
+
+ final case class ActorCommandElement(cmd: ActorCommand) extends DPInputQueueElement
+
+ final case class WorkerReplayInitialization(
+ restoreConfOpt: Option[StateRestoreConfig] = None,
+ faultToleranceConfOpt: Option[FaultToleranceConfig] = None
+ )
+
+ final case class StateRestoreConfig(
+ readFrom: URI,
+ replayDestination: EmbeddedControlMessageIdentity
+ )
+
+ final case class FaultToleranceConfig(writeTo: URI)
+}
+
+class WorkflowWorker(
+ workerConfig: WorkerConfig,
+ replayInitialization: WorkerReplayInitialization
+) extends WorkflowActor(replayInitialization.faultToleranceConfOpt, workerConfig.workerId) {
+ val inputQueue: LinkedBlockingQueue[DPInputQueueElement] =
+ new LinkedBlockingQueue()
+ // Internal inputQueue is passed to dp.InputManager because input port materialization
+ // reader threads need to put input data read from materialization into this queue.
+ var dp = new DataProcessor(workerConfig.workerId, logManager.sendCommitted, inputQueue)
+ val timerService = new WorkerTimerService(actorService)
+
+ var dpThread: DPThread = _
+
+ val recordedInputs =
+ new mutable.HashMap[EmbeddedControlMessageIdentity, mutable.ArrayBuffer[WorkflowFIFOMessage]]()
+
+ override def initState(): Unit = {
+ dp.initTimerService(timerService)
+ if (replayInitialization.restoreConfOpt.isDefined) {
+ context.parent ! ReplayStatusUpdate(actorId, status = true)
+ setupReplay(
+ dp,
+ replayInitialization.restoreConfOpt.get,
+ () => {
+ logger.info("replay completed!")
+ context.parent ! ReplayStatusUpdate(actorId, status = false)
+ }
+ )
+ }
+ // dp is ready
+ dpThread = new DPThread(workerConfig.workerId, dp, logManager, inputQueue)
+ dpThread.start()
+ }
+
+ def handleDirectInvocation: Receive = {
+ case c: ControlInvocation =>
+ inputQueue.put(TimerBasedControlElement(c))
+ }
+
+ def handleTriggerClosure: Receive = {
+ case t: MainThreadDelegateMessage =>
+ t.closure(this)
+ }
+
+ def handleActorCommand: Receive = {
+ case c: ActorCommand =>
+ println(c)
+ }
+
+ override def preRestart(reason: Throwable, message: Option[Any]): Unit = {
+ super.preRestart(reason, message)
+ logger.error(s"Encountered fatal error, worker is shutting done.", reason)
+ postStop()
+ }
+
+ override def receive: Receive = {
+ super.receive orElse handleDirectInvocation orElse handleTriggerClosure
+ }
+
+ override def handleInputMessage(id: Long, workflowMsg: WorkflowFIFOMessage): Unit = {
+ inputQueue.put(FIFOMessageElement(workflowMsg))
+ recordedInputs.values.foreach(_.append(workflowMsg))
+ sender() ! NetworkAck(id, getInMemSize(workflowMsg), getQueuedCredit(workflowMsg.channelId))
+ }
+
+ /** flow-control */
+ override def getQueuedCredit(channelId: ChannelIdentity): Long = {
+ dp.getQueuedCredit(channelId)
+ }
+
+ override def postStop(): Unit = {
+ super.postStop()
+ timerService.stopAdaptiveBatching()
+ dpThread.stop()
+ logManager.terminate()
+ }
+
+ override def handleBackpressure(isBackpressured: Boolean): Unit = {
+ inputQueue.put(ActorCommandElement(Backpressure(isBackpressured)))
+ }
+
+ override def loadFromCheckpoint(chkpt: CheckpointState): Unit = {
+ logger.info("start loading from checkpoint.")
+ val inflightMessages: mutable.ArrayBuffer[WorkflowFIFOMessage] =
+ chkpt.load(SerializedState.IN_FLIGHT_MSG_KEY)
+ logger.info("inflight messages restored.")
+ val dpState: DataProcessor = chkpt.load(SerializedState.DP_STATE_KEY)
+ logger.info("dp state restored")
+ val queuedMessages: mutable.ArrayBuffer[WorkflowFIFOMessage] =
+ chkpt.load(SerializedState.DP_QUEUED_MSG_KEY)
+ logger.info("queued messages restored.")
+ val outputMessages: Array[WorkflowFIFOMessage] = chkpt.load(SerializedState.OUTPUT_MSG_KEY)
+ logger.info("output messages restored.")
+ dp = dpState // overwrite dp state
+ dp.outputHandler = logManager.sendCommitted
+ dp.initTimerService(timerService)
+ logger.info("start re-initialize executor from checkpoint.")
+ val (executor, iter) = dp.serializationManager.restoreExecutorState(chkpt)
+ dp.executor = executor
+ logger.info("re-initialize executor done.")
+ dp.outputManager.outputIterator.setTupleOutput(iter)
+ logger.info("set tuple output done.")
+ queuedMessages.foreach(msg => inputQueue.put(FIFOMessageElement(msg)))
+ inflightMessages.foreach(msg => inputQueue.put(FIFOMessageElement(msg)))
+ outputMessages.foreach(transferService.send)
+ logger.info("restored all messages done.")
+ context.parent ! ReplayStatusUpdate(actorId, status = false)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/InputPortMaterializationReaderThread.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/InputPortMaterializationReaderThread.scala
new file mode 100644
index 00000000000..10fbbc44a2c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/InputPortMaterializationReaderThread.scala
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.managers
+
+import io.grpc.MethodDescriptor
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.storage.model.VirtualDocument
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.architecture.messaginglayer.OutputManager.toPartitioner
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.{
+ NO_ALIGNMENT,
+ PORT_ALIGNMENT
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.{
+ METHOD_END_CHANNEL,
+ METHOD_START_CHANNEL
+}
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.Partitioning
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ DPInputQueueElement,
+ FIFOMessageElement
+}
+import org.apache.texera.amber.engine.common.ambermessage.{DataFrame, WorkflowFIFOMessage}
+import org.apache.texera.amber.util.VirtualIdentityUtils.getFromActorIdForInputPortStorage
+
+import java.net.URI
+import java.util.concurrent.LinkedBlockingQueue
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
+import scala.collection.mutable.ArrayBuffer
+
+class InputPortMaterializationReaderThread(
+ uri: URI,
+ inputMessageQueue: LinkedBlockingQueue[DPInputQueueElement],
+ workerActorId: ActorVirtualIdentity,
+ partitioning: Partitioning
+) extends Thread {
+
+ private val sequenceNum = new AtomicLong()
+ private val buffer = new ArrayBuffer[Tuple]()
+ private lazy val channelId = {
+ // A unique channel between this thread (dummy actor) and the worker actor.
+ val fromActorId: ActorVirtualIdentity =
+ getFromActorIdForInputPortStorage(uri.toString, workerActorId)
+ ChannelIdentity(fromActorId, workerActorId, isControl = false)
+ }
+ private val partitioner = toPartitioner(partitioning, workerActorId)
+ private val batchSize = ApplicationConfig.defaultDataTransferBatchSize
+ private val isFinished = new AtomicBoolean(false)
+
+ /**
+ * Whether the reader thread has completed.
+ */
+ def finished: Boolean = isFinished.get()
+
+ /**
+ * Read from the materialization stoage, and mimcs the behavior of an upstream worker's output manager.
+ */
+ override def run(): Unit = {
+ // Notify the input port of start of input channel
+ emitECM(METHOD_START_CHANNEL, NO_ALIGNMENT)
+ try {
+ val materialization: VirtualDocument[Tuple] = DocumentFactory
+ .openDocument(uri)
+ ._1
+ .asInstanceOf[VirtualDocument[Tuple]]
+ val storageReadIterator = materialization.get()
+ // Produce tuples
+ while (storageReadIterator.hasNext) {
+ val tuple = storageReadIterator.next()
+ if (
+ partitioner
+ .getBucketIndex(tuple)
+ .toList
+ .exists(bucketIndex => partitioner.allReceivers(bucketIndex) == workerActorId)
+ ) {
+ buffer.append(tuple)
+ if (buffer.size >= batchSize) {
+ flush()
+ }
+ }
+ }
+ // Flush any remaining tuples in the buffer.
+ if (buffer.nonEmpty) flush()
+ emitECM(METHOD_END_CHANNEL, PORT_ALIGNMENT)
+ isFinished.set(true)
+ } catch {
+ case e: Exception =>
+ throw new RuntimeException(s"Error reading input port materializations: ${e.getMessage}", e)
+ }
+ }
+
+ /**
+ * Puts an ECM into the internal queue.
+ */
+ private def emitECM(
+ method: MethodDescriptor[EmptyRequest, EmptyReturn],
+ alignment: EmbeddedControlMessageType
+ ): Unit = {
+ flush()
+ val ecm = EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity(method.getBareMethodName),
+ alignment,
+ Seq(),
+ Map(
+ workerActorId.name ->
+ ControlInvocation(
+ method.getBareMethodName,
+ EmptyRequest(),
+ AsyncRPCContext(ActorVirtualIdentity(""), ActorVirtualIdentity("")),
+ -1
+ )
+ )
+ )
+ val fifoMessage = WorkflowFIFOMessage(channelId, getSequenceNumber, ecm)
+ val inputQueueElement = FIFOMessageElement(fifoMessage)
+ inputMessageQueue.put(inputQueueElement)
+ }
+
+ /**
+ * Flush the current batch into a DataFrame and enqueue it.
+ */
+ private def flush(): Unit = {
+ if (buffer.isEmpty) return
+ val dataPayload = DataFrame(buffer.toArray) // Mimics flush logic in NetworkOutputBuffer.
+ val fifoMessage = WorkflowFIFOMessage(channelId, getSequenceNumber, dataPayload)
+ val inputQueueElement = FIFOMessageElement(fifoMessage)
+ inputMessageQueue.put(inputQueueElement)
+ buffer.clear()
+ }
+
+ private def getSequenceNumber = {
+ sequenceNum.getAndIncrement()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/OutputPortResultWriterThread.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/OutputPortResultWriterThread.scala
new file mode 100644
index 00000000000..28e5d2af667
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/OutputPortResultWriterThread.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.managers
+
+import com.google.common.collect.Queues
+import org.apache.texera.amber.core.storage.model.BufferedItemWriter
+import org.apache.texera.amber.core.tuple.Tuple
+
+import java.util.concurrent.LinkedBlockingQueue
+
+sealed trait TerminateSignal
+case object PortStorageWriterTerminateSignal extends TerminateSignal
+
+class OutputPortResultWriterThread(
+ bufferedItemWriter: BufferedItemWriter[Tuple]
+) extends Thread {
+
+ val queue: LinkedBlockingQueue[Either[Tuple, TerminateSignal]] =
+ Queues.newLinkedBlockingQueue[Either[Tuple, TerminateSignal]]()
+
+ override def run(): Unit = {
+ var internalStop = false
+ while (!internalStop) {
+ val queueContent = queue.take()
+ queueContent match {
+ case Left(tuple) => bufferedItemWriter.putOne(tuple)
+ case Right(_) => internalStop = true
+ }
+ }
+ bufferedItemWriter.close()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/SerializationManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/SerializationManager.scala
new file mode 100644
index 00000000000..b4afe510306
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/SerializationManager.scala
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.managers
+
+import org.apache.texera.amber.core.executor._
+import org.apache.texera.amber.core.tuple.TupleLike
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.InitializeExecutorRequest
+import org.apache.texera.amber.engine.common.{AmberLogging, CheckpointState, CheckpointSupport}
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+class SerializationManager(val actorId: ActorVirtualIdentity) extends AmberLogging {
+
+ @transient private var serializationCall: () => Unit = _
+ private var execInitMsg: InitializeExecutorRequest = _
+
+ def setOpInitialization(msg: InitializeExecutorRequest): Unit = {
+ execInitMsg = msg
+ }
+
+ def restoreExecutorState(
+ chkpt: CheckpointState
+ ): (OperatorExecutor, Iterator[(TupleLike, Option[PortIdentity])]) = {
+ val workerIdx = VirtualIdentityUtils.getWorkerIndex(actorId)
+ val workerCount = execInitMsg.totalWorkerCount
+ val executor = execInitMsg.opExecInitInfo match {
+ case OpExecWithClassName(className, descString) =>
+ ExecFactory.newExecFromJavaClassName(className, descString, workerIdx, workerCount)
+ case OpExecWithCode(code, language) => ExecFactory.newExecFromJavaCode(code)
+ case _ => throw new UnsupportedOperationException("Unsupported OpExec type")
+ }
+
+ val iter = executor match {
+ case support: CheckpointSupport =>
+ support.deserializeState(chkpt)
+ case _ => Iterator.empty
+ }
+ (executor, iter)
+ }
+
+ def registerSerialization(call: () => Unit): Unit = {
+ serializationCall = call
+ }
+
+ def applySerialization(): Unit = {
+ if (serializationCall != null) {
+ serializationCall()
+ serializationCall = null
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/StatisticsManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/StatisticsManager.scala
new file mode 100644
index 00000000000..8ae0419f0a3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/StatisticsManager.scala
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.managers
+
+import org.apache.texera.amber.core.executor.OperatorExecutor
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.worker.statistics.{
+ PortTupleMetricsMapping,
+ TupleMetrics,
+ WorkerStatistics
+}
+
+import scala.collection.mutable
+
+class StatisticsManager {
+ // DataProcessor
+ private val inputStatistics: mutable.Map[PortIdentity, (Long, Long)] =
+ mutable.Map.empty.withDefaultValue((0L, 0L))
+ private val outputStatistics: mutable.Map[PortIdentity, (Long, Long)] =
+ mutable.Map.empty.withDefaultValue((0L, 0L))
+ private var dataProcessingTime: Long = 0L
+ private var totalExecutionTime: Long = 0L
+ private var workerStartTime: Long = 0L
+
+ // AmberProcessor
+ private var controlProcessingTime: Long = 0L
+
+ /**
+ * Retrieves the current statistics for the operator.
+ * @param operator the operator executor
+ * @return a WorkerStatistics object containing the statistics
+ */
+ def getStatistics(operator: OperatorExecutor): WorkerStatistics = {
+ WorkerStatistics(
+ inputStatistics.map {
+ case (portId, (tupleCount, tupleSize)) =>
+ PortTupleMetricsMapping(portId, TupleMetrics(tupleCount, tupleSize))
+ }.toSeq,
+ outputStatistics.map {
+ case (portId, (tupleCount, tupleSize)) =>
+ PortTupleMetricsMapping(portId, TupleMetrics(tupleCount, tupleSize))
+ }.toSeq,
+ dataProcessingTime,
+ controlProcessingTime,
+ totalExecutionTime - dataProcessingTime - controlProcessingTime
+ )
+ }
+
+ /**
+ * Calculates the total number of input tuples.
+ * @return the total input tuple count
+ */
+ def getInputTupleCount: Long = inputStatistics.values.map(_._1).sum
+
+ /**
+ * Calculates the total number of output tuples.
+ * @return the total output tuple count
+ */
+ def getOutputTupleCount: Long = outputStatistics.values.map(_._1).sum
+
+ /**
+ * Increases the input statistics for a given port.
+ * @param portId the port identity
+ * @param size the size of the tuple
+ */
+ def increaseInputStatistics(portId: PortIdentity, size: Long): Unit = {
+ require(size >= 0, "Tuple size must be non-negative")
+ val (count, totalSize) = inputStatistics(portId)
+ inputStatistics.update(portId, (count + 1, totalSize + size))
+ }
+
+ /**
+ * Increases the output statistics for a given port.
+ * @param portId the port identity
+ * @param size the size of the tuple
+ */
+ def increaseOutputStatistics(portId: PortIdentity, size: Long): Unit = {
+ require(size >= 0, "Tuple size must be non-negative")
+ val (count, totalSize) = outputStatistics(portId)
+ outputStatistics.update(portId, (count + 1, totalSize + size))
+ }
+
+ /**
+ * Increases the data processing time.
+ * @param time the time to add
+ */
+ def increaseDataProcessingTime(time: Long): Unit = {
+ require(time >= 0, "Time must be non-negative")
+ dataProcessingTime += time
+ }
+
+ /**
+ * Increases the control processing time.
+ * @param time the time to add
+ */
+ def increaseControlProcessingTime(time: Long): Unit = {
+ require(time >= 0, "Time must be non-negative")
+ controlProcessingTime += time
+ }
+
+ /**
+ * Updates the total execution time.
+ * @param time the current time
+ */
+ def updateTotalExecutionTime(time: Long): Unit = {
+ require(
+ time >= workerStartTime,
+ "Current time must be greater than or equal to worker start time"
+ )
+ totalExecutionTime = time - workerStartTime
+ }
+
+ /**
+ * Initializes the worker start time.
+ * @param time the start time
+ */
+ def initializeWorkerStartTime(time: Long): Unit = {
+ workerStartTime = time
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AddInputChannelHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AddInputChannelHandler.scala
new file mode 100644
index 00000000000..685f244aeee
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AddInputChannelHandler.scala
@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AddInputChannelRequest,
+ AsyncRPCContext
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.{
+ PAUSED,
+ READY,
+ RUNNING
+}
+
+trait AddInputChannelHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def addInputChannel(
+ msg: AddInputChannelRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ dp.inputGateway.getChannel(msg.channelId).setPortId(msg.portId)
+ dp.inputManager.getPort(msg.portId).channels.add(msg.channelId)
+ dp.stateManager.assertState(READY, RUNNING, PAUSED)
+ EmptyReturn()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AddPartitioningHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AddPartitioningHandler.scala
new file mode 100644
index 00000000000..c2da82214bd
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AddPartitioningHandler.scala
@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AddPartitioningRequest,
+ AsyncRPCContext
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.{
+ PAUSED,
+ READY,
+ RUNNING
+}
+
+trait AddPartitioningHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def addPartitioning(
+ msg: AddPartitioningRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ dp.stateManager.assertState(READY, RUNNING, PAUSED)
+ dp.outputManager.addPartitionerWithPartitioning(msg.tag, msg.partitioning)
+ EmptyReturn()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AssignPortHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AssignPortHandler.scala
new file mode 100644
index 00000000000..fe959733abb
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/AssignPortHandler.scala
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.tuple.Schema
+import org.apache.texera.amber.core.virtualidentity.ChannelIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AssignPortRequest,
+ AsyncRPCContext
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.{
+ PAUSED,
+ READY,
+ RUNNING
+}
+import org.apache.texera.amber.util.VirtualIdentityUtils.getFromActorIdForInputPortStorage
+
+import java.net.URI
+
+trait AssignPortHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def assignPort(msg: AssignPortRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = {
+ val schema = Schema.fromRawSchema(msg.schema)
+ if (msg.input) {
+ val inputPortURIStrs = msg.storageUris.toList
+ val inputPortURIs = inputPortURIStrs.map(uriStr => URI.create(uriStr))
+ val partitionings = msg.partitionings.toList
+ dp.inputManager.addPort(msg.portId, schema, inputPortURIs, partitionings)
+ inputPortURIStrs.foreach { uriStr =>
+ val toActorId = ctx.receiver
+ val fromActorId = getFromActorIdForInputPortStorage(uriStr, toActorId)
+ val channelId =
+ ChannelIdentity(fromWorkerId = fromActorId, toWorkerId = toActorId, isControl = false)
+ // Same as AddInputChannelHandler
+ dp.inputGateway.getChannel(channelId).setPortId(msg.portId)
+ dp.inputManager.getPort(msg.portId).channels.add(channelId)
+ dp.stateManager.assertState(READY, RUNNING, PAUSED)
+ }
+ } else {
+ val storageURIOption: Option[URI] = msg.storageUris.head match {
+ case "" => None
+ case uriString => Some(URI.create(uriString))
+ }
+ dp.outputManager.addPort(msg.portId, schema, storageURIOption)
+ }
+ EmptyReturn()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndChannelHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndChannelHandler.scala
new file mode 100644
index 00000000000..7794342690b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndChannelHandler.scala
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.tuple.FinalizePort
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.error.ErrorUtils.safely
+
+trait EndChannelHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def endChannel(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ val channelId = dp.inputManager.currentChannelId
+ val portId = dp.inputGateway.getChannel(channelId).getPortId
+ dp.inputManager.getPort(portId).completed = true
+ dp.inputManager.initBatch(channelId, Array.empty)
+ try {
+ val outputState = dp.executor.produceStateOnFinish(portId.id)
+ if (outputState.isDefined) {
+ dp.outputManager.emitState(outputState.get)
+ }
+ dp.outputManager.outputIterator.setTupleOutput(
+ dp.executor.onFinishMultiPort(portId.id)
+ )
+ } catch safely {
+ case e =>
+ // forward input tuple to the user and pause DP thread
+ dp.handleExecutorException(e)
+ }
+
+ dp.outputManager.outputIterator.appendSpecialTupleToEnd(
+ FinalizePort(portId, input = true)
+ )
+
+ if (dp.inputManager.getAllPorts.forall(portId => dp.inputManager.isPortCompleted(portId))) {
+ // Need this check for handling input port dependency relationships.
+ // See documentation of isMissingOutputPort
+ if (!dp.outputManager.isMissingOutputPort) {
+ // assuming all the output ports finalize after all input ports are finalized.
+ dp.outputManager.finalizeOutput()
+ }
+ }
+ EmptyReturn()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala
new file mode 100644
index 00000000000..0504e66f52b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+
+/**
+ * The EndWorker control messages is needed to ensure all the other control messages in a worker
+ * are processed before worker termination.
+ */
+trait EndHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ /**
+ * The response of endWorker to the controller indicates that this worker has finished not only
+ * the data processing logic, but also , but also the processing of all the control messages.
+ */
+ override def endWorker(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ // Ensure this is really the last message.
+ if (!dp.inputManager.inputMessageQueue.isEmpty) {
+ logger.warn(
+ s"Received EndHandler before all messages are processed. Unprocessed messages: " +
+ s"${dp.inputManager.inputMessageQueue.peek()}"
+ )
+ return Future.exception(new IllegalStateException("worker still has unprocessed messages"))
+ }
+ // Now we can safely acknowledge that this worker can be terminated.
+ EmptyReturn()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FinalizeCheckpointHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FinalizeCheckpointHandler.scala
new file mode 100644
index 00000000000..96d7b4dfaff
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FinalizeCheckpointHandler.scala
@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ FinalizeCheckpointRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.FinalizeCheckpointResponse
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage
+import org.apache.texera.amber.engine.architecture.worker.{
+ DataProcessorRPCHandlerInitializer,
+ WorkflowWorker
+}
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.engine.common.{CheckpointState, CheckpointSupport, SerializedState}
+
+import java.net.URI
+import java.util.concurrent.CompletableFuture
+import scala.collection.mutable.ArrayBuffer
+
+trait FinalizeCheckpointHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def finalizeCheckpoint(
+ msg: FinalizeCheckpointRequest,
+ ctx: AsyncRPCContext
+ ): Future[FinalizeCheckpointResponse] = {
+ val checkpointSize = if (dp.ecmManager.checkpoints.contains(msg.checkpointId)) {
+ val waitFuture = new CompletableFuture[Unit]()
+ val chkpt = dp.ecmManager.checkpoints(msg.checkpointId)
+ val closure = (worker: WorkflowWorker) => {
+ logger.info(s"Main thread: start to serialize recorded messages.")
+ chkpt.save(
+ SerializedState.IN_FLIGHT_MSG_KEY,
+ worker.recordedInputs.getOrElse(msg.checkpointId, new ArrayBuffer())
+ )
+ worker.recordedInputs.remove(msg.checkpointId)
+ logger.info(s"Main thread: recorded messages serialized.")
+ waitFuture.complete(())
+ ()
+ }
+ // TODO: find a way to skip logging for the following output?
+ dp.outputHandler(
+ Left(MainThreadDelegateMessage(closure))
+ ) //this will create duplicate log records!
+ waitFuture.get()
+ logger.info(s"Start to write checkpoint to storage. Destination: ${msg.writeTo}")
+ val storage = SequentialRecordStorage.getStorage[CheckpointState](Some(new URI(msg.writeTo)))
+ val writer = storage.getWriter(actorId.name.replace("Worker:", ""))
+ writer.writeRecord(chkpt)
+ writer.flush()
+ writer.close()
+ logger.info(s"Checkpoint finalized, total size = ${chkpt.size()} bytes")
+ chkpt.size()
+ } else {
+ logger.info(s"Checkpoint is estimation-only. report estimated size.")
+ dp.executor match {
+ case support: CheckpointSupport =>
+ support.getEstimatedCheckpointCost
+ case _ => 0L
+ } // for estimation
+ }
+ FinalizeCheckpointResponse(checkpointSize)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FlushNetworkBufferHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FlushNetworkBufferHandler.scala
new file mode 100644
index 00000000000..446ac5b3cdc
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FlushNetworkBufferHandler.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+
+trait FlushNetworkBufferHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def flushNetworkBuffer(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ dp.outputManager.flush()
+ EmptyReturn()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/InitializeExecutorHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/InitializeExecutorHandler.scala
new file mode 100644
index 00000000000..212a980e5ed
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/InitializeExecutorHandler.scala
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ InitializeExecutorRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+trait InitializeExecutorHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def initializeExecutor(
+ req: InitializeExecutorRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ dp.serializationManager.setOpInitialization(req)
+ val workerIdx = VirtualIdentityUtils.getWorkerIndex(actorId)
+ cachedTotalWorkerCount = req.totalWorkerCount
+ setupExecutor(req.opExecInitInfo, workerIdx, cachedTotalWorkerCount)
+ EmptyReturn()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/OpenExecutorHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/OpenExecutorHandler.scala
new file mode 100644
index 00000000000..81298a58d88
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/OpenExecutorHandler.scala
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+
+trait OpenExecutorHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def openExecutor(request: EmptyRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = {
+ dp.executor.open()
+ EmptyReturn()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PauseHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PauseHandler.scala
new file mode 100644
index 00000000000..cec7ca87e63
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PauseHandler.scala
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkerStateResponse
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.{
+ PAUSED,
+ READY,
+ RUNNING
+}
+import org.apache.texera.amber.engine.architecture.worker.{
+ DataProcessorRPCHandlerInitializer,
+ UserPause
+}
+
+trait PauseHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def pauseWorker(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[WorkerStateResponse] = {
+ if (dp.stateManager.confirmState(RUNNING, READY)) {
+ dp.pauseManager.pause(UserPause)
+ dp.stateManager.transitTo(PAUSED)
+ }
+ WorkerStateResponse(dp.stateManager.getCurrentState)
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PrepareCheckpointHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PrepareCheckpointHandler.scala
new file mode 100644
index 00000000000..3333f54e90d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PrepareCheckpointHandler.scala
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.EmbeddedControlMessageIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ PrepareCheckpointRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage
+import org.apache.texera.amber.engine.architecture.worker.{
+ DataProcessorRPCHandlerInitializer,
+ WorkflowWorker
+}
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.{CheckpointState, CheckpointSupport, SerializedState}
+
+import java.util.concurrent.CompletableFuture
+import scala.collection.mutable
+
+trait PrepareCheckpointHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def prepareCheckpoint(
+ msg: PrepareCheckpointRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ logger.info("Start to take checkpoint")
+ if (!msg.estimationOnly) {
+ dp.serializationManager.registerSerialization(() => {
+ serializeWorkerState(msg.checkpointId)
+ })
+ } else {
+ logger.info(s"Checkpoint is estimation-only. do nothing.")
+ }
+ EmptyReturn()
+ }
+
+ private def serializeWorkerState(checkpointId: EmbeddedControlMessageIdentity): Unit = {
+ val chkpt = new CheckpointState()
+ // 1. serialize DP state
+ chkpt.save(SerializedState.DP_STATE_KEY, this.dp)
+ // checkpoint itself should not be serialized, thus we register it after serialization
+ dp.ecmManager.checkpoints(checkpointId) = chkpt
+ logger.info("Serialized DP state")
+ // 2. serialize operator state
+ dp.executor match {
+ case support: CheckpointSupport =>
+ dp.outputManager.outputIterator.setTupleOutput(
+ support.serializeState(dp.outputManager.outputIterator.outputIter, chkpt)
+ )
+ logger.info("Serialized operator state")
+ case _ =>
+ logger.info("Operator does not support checkpoint, skip")
+ }
+ // 3. record inflight messages
+ logger.info("Begin collecting inflight messages")
+ val waitFuture = new CompletableFuture[Unit]()
+ val closure = (worker: WorkflowWorker) => {
+ val queuedMsgs = mutable.ArrayBuffer[WorkflowFIFOMessage]()
+ worker.inputQueue.forEach {
+ case WorkflowWorker.FIFOMessageElement(msg) => queuedMsgs.append(msg)
+ case WorkflowWorker.TimerBasedControlElement(control) => // skip
+ case WorkflowWorker.ActorCommandElement(cmd) => // skip
+ }
+ chkpt.save(SerializedState.DP_QUEUED_MSG_KEY, queuedMsgs)
+ // get all output messages from worker.transferService
+ chkpt.save(
+ SerializedState.OUTPUT_MSG_KEY,
+ worker.transferService.getAllUnAckedMessages.toArray
+ )
+ logger.info("Main thread: serialized queued and output messages.")
+ // start to record input messages on main thread
+ worker.recordedInputs(checkpointId) = new mutable.ArrayBuffer[WorkflowFIFOMessage]()
+ logger.info("Main thread: start recording for input messages from now on.")
+ waitFuture.complete(())
+ ()
+ }
+ dp.outputHandler(Left(MainThreadDelegateMessage(closure)))
+ waitFuture.get()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/QueryStatisticsHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/QueryStatisticsHandler.scala
new file mode 100644
index 00000000000..74d8d14faa6
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/QueryStatisticsHandler.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkerMetricsResponse
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerMetrics
+
+trait QueryStatisticsHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def queryStatistics(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[WorkerMetricsResponse] = {
+ WorkerMetricsResponse(WorkerMetrics(dp.stateManager.getCurrentState, dp.collectStatistics()))
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/ResumeHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/ResumeHandler.scala
new file mode 100644
index 00000000000..434c50c914c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/ResumeHandler.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkerStateResponse
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.{PAUSED, RUNNING}
+import org.apache.texera.amber.engine.architecture.worker.{
+ DataProcessorRPCHandlerInitializer,
+ UserPause
+}
+
+trait ResumeHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def resumeWorker(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[WorkerStateResponse] = {
+ if (dp.stateManager.getCurrentState == PAUSED) {
+ dp.pauseManager.resume(UserPause)
+ dp.stateManager.transitTo(RUNNING)
+ dp.adaptiveBatchingMonitor.resumeAdaptiveBatching()
+ }
+ WorkerStateResponse(dp.stateManager.getCurrentState)
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/RetrieveStateHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/RetrieveStateHandler.scala
new file mode 100644
index 00000000000..2b0721a2525
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/RetrieveStateHandler.scala
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+
+trait RetrieveStateHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def retrieveState(request: EmptyRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = {
+ EmptyReturn() // TODO: add implementation
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartChannelHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartChannelHandler.scala
new file mode 100644
index 00000000000..01cbb858bd7
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartChannelHandler.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.NO_ALIGNMENT
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_START_CHANNEL
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.error.ErrorUtils.safely
+
+trait StartChannelHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def startChannel(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ val portId = dp.inputGateway.getChannel(dp.inputManager.currentChannelId).getPortId
+ dp.sendECMToDataChannels(METHOD_START_CHANNEL, NO_ALIGNMENT)
+ try {
+ val outputState = dp.executor.produceStateOnStart(portId.id)
+ if (outputState.isDefined) {
+ dp.outputManager.emitState(outputState.get)
+ }
+ } catch safely {
+ case e =>
+ dp.handleExecutorException(e)
+ }
+ EmptyReturn()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartHandler.scala
new file mode 100644
index 00000000000..5d1bf8ccbde
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartHandler.scala
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.core.executor.SourceOperatorExecutor
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkerStateResponse
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.{READY, RUNNING}
+
+trait StartHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def startWorker(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[WorkerStateResponse] = {
+ logger.info("Starting the worker.")
+ if (dp.executor.isInstanceOf[SourceOperatorExecutor]) {
+ val channelId =
+ ChannelIdentity(ActorVirtualIdentity("SOURCE_STARTER"), actorId, isControl = false)
+ dp.stateManager.assertState(READY)
+ dp.stateManager.transitTo(RUNNING)
+ // for source operator: add a virtual input channel just for kicking off the execution
+ dp.inputManager.addPort(
+ PortIdentity(),
+ null,
+ urisToRead = List.empty,
+ partitionings = List.empty
+ )
+ dp.inputManager.currentChannelId = channelId
+ dp.inputGateway.getChannel(channelId).setPortId(PortIdentity())
+ startChannel(request, ctx)
+ endChannel(request, ctx)
+ WorkerStateResponse(dp.stateManager.getCurrentState)
+ } else if (dp.inputManager.getInputPortReaderThreads.nonEmpty) {
+ // This means the worker should read from materialized storage for its input ports.
+ // Start the reader threads
+ dp.inputManager.startInputPortReaderThreads()
+ WorkerStateResponse(dp.stateManager.getCurrentState)
+ } else {
+ throw new WorkflowRuntimeException(
+ s"non-source worker $actorId received unexpected StartWorker!"
+ )
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/UpdateExecutorHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/UpdateExecutorHandler.scala
new file mode 100644
index 00000000000..8ed9ebdc595
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/UpdateExecutorHandler.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ UpdateExecutorRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+trait UpdateExecutorHandler {
+ this: DataProcessorRPCHandlerInitializer =>
+
+ override def updateExecutor(
+ request: UpdateExecutorRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ val workerIdx = VirtualIdentityUtils.getWorkerIndex(actorId)
+ // Close the existing executor (if any) before replacing it to avoid resource leaks.
+ val oldExecutor = dp.executor
+ if (oldExecutor != null) {
+ oldExecutor.close()
+ }
+ setupExecutor(request.newExecInitInfo, workerIdx, cachedTotalWorkerCount)
+ dp.executor.open()
+ EmptyReturn()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberConfig.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberConfig.scala
new file mode 100644
index 00000000000..2b05cc771af
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberConfig.scala
@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import org.apache.pekko.actor.Address
+
+object AmberConfig {
+ var masterNodeAddr: Address = Address("pekko", "Amber", "localhost", 2552)
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberKryoInitializer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberKryoInitializer.scala
new file mode 100644
index 00000000000..1721764d414
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberKryoInitializer.scala
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import com.esotericsoftware.kryo.kryo5.serializers.ClosureSerializer
+import io.altoo.serialization.kryo.pekko.DefaultKryoInitializer
+import io.altoo.serialization.kryo.scala.serializer.ScalaKryo
+
+import java.lang.invoke.SerializedLambda
+
+class AmberKryoInitializer extends DefaultKryoInitializer {
+ override def preInit(kryo: ScalaKryo): Unit = {
+ kryo.register(classOf[SerializedLambda])
+ kryo.register(classOf[ClosureSerializer.Closure], new ClosureSerializer())
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberLogging.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberLogging.scala
new file mode 100644
index 00000000000..0fd16475d9f
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberLogging.scala
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import com.typesafe.scalalogging.Logger
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.util.VirtualIdentityUtils
+import org.slf4j.LoggerFactory
+
+trait AmberLogging {
+
+ @transient
+ protected lazy val logger: Logger = Logger(
+ LoggerFactory.getLogger(
+ s"${VirtualIdentityUtils.toShorterString(actorId)}] [${getClass.getSimpleName}"
+ )
+ )
+
+ def actorId: ActorVirtualIdentity
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberRuntime.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberRuntime.scala
new file mode 100644
index 00000000000..7078f766a63
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/AmberRuntime.scala
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import org.apache.pekko.actor.{ActorSystem, Address, Cancellable, DeadLetter, Props}
+import org.apache.pekko.serialization.{Serialization, SerializationExtension}
+import com.typesafe.config.{Config, ConfigFactory}
+import org.apache.texera.amber.clustering.ClusterListener
+import org.apache.texera.amber.config.AkkaConfig
+import org.apache.texera.amber.engine.architecture.messaginglayer.DeadLetterMonitorActor
+
+import java.io.{BufferedReader, InputStreamReader}
+import java.net.URL
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration.FiniteDuration
+
+object AmberRuntime {
+
+ private var _serde: Serialization = _
+ private var _actorSystem: ActorSystem = _
+
+ def serde: Serialization = {
+ if (_serde == null) {
+ if (_actorSystem == null) {
+ _serde = SerializationExtension(ActorSystem("Amber", akkaConfig))
+ } else {
+ _serde = SerializationExtension(_actorSystem)
+ }
+ }
+ _serde
+ }
+
+ def actorSystem: ActorSystem = {
+ _actorSystem
+ }
+
+ def scheduleCallThroughActorSystem(delay: FiniteDuration)(call: => Unit): Cancellable = {
+ _actorSystem.scheduler.scheduleOnce(delay)(call)
+ }
+
+ def scheduleRecurringCallThroughActorSystem(initialDelay: FiniteDuration, delay: FiniteDuration)(
+ call: => Unit
+ ): Cancellable = {
+ _actorSystem.scheduler.scheduleWithFixedDelay(initialDelay, delay)(() => call)
+ }
+
+ private def getNodeIpAddress: String = {
+ try {
+ val query = new URL("http://checkip.amazonaws.com")
+ val in = new BufferedReader(new InputStreamReader(query.openStream()))
+ in.readLine()
+ } catch {
+ case e: Exception => throw e
+ }
+ }
+
+ def startActorMaster(clusterMode: Boolean): Unit = {
+ var localIpAddress = "localhost"
+ if (clusterMode) {
+ localIpAddress = getNodeIpAddress
+ }
+
+ val masterConfig = ConfigFactory
+ .parseString(s"""
+ pekko.remote.artery.canonical.port = 2552
+ pekko.remote.artery.canonical.hostname = $localIpAddress
+ pekko.cluster.seed-nodes = [ "pekko://Amber@$localIpAddress:2552" ]
+ """)
+ .withFallback(akkaConfig)
+ .resolve()
+ AmberConfig.masterNodeAddr = createMasterAddress(localIpAddress)
+ createAmberSystem(masterConfig)
+ }
+
+ def akkaConfig: Config = AkkaConfig.akkaConfig
+
+ private def createMasterAddress(addr: String): Address = Address("pekko", "Amber", addr, 2552)
+
+ def startActorWorker(mainNodeAddress: Option[String]): Unit = {
+ val addr = mainNodeAddress.getOrElse("localhost")
+ var localIpAddress = "localhost"
+ if (mainNodeAddress.isDefined) {
+ localIpAddress = getNodeIpAddress
+ }
+ val workerConfig = ConfigFactory
+ .parseString(s"""
+ pekko.remote.artery.canonical.hostname = $localIpAddress
+ pekko.remote.artery.canonical.port = 0
+ pekko.cluster.seed-nodes = [ "pekko://Amber@$addr:2552" ]
+ """)
+ .withFallback(akkaConfig)
+ .resolve()
+ AmberConfig.masterNodeAddr = createMasterAddress(addr)
+ createAmberSystem(workerConfig)
+ }
+
+ private def createAmberSystem(actorSystemConf: Config): Unit = {
+ _actorSystem = ActorSystem("Amber", actorSystemConf)
+ _actorSystem.actorOf(Props[ClusterListener](), "cluster-info")
+ val deadLetterMonitorActor =
+ _actorSystem.actorOf(Props[DeadLetterMonitorActor](), name = "dead-letter-monitor-actor")
+ _actorSystem.eventStream.subscribe(deadLetterMonitorActor, classOf[DeadLetter])
+ _serde = SerializationExtension(_actorSystem)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/CheckpointState.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/CheckpointState.scala
new file mode 100644
index 00000000000..0f1159a2ef7
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/CheckpointState.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import scala.collection.mutable
+
+class CheckpointState {
+
+ private val states = new mutable.HashMap[String, SerializedState]()
+
+ def save[T <: Any](key: String, state: T): Unit = {
+ states(key) = SerializedState.fromObject(state.asInstanceOf[AnyRef], AmberRuntime.serde)
+ }
+
+ def has(key: String): Boolean = {
+ states.contains(key)
+ }
+
+ def load[T <: Any](key: String): T = {
+ if (states.contains(key)) {
+ states(key).toObject(AmberRuntime.serde).asInstanceOf[T]
+ } else {
+ throw new RuntimeException(s"no state saved for key = $key")
+ }
+ }
+
+ def size(): Long = {
+ states.filter(_._2 != null).map(_._2.size()).sum
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/CheckpointSupport.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/CheckpointSupport.scala
new file mode 100644
index 00000000000..d409e5da91e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/CheckpointSupport.scala
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import org.apache.texera.amber.core.tuple.TupleLike
+import org.apache.texera.amber.core.workflow.PortIdentity
+
+trait CheckpointSupport {
+ def serializeState(
+ currentIteratorState: Iterator[(TupleLike, Option[PortIdentity])],
+ checkpoint: CheckpointState
+ ): Iterator[(TupleLike, Option[PortIdentity])]
+
+ def deserializeState(
+ checkpoint: CheckpointState
+ ): Iterator[(TupleLike, Option[PortIdentity])]
+
+ def getEstimatedCheckpointCost: Long
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/ElidableStatement.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/ElidableStatement.scala
new file mode 100644
index 00000000000..bf78b290c31
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/ElidableStatement.scala
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import scala.annotation.elidable
+import scala.annotation.elidable._
+
+object ElidableStatement {
+
+ @elidable(FINEST) def finest(operations: => Unit): Unit = operations
+
+ @elidable(FINER) def finer(operations: => Unit): Unit = operations
+
+ @elidable(FINE) def fine(operations: => Unit): Unit = operations
+
+ @elidable(INFO) def info(operations: => Unit): Unit = operations
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/FriesReconfigurationAlgorithm.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/FriesReconfigurationAlgorithm.scala
new file mode 100644
index 00000000000..c13e7801190
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/FriesReconfigurationAlgorithm.scala
@@ -0,0 +1,127 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
+import org.apache.texera.amber.core.workflow.PhysicalPlan
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ UpdateExecutorRequest,
+ WorkflowReconfigureRequest
+}
+import org.apache.texera.amber.engine.architecture.scheduling.{Region, WorkflowExecutionCoordinator}
+import org.jgrapht.alg.connectivity.ConnectivityInspector
+
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters.SetHasAsScala
+
+object FriesReconfigurationAlgorithm {
+
+ case class FriesComponent(
+ sources: Set[PhysicalOpIdentity],
+ scope: Set[PhysicalOpIdentity],
+ reconfigurations: Set[UpdateExecutorRequest]
+ )
+
+ private def getOneToManyOperators(region: Region): Set[PhysicalOpIdentity] = {
+ region.getOperators.filter(op => op.isOneToManyOp).map(op => op.id)
+ }
+
+ def getReconfigurations(
+ workflowExecutionCoordinator: WorkflowExecutionCoordinator,
+ reconfiguration: WorkflowReconfigureRequest
+ ): Set[FriesComponent] = {
+ // independently schedule reconfigurations for each region:
+ workflowExecutionCoordinator.getExecutingRegions
+ .flatMap(region => computeMCS(region, reconfiguration, reconfiguration.reconfigurationId))
+ }
+
+ private def computeMCS(
+ region: Region,
+ reconfiguration: WorkflowReconfigureRequest,
+ epochMarkerId: String
+ ): List[FriesComponent] = {
+
+ // add all reconfiguration operators to M
+ val reconfigOps = reconfiguration.reconfiguration.map(req => req.targetOpId).toSet
+ val M = mutable.Set.empty ++ reconfigOps
+
+ // for each one-to-many operator, add it to M if its downstream has a reconfiguration operator
+ val oneToManyOperators = getOneToManyOperators(region)
+ oneToManyOperators.foreach(oneToManyOp => {
+ val intersection = region.dag.getDescendants(oneToManyOp).asScala.intersect(reconfigOps)
+ if (intersection.nonEmpty) {
+ M += oneToManyOp
+ }
+ })
+
+ // compute MCS based on M
+ var forwardVertices: Set[PhysicalOpIdentity] = Set()
+ var backwardVertices: Set[PhysicalOpIdentity] = Set()
+
+ val topologicalOps = region.topologicalIterator().toList
+ val reverseTopologicalOps = topologicalOps.reverse
+
+ topologicalOps.foreach(opId => {
+ val op = region.getOperator(opId)
+ val parents = op.inputPorts.flatMap(_._2._2).map(_.fromOpId)
+ val fromParent: Boolean = parents.exists(p => forwardVertices.contains(p))
+ if (M.contains(opId) || fromParent) {
+ forwardVertices += opId
+ }
+ })
+
+ reverseTopologicalOps.foreach(opId => {
+ val op = region.getOperator(opId)
+ val children = op.outputPorts.flatMap(_._2._2).map(_.toOpId)
+ val fromChildren: Boolean = children.exists(p => backwardVertices.contains(p))
+ if (M.contains(opId) || fromChildren) {
+ backwardVertices += opId
+ }
+ })
+
+ val resultMCSOpIds = forwardVertices.intersect(backwardVertices)
+ val newLinks =
+ region.getLinks.filter(link =>
+ resultMCSOpIds.contains(link.fromOpId) && resultMCSOpIds.contains(link.toOpId)
+ )
+ val mcsPlan = PhysicalPlan(resultMCSOpIds.map(opId => region.getOperator(opId)), newLinks)
+
+ // find the MCS components,
+ // for each component, send an epoch marker to each of its source operators
+ val epochMarkers = new ArrayBuffer[FriesComponent]()
+
+ val connectedSets = new ConnectivityInspector(mcsPlan.dag).connectedSets()
+ connectedSets.forEach(component => {
+ val componentSet = component.asScala.toSet
+ val componentPlan = mcsPlan.getSubPlan(componentSet)
+ val reconfigCommands =
+ reconfiguration.reconfiguration
+ .filter(req => component.contains(req.targetOpId))
+ .toSet
+
+ // find the source operators of the component
+ val sources = componentSet.intersect(mcsPlan.getSourceOperatorIds)
+ epochMarkers += FriesComponent(sources, componentPlan.operators.map(_.id), reconfigCommands)
+ })
+ epochMarkers.toList
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/FutureBijection.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/FutureBijection.scala
new file mode 100644
index 00000000000..41f8113ef52
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/FutureBijection.scala
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import com.twitter.util.{Return, Throw, Future => TwitterFuture, Promise => TwitterPromise}
+
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.{Future => ScalaFuture, Promise => ScalaPromise}
+import scala.util.{Failure, Success}
+
+object FutureBijection {
+
+ /** Convert from a Twitter Future to a Scala Future */
+ implicit class RichTwitterFuture[A](val tf: TwitterFuture[A]) extends AnyVal {
+ def asScala: ScalaFuture[A] = {
+ val promise: ScalaPromise[A] = ScalaPromise()
+ tf.respond {
+ case Return(value) => promise.success(value)
+ case Throw(exception) => promise.failure(exception)
+ }
+ promise.future
+ }
+ }
+
+ /** Convert from a Scala Future to a Twitter Future */
+ implicit class RichScalaFuture[A](val sf: ScalaFuture[A]) extends AnyVal {
+ def asTwitter(): TwitterFuture[A] = {
+ val promise: TwitterPromise[A] = new TwitterPromise[A]()
+ sf.onComplete {
+ case Success(value) => promise.setValue(value)
+ case Failure(exception) => promise.setException(exception)
+ }
+ promise
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/SerializedState.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/SerializedState.scala
new file mode 100644
index 00000000000..c1a7bae6b88
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/SerializedState.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import org.apache.pekko.serialization.{Serialization, Serializers}
+
+object SerializedState {
+
+ val CP_STATE_KEY = "Amber_CPState"
+ val DP_STATE_KEY = "Amber_DPState"
+ val IN_FLIGHT_MSG_KEY = "Amber_Inflight_Messages"
+ val DP_QUEUED_MSG_KEY = "Amber_DP_Queued_Messages"
+ val OUTPUT_MSG_KEY = "Amber_Output_Messages"
+
+ def fromObject[T <: AnyRef](obj: T, serialization: Serialization): SerializedState = {
+ val bytes = serialization.serialize(obj).get
+ val ser = serialization.findSerializerFor(obj)
+ val manifest = Serializers.manifestFor(ser, obj)
+ SerializedState(bytes, ser.identifier, manifest)
+ }
+}
+
+case class SerializedState(bytes: Array[Byte], serializerId: Int, manifest: String) {
+
+ def toObject[T <: AnyRef](serialization: Serialization): T = {
+ serialization.deserialize(bytes, serializerId, manifest).get.asInstanceOf[T]
+ }
+
+ def size(): Long = {
+ bytes.length
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/Utils.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/Utils.scala
new file mode 100644
index 00000000000..079640317c2
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/Utils.scala
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+
+import java.nio.file.{Files, Path, Paths}
+import java.util.concurrent.locks.Lock
+import scala.annotation.tailrec
+
+object Utils extends LazyLogging {
+
+ /**
+ * Gets the real path of the amber home directory by:
+ * 1): check if the current directory is texera/amber
+ * if it's not then:
+ * 2): search the siblings and children to find the texera home path
+ *
+ * @return the real absolute path to amber home directory
+ */
+ lazy val amberHomePath: Path = {
+ val currentWorkingDirectory = Paths.get(".").toRealPath()
+ // check if the current directory is the amber home path
+ if (isAmberHomePath(currentWorkingDirectory)) {
+ currentWorkingDirectory
+ } else {
+ // from current path's directory, search its children to find amber home path
+ // current max depth is set to 2 (current path's siblings and direct children)
+ val searchChildren = Files
+ .walk(currentWorkingDirectory, 2)
+ .filter((path: Path) => isAmberHomePath(path))
+ .findAny
+ if (searchChildren.isPresent) {
+ searchChildren.get
+ } else {
+ throw new RuntimeException(
+ "Finding texera home path failed. Current working directory is " + currentWorkingDirectory
+ )
+ }
+ }
+ }
+ val AMBER_HOME_FOLDER_NAME = "amber";
+
+ /**
+ * Retry the given logic with a backoff time interval. The attempts are executed sequentially, thus blocking the thread.
+ * Backoff time is doubled after each attempt.
+ *
+ * @param attempts total number of attempts. if n <= 1 then it will not retry at all, decreased by 1 for each recursion.
+ * @param baseBackoffTimeInMS time to wait before next attempt, started with the base time, and doubled after each attempt.
+ * @param fn the target function to execute.
+ * @tparam T any return type from the provided function fn.
+ * @return the provided function fn's return, or any exception that still being raised after n attempts.
+ */
+ @tailrec
+ def retry[T](attempts: Int, baseBackoffTimeInMS: Long)(fn: => T): T = {
+ try {
+ fn
+ } catch {
+ case e: Throwable =>
+ if (attempts > 1) {
+ logger.warn(
+ "retrying after " + baseBackoffTimeInMS + "ms, number of attempts left: " + (attempts - 1),
+ e
+ )
+ Thread.sleep(baseBackoffTimeInMS)
+ retry(attempts - 1, baseBackoffTimeInMS * 2)(fn)
+ } else throw e
+ }
+ }
+
+ private def isAmberHomePath(path: Path): Boolean = {
+ path.toRealPath().endsWith(AMBER_HOME_FOLDER_NAME)
+ }
+
+ def aggregatedStateToString(state: WorkflowAggregatedState): String = {
+ state match {
+ case WorkflowAggregatedState.UNINITIALIZED => "Uninitialized"
+ case WorkflowAggregatedState.READY => "Initializing"
+ case WorkflowAggregatedState.RUNNING => "Running"
+ case WorkflowAggregatedState.PAUSING => "Pausing"
+ case WorkflowAggregatedState.PAUSED => "Paused"
+ case WorkflowAggregatedState.RESUMING => "Resuming"
+ case WorkflowAggregatedState.COMPLETED => "Completed"
+ case WorkflowAggregatedState.TERMINATED => "Terminated"
+ case WorkflowAggregatedState.FAILED => "Failed"
+ case WorkflowAggregatedState.KILLED => "Killed"
+ case WorkflowAggregatedState.UNKNOWN => "Unknown"
+ case WorkflowAggregatedState.Unrecognized(unrecognizedValue) =>
+ s"Unrecognized($unrecognizedValue)"
+ }
+ }
+
+ def stringToAggregatedState(str: String): WorkflowAggregatedState = {
+ str.trim.toLowerCase match {
+ case "uninitialized" => WorkflowAggregatedState.UNINITIALIZED
+ case "ready" => WorkflowAggregatedState.READY
+ case "initializing" => WorkflowAggregatedState.READY // accept alias
+ case "running" => WorkflowAggregatedState.RUNNING
+ case "pausing" => WorkflowAggregatedState.PAUSING
+ case "paused" => WorkflowAggregatedState.PAUSED
+ case "resuming" => WorkflowAggregatedState.RESUMING
+ case "completed" => WorkflowAggregatedState.COMPLETED
+ case "failed" => WorkflowAggregatedState.FAILED
+ case "killed" => WorkflowAggregatedState.KILLED
+ case "terminated" => WorkflowAggregatedState.TERMINATED
+ case "unknown" => WorkflowAggregatedState.UNKNOWN
+ case other => throw new IllegalArgumentException(s"Unrecognized state: $other")
+ }
+ }
+
+ /**
+ * @param state indicates the workflow state
+ * @return code indicates the status of the execution in the DB it is 0 by default for any unused states.
+ * This code is stored in the DB and read in the frontend.
+ * If these codes are changed, they also have to be changed in the frontend `ngbd-modal-workflow-executions.component.ts`
+ */
+ def maptoStatusCode(state: WorkflowAggregatedState): Byte = {
+ state match {
+ case WorkflowAggregatedState.UNINITIALIZED => 0
+ case WorkflowAggregatedState.READY => 0
+ case WorkflowAggregatedState.RUNNING => 1
+ case WorkflowAggregatedState.PAUSED => 2
+ case WorkflowAggregatedState.COMPLETED => 3
+ case WorkflowAggregatedState.FAILED => 4
+ case WorkflowAggregatedState.KILLED => 5
+ case other => -1
+ }
+ }
+
+ def withLock[X](instructions: => X)(implicit lock: Lock): X = {
+ lock.lock()
+ try {
+ instructions
+ } catch {
+ case e: Throwable =>
+ throw e
+ } finally {
+ lock.unlock()
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/DataPayload.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/DataPayload.scala
new file mode 100644
index 00000000000..54f577a0beb
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/DataPayload.scala
@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.ambermessage
+
+import org.apache.texera.amber.core.state.State
+import org.apache.texera.amber.core.tuple.Tuple
+
+sealed trait DataPayload extends WorkflowFIFOMessagePayload {}
+
+final case class StateFrame(frame: State) extends DataPayload
+
+final case class DataFrame(frame: Array[Tuple]) extends DataPayload {
+ val inMemSize: Long = {
+ frame.map(_.inMemSize).sum
+ }
+
+ override def equals(obj: Any): Boolean = {
+ if (!obj.isInstanceOf[DataFrame]) return false
+ val other = obj.asInstanceOf[DataFrame]
+ if (other eq null) return false
+ if (frame.length != other.frame.length) {
+ return false
+ }
+ var i = 0
+ while (i < frame.length) {
+ if (frame(i) != other.frame(i)) {
+ return false
+ }
+ i += 1
+ }
+ true
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/DirectControlMessagePayload.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/DirectControlMessagePayload.scala
new file mode 100644
index 00000000000..1f95cbcbc5e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/DirectControlMessagePayload.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.ambermessage
+
+trait DirectControlMessagePayload extends WorkflowFIFOMessagePayload
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/RecoveryPayload.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/RecoveryPayload.scala
new file mode 100644
index 00000000000..dcd83f3a3fc
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/RecoveryPayload.scala
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.ambermessage
+
+import org.apache.pekko.actor.{ActorRef, Address}
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+
+sealed trait RecoveryPayload extends Serializable {}
+
+// Notify controller on worker recovery starts/ends
+final case class UpdateRecoveryStatus(isRecovering: Boolean) extends RecoveryPayload
+
+// Notify upstream worker to resend output to another worker for recovery
+final case class ResendOutputTo(vid: ActorVirtualIdentity, ref: ActorRef) extends RecoveryPayload
+
+// Notify controller when the machine fails and triggers recovery
+final case class NotifyFailedNode(addr: Address) extends RecoveryPayload
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/WorkflowFIFOMessagePayload.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/WorkflowFIFOMessagePayload.scala
new file mode 100644
index 00000000000..e77ed2a2b88
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/WorkflowFIFOMessagePayload.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.ambermessage
+
+trait WorkflowFIFOMessagePayload extends Serializable
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/WorkflowMessage.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/WorkflowMessage.scala
new file mode 100644
index 00000000000..a54d4c20310
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/WorkflowMessage.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.ambermessage
+
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+
+case object WorkflowMessage {
+ def getInMemSize(msg: WorkflowMessage): Long = {
+ msg match {
+ case dataMsg: WorkflowFIFOMessage =>
+ dataMsg.payload match {
+ case df: DataFrame => df.inMemSize
+ case _ => 200L
+ }
+ case _ => 200L
+ }
+ }
+}
+
+sealed trait WorkflowMessage extends Serializable
+
+case class WorkflowFIFOMessage(
+ channelId: ChannelIdentity,
+ sequenceNumber: Long,
+ payload: WorkflowFIFOMessagePayload
+) extends WorkflowMessage
+
+case class WorkflowRecoveryMessage(
+ from: ActorVirtualIdentity,
+ payload: RecoveryPayload
+)
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/client/AmberClient.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/client/AmberClient.scala
new file mode 100644
index 00000000000..91d2ba06ba5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/client/AmberClient.scala
@@ -0,0 +1,150 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.client
+
+import org.apache.pekko.actor.{ActorSystem, Address, PoisonPill, Props}
+import org.apache.pekko.pattern._
+import org.apache.pekko.util.Timeout
+import com.twitter.util.{Future, Promise}
+import io.reactivex.rxjava3.core.Observable
+import io.reactivex.rxjava3.disposables.Disposable
+import io.reactivex.rxjava3.subjects.PublishSubject
+import org.apache.texera.amber.core.workflow.{PhysicalPlan, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.controller.ControllerConfig
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.ControlRequest
+import org.apache.texera.amber.engine.architecture.rpc.controllerservice.ControllerServiceFs2Grpc
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.ControlReturn
+import org.apache.texera.amber.engine.common.FutureBijection._
+import org.apache.texera.amber.engine.common.ambermessage.{
+ NotifyFailedNode,
+ WorkflowRecoveryMessage
+}
+import org.apache.texera.amber.engine.common.client.ClientActor.{
+ CommandRequest,
+ InitializeRequest,
+ ObservableRequest
+}
+import org.apache.texera.amber.engine.common.virtualidentity.util.CLIENT
+
+import java.lang.reflect.{InvocationHandler, Method, Proxy}
+import scala.collection.mutable
+import scala.concurrent.Await
+import scala.concurrent.duration.DurationInt
+import scala.reflect.ClassTag
+
+class AmberClient(
+ system: ActorSystem,
+ workflowContext: WorkflowContext,
+ physicalPlan: PhysicalPlan,
+ controllerConfig: ControllerConfig,
+ errorHandler: Throwable => Unit
+) {
+
+ private val clientActor = system.actorOf(Props(new ClientActor))
+ private implicit val timeout: Timeout = Timeout(1.minute)
+ private val registeredObservables = new mutable.HashMap[Class[_], Observable[_]]()
+ @volatile private var isActive = true
+
+ Await.result(
+ clientActor ? InitializeRequest(
+ workflowContext,
+ physicalPlan,
+ controllerConfig
+ ),
+ 10.seconds
+ )
+
+ def shutdown(): Unit = {
+ if (isActive) {
+ isActive = false
+ clientActor ! PoisonPill
+ }
+ }
+
+ val controllerInterface: ControllerServiceFs2Grpc[Future, Unit] =
+ createProxy[ControllerServiceFs2Grpc[Future, Unit]]()
+
+ private def createProxy[T]()(implicit ct: ClassTag[T]): T = {
+ val handler = new InvocationHandler {
+
+ override def invoke(proxy: Any, method: Method, args: Array[AnyRef]): AnyRef = {
+ val req = args(0).asInstanceOf[ControlRequest]
+ val p = Promise[ControlReturn]()
+ clientActor ! CommandRequest(method.getName, req, p)
+ p
+ }
+ }
+
+ Proxy
+ .newProxyInstance(
+ getClassLoader(ct.runtimeClass),
+ Array(ct.runtimeClass),
+ handler
+ )
+ .asInstanceOf[T]
+ }
+
+ private def getClassLoader(cls: Class[_]): ClassLoader = {
+ Option(cls.getClassLoader).getOrElse(ClassLoader.getSystemClassLoader)
+ }
+
+ def notifyNodeFailure(address: Address): Future[Any] = {
+ if (!isActive) {
+ Future[Any](())
+ } else {
+ (clientActor ? WorkflowRecoveryMessage(CLIENT, NotifyFailedNode(address))).asTwitter()
+ }
+ }
+
+ def registerCallback[T](callback: T => Unit)(implicit ct: ClassTag[T]): Disposable = {
+ if (!isActive) {
+ throw new RuntimeException("amber runtime environment is not active")
+ }
+ assert(
+ clientActor.path.address.hasLocalScope,
+ "get observable with a remote client actor is not supported"
+ )
+ val clazz = ct.runtimeClass
+ val observable =
+ if (registeredObservables.contains(clazz)) {
+ registeredObservables(clazz).asInstanceOf[Observable[T]]
+ } else {
+ val sub = PublishSubject.create[T]()
+ val req = ObservableRequest({
+ case x: T =>
+ sub.onNext(x)
+ })
+ Await.result(clientActor ? req, atMost = 2.seconds)
+ val ob = sub.onTerminateDetach
+ registeredObservables(clazz) = ob
+ ob
+ }
+ observable.subscribe { evt: T =>
+ {
+ try {
+ callback(evt)
+ } catch {
+ case t: Throwable => errorHandler(t)
+ }
+ }
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/client/ClientActor.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/client/ClientActor.scala
new file mode 100644
index 00000000000..9df86ec3ffe
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/client/ClientActor.scala
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.client
+
+import org.apache.pekko.actor.{Actor, ActorRef}
+import org.apache.pekko.pattern.StatusReply.Ack
+import com.twitter.util.Promise
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow.{PhysicalPlan, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.{
+ CreditRequest,
+ CreditResponse,
+ NetworkAck,
+ NetworkMessage
+}
+import org.apache.texera.amber.engine.architecture.controller.{
+ ClientEvent,
+ Controller,
+ ControllerConfig
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ ControlRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ ControlError,
+ ControlReturn,
+ ReturnInvocation
+}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataPayload,
+ DirectControlMessagePayload,
+ WorkflowFIFOMessage,
+ WorkflowRecoveryMessage
+}
+import org.apache.texera.amber.engine.common.client.ClientActor.{
+ ClosureRequest,
+ CommandRequest,
+ InitializeRequest,
+ ObservableRequest
+}
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient
+import org.apache.texera.amber.engine.common.virtualidentity.util.{CLIENT, CONTROLLER}
+import org.apache.texera.amber.error.ErrorUtils.reconstructThrowable
+
+import scala.collection.mutable
+
+// TODO: Rename or refactor it since it has mixed duties (send/receive messages, execute callbacks)
+private[client] object ClientActor {
+ case class InitializeRequest(
+ workflowContext: WorkflowContext,
+ physicalPlan: PhysicalPlan,
+ controllerConfig: ControllerConfig
+ )
+
+ case class ObservableRequest(pf: PartialFunction[Any, Unit])
+
+ case class ClosureRequest[T](closure: () => T)
+
+ case class CommandRequest(
+ methodName: String,
+ command: ControlRequest,
+ promise: Promise[ControlReturn]
+ )
+}
+
+private[client] class ClientActor extends Actor with AmberLogging {
+ var actorId: ActorVirtualIdentity = ActorVirtualIdentity("Client")
+ var controller: ActorRef = _
+ var controlId = 0L
+ val promiseMap = new mutable.LongMap[Promise[ControlReturn]]()
+ var handlers: PartialFunction[Any, Unit] = PartialFunction.empty
+
+ private def getQueuedCredit(channelId: ChannelIdentity): Long = {
+ 0L // client does not have queued credits
+ }
+
+ private def handleClientEvent(evt: ClientEvent): Unit = {
+ if (handlers.isDefinedAt(evt)) {
+ handlers(evt)
+ }
+ }
+
+ override def receive: Receive = {
+ case InitializeRequest(workflowContext, physicalPlan, controllerConfig) =>
+ assert(controller == null)
+ controller = context.actorOf(
+ Controller.props(workflowContext, physicalPlan, controllerConfig)
+ )
+ sender() ! Ack
+ case CreditRequest(channelId: ChannelIdentity) =>
+ sender() ! CreditResponse(channelId, getQueuedCredit(channelId))
+ case ClosureRequest(closure) =>
+ try {
+ sender() ! closure()
+ } catch {
+ case e: Throwable =>
+ sender() ! e
+ }
+ case commandRequest: CommandRequest =>
+ controller ! AsyncRPCClient.ControlInvocation(
+ commandRequest.methodName,
+ commandRequest.command,
+ AsyncRPCContext(CLIENT, CONTROLLER),
+ controlId
+ )
+ promiseMap(controlId) = commandRequest.promise
+ controlId += 1
+ case req: ObservableRequest =>
+ handlers = req.pf orElse handlers
+ sender() ! scala.runtime.BoxedUnit.UNIT
+ case NetworkMessage(
+ mId,
+ fifoMsg @ WorkflowFIFOMessage(_, _, payload)
+ ) =>
+ sender() ! NetworkAck(mId, getInMemSize(fifoMsg), getQueuedCredit(fifoMsg.channelId))
+ payload match {
+ case payload: DirectControlMessagePayload =>
+ payload match {
+ case ReturnInvocation(originalCommandID, controlReturn) =>
+ if (promiseMap.contains(originalCommandID)) {
+ controlReturn match {
+ case t: ControlError =>
+ promiseMap(originalCommandID).setException(reconstructThrowable(t))
+ case other =>
+ promiseMap(originalCommandID).setValue(other)
+ }
+ promiseMap.remove(originalCommandID)
+ }
+ case o => logger.warn(s"Amber Client should not receive control invocation: $o")
+ }
+ case _: DataPayload => ???
+ case event: ClientEvent => handleClientEvent(event)
+ case msg => logger.info(s"Amber Client received: $msg")
+ }
+ case x: WorkflowRecoveryMessage =>
+ sender() ! Ack
+ controller ! x
+ case other =>
+ logger.warn("client actor cannot handle " + other) //skip
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCClient.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCClient.scala
new file mode 100644
index 00000000000..f7e26803b47
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCClient.scala
@@ -0,0 +1,218 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.rpc
+
+import com.twitter.util.{Future, Promise}
+import io.grpc.MethodDescriptor
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.architecture.controller.ClientEvent
+import org.apache.texera.amber.engine.architecture.messaginglayer.{
+ NetworkInputGateway,
+ NetworkOutputGateway
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controllerservice.ControllerServiceFs2Grpc
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ ControlError,
+ ControlReturn,
+ ReturnInvocation,
+ WorkerMetricsResponse
+}
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceFs2Grpc
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.createProxy
+import org.apache.texera.amber.engine.common.virtualidentity.util.CLIENT
+import org.apache.texera.amber.error.ErrorUtils.reconstructThrowable
+
+import java.lang.reflect.{InvocationHandler, Method, Proxy}
+import scala.collection.mutable
+import scala.reflect.ClassTag
+
+/** Motivation of having a separate module to handle control messages as RPCs:
+ * In the old design, every control message and its response are handled by
+ * message passing. That means developers need to manually send response back
+ * and write proper handlers on the sender side.
+ * Writing control messages becomes tedious if we use this way.
+ *
+ * So we want to implement rpc model on top of message passing.
+ * rpc (request-response)
+ * remote.callFunctionX().then(response => {
+ * })
+ * user-api: promise
+ *
+ * goal: request-response model with multiplexing
+ * client: initiate request
+ * (web browser, actor that invoke control command)
+ * server: handle request, return response
+ * (web server, actor that handles control command)
+ */
+object AsyncRPCClient {
+
+ final val IgnoreReply = -1
+ final val IgnoreReplyAndDoNotLog = -2
+
+ object ControlInvocation {
+ def apply(
+ method: MethodDescriptor[_, _],
+ payload: ControlRequest,
+ context: AsyncRPCContext,
+ commandID: Long
+ ): ControlInvocation = {
+ new ControlInvocation(method.getBareMethodName, payload, context, commandID)
+ }
+
+ def apply(
+ methodName: String,
+ payload: ControlRequest,
+ context: AsyncRPCContext,
+ commandID: Long
+ ): ControlInvocation = {
+ new ControlInvocation(methodName, payload, context, commandID)
+ }
+ }
+
+ /**
+ * Creates a dynamic proxy for the specified type `T`, which intercepts method calls
+ * and sends them as ControlInvocation messages via the provided output gateway.
+ */
+ def createProxy[T](
+ createPromise: () => (Promise[ControlReturn], Long),
+ outputGateway: NetworkOutputGateway
+ )(implicit ct: ClassTag[T]): T = {
+ val handler = new InvocationHandler {
+
+ override def invoke(proxy: Any, method: Method, args: Array[AnyRef]): AnyRef = {
+ val (p, pid) = createPromise()
+ val context = args(1).asInstanceOf[AsyncRPCContext]
+ val msg = args(0).asInstanceOf[ControlRequest]
+ outputGateway.sendTo(context.receiver, ControlInvocation(method.getName, msg, context, pid))
+ p
+ }
+ }
+
+ Proxy
+ .newProxyInstance(
+ getClassLoader(ct.runtimeClass),
+ Array(ct.runtimeClass),
+ handler
+ )
+ .asInstanceOf[T]
+ }
+
+ // Helper to get the correct class loader
+ private def getClassLoader(cls: Class[_]): ClassLoader = {
+ Option(cls.getClassLoader).getOrElse(ClassLoader.getSystemClassLoader)
+ }
+
+}
+
+class AsyncRPCClient(
+ val inputGateway: NetworkInputGateway,
+ val outputGateway: NetworkOutputGateway,
+ val actorId: ActorVirtualIdentity
+) extends AmberLogging {
+
+ private val unfulfilledPromises = mutable.HashMap[Long, Promise[ControlReturn]]()
+ private var promiseID = 0L
+ @transient lazy val controllerInterface: ControllerServiceFs2Grpc[Future, AsyncRPCContext] =
+ createProxy[ControllerServiceFs2Grpc[Future, AsyncRPCContext]](createPromise, outputGateway)
+ @transient lazy val workerInterface: WorkerServiceFs2Grpc[Future, AsyncRPCContext] =
+ createProxy[WorkerServiceFs2Grpc[Future, AsyncRPCContext]](createPromise, outputGateway)
+
+ def mkContext(to: ActorVirtualIdentity): AsyncRPCContext = AsyncRPCContext(actorId, to)
+
+ protected def createPromise(): (Promise[ControlReturn], Long) = {
+ promiseID += 1
+ val promise = new Promise[ControlReturn]()
+ unfulfilledPromises(promiseID) = promise
+ (promise, promiseID)
+ }
+
+ def createInvocation(
+ methodName: String,
+ message: ControlRequest,
+ context: AsyncRPCContext
+ ): (ControlInvocation, Future[ControlReturn]) = {
+ val (p, pid) = createPromise()
+ (ControlInvocation(methodName, message, context, pid), p)
+ }
+
+ def sendECMToChannel(
+ ecmId: EmbeddedControlMessageIdentity,
+ ecmType: EmbeddedControlMessageType,
+ scope: Set[ChannelIdentity],
+ cmdMapping: Map[String, ControlInvocation],
+ channelId: ChannelIdentity
+ ): Unit = {
+ logger.debug(s"send ECM: $ecmId to $channelId")
+ outputGateway.sendTo(
+ channelId,
+ EmbeddedControlMessage(ecmId, ecmType, scope.toSeq, cmdMapping)
+ )
+ }
+
+ def sendToClient(clientEvent: ClientEvent): Unit = {
+ outputGateway.sendTo(
+ ChannelIdentity(actorId, CLIENT, isControl = true),
+ clientEvent
+ )
+ }
+
+ def fulfillPromise(ret: ReturnInvocation): Unit = {
+ if (unfulfilledPromises.contains(ret.commandId)) {
+ val p = unfulfilledPromises(ret.commandId)
+ ret.returnValue match {
+ case err: ControlError =>
+ p.setException(reconstructThrowable(err))
+ case other =>
+ p.setValue(other)
+ }
+ unfulfilledPromises.remove(ret.commandId)
+ }
+ }
+
+ def logControlReply(ret: ReturnInvocation, channelId: ChannelIdentity): Unit = {
+ if (ret.commandId == AsyncRPCClient.IgnoreReplyAndDoNotLog) {
+ return
+ }
+ if (ret.returnValue != null) {
+ if (ret.returnValue.isInstanceOf[WorkerMetricsResponse]) {
+ return
+ }
+ logger.debug(
+ s"receive reply: ${ret.returnValue.getClass.getSimpleName} from $channelId (controlID: ${ret.commandId})"
+ )
+ ret.returnValue match {
+ case err: ControlError =>
+ logger.error(s"received error from $channelId", err)
+ case _ =>
+ }
+ } else {
+ logger.info(
+ s"receive reply: null from $channelId (controlID: ${ret.commandId})"
+ )
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCHandlerInitializer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCHandlerInitializer.scala
new file mode 100644
index 00000000000..f658e0875f7
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCHandlerInitializer.scala
@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.rpc
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import org.apache.texera.amber.engine.architecture.controller.ClientEvent
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controllerservice.ControllerServiceFs2Grpc
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceFs2Grpc
+
+import scala.language.implicitConversions
+
+class AsyncRPCHandlerInitializer(
+ ctrlSource: AsyncRPCClient,
+ ctrlReceiver: AsyncRPCServer
+) {
+ implicit def returnAsFuture[R](ret: R): Future[R] = Future[R](ret)
+
+ implicit def actorIdAsContext(to: ActorVirtualIdentity): AsyncRPCContext = mkContext(to)
+
+ implicit def stringToResponse(s: String): StringResponse = StringResponse(s)
+
+ implicit def intToResponse(i: Int): IntResponse = IntResponse(i)
+
+ // register all handlers
+ ctrlReceiver.handler = this
+
+ def controllerInterface: ControllerServiceFs2Grpc[Future, AsyncRPCContext] =
+ ctrlSource.controllerInterface
+
+ def workerInterface: WorkerServiceFs2Grpc[Future, AsyncRPCContext] = ctrlSource.workerInterface
+
+ def mkContext(to: ActorVirtualIdentity): AsyncRPCContext = ctrlSource.mkContext(to)
+
+ def sendECM(
+ ecmId: EmbeddedControlMessageIdentity,
+ ecmType: EmbeddedControlMessageType,
+ scope: Set[ChannelIdentity],
+ cmdMapping: Map[String, ControlInvocation],
+ to: ChannelIdentity
+ ): Unit = {
+ ctrlSource.sendECMToChannel(ecmId, ecmType, scope, cmdMapping, to)
+ }
+
+ def sendToClient(clientEvent: ClientEvent): Unit = {
+ ctrlSource.sendToClient(clientEvent)
+ }
+
+ def createInvocation(
+ methodName: String,
+ payload: ControlRequest,
+ to: ActorVirtualIdentity
+ ): (ControlInvocation, Future[ControlReturn]) =
+ ctrlSource.createInvocation(methodName, payload, ctrlSource.mkContext(to))
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCServer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCServer.scala
new file mode 100644
index 00000000000..e9a3e2cc455
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/rpc/AsyncRPCServer.scala
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.rpc
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.messaginglayer.NetworkOutputGateway
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ ControlInvocation,
+ ControlRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ ControlReturn,
+ ReturnInvocation
+}
+import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.error.ErrorUtils.mkControlError
+
+import java.lang.reflect.Method
+import scala.collection.mutable
+
+class AsyncRPCServer(
+ outputGateway: NetworkOutputGateway,
+ val actorId: ActorVirtualIdentity
+) extends AmberLogging {
+
+ // variable to hold all handler implementations
+ var handler: AnyRef = _
+
+ // retrieve a mapping from method name to implementation
+ // used transient lazy val to avoid serialization.
+ @transient
+ private lazy val methodsByName: Map[String, Method] = {
+ val mapping = mutable.HashMap[String, Method]()
+ handler.getClass.getMethods.foreach { method =>
+ mapping(method.getName.toLowerCase) = method
+ }
+ mapping.toMap
+ }
+
+ def receive(request: ControlInvocation, senderID: ActorVirtualIdentity): Unit = {
+ val methodName = request.methodName.toLowerCase
+ val requestArg = request.command
+ val contextArg = request.context
+ val id = request.commandId
+ logger.debug(
+ s"receive command: ${methodName} with payload ${requestArg} from $senderID (controlID: ${id})"
+ )
+ methodsByName.get(methodName) match {
+ case Some(method) =>
+ invokeMethod(method, requestArg, contextArg, id, senderID)
+ case None =>
+ logger.error(s"No methods found with name $methodName")
+ }
+ }
+
+ private def invokeMethod(
+ method: Method,
+ requestArg: ControlRequest,
+ contextArg: AsyncRPCContext,
+ id: Long,
+ senderID: ActorVirtualIdentity
+ ): Unit = {
+ try {
+ val result =
+ try {
+ method.invoke(handler, requestArg, contextArg)
+ } catch {
+ case e: java.lang.reflect.InvocationTargetException =>
+ throw Option(e.getCause).getOrElse(e)
+ case e: Throwable => throw e
+ }
+ result
+ .asInstanceOf[Future[ControlReturn]]
+ .onSuccess { ret =>
+ returnResult(senderID, id, ret)
+ }
+ .onFailure { err =>
+ logger.error("Exception occurred", err)
+ returnResult(senderID, id, mkControlError(err))
+ }
+
+ } catch {
+ case err: Throwable =>
+ // if error occurs, return it to the sender.
+ logger.error("Exception occurred", err)
+ returnResult(senderID, id, mkControlError(err))
+ // if throw this exception right now, the above message might not be able
+ // to be sent out. We do not throw for now.
+ // throw err
+ }
+ }
+
+ @inline
+ private def noReplyNeeded(id: Long): Boolean = id < 0
+
+ @inline
+ private def returnResult(sender: ActorVirtualIdentity, id: Long, ret: ControlReturn): Unit = {
+ if (noReplyNeeded(id)) {
+ return
+ }
+ outputGateway.sendTo(sender, ReturnInvocation(id, ret))
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/statetransition/StateManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/statetransition/StateManager.scala
new file mode 100644
index 00000000000..9bd951389ea
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/statetransition/StateManager.scala
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.statetransition
+
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.common.statetransition.StateManager.{
+ InvalidStateException,
+ InvalidTransitionException
+}
+
+object StateManager {
+
+ case class InvalidStateException(msg: String, actorId: ActorVirtualIdentity)
+ extends WorkflowRuntimeException(msg, Some(actorId))
+
+ case class InvalidTransitionException(msg: String, actorId: ActorVirtualIdentity)
+ extends WorkflowRuntimeException(msg, Some(actorId))
+}
+
+class StateManager[T](
+ actorId: ActorVirtualIdentity,
+ stateTransitionGraph: Map[T, Set[T]],
+ initialState: T
+) extends Serializable {
+
+ private var currentState: T = initialState
+
+ def assertState(state: T): Unit = {
+ if (currentState != state) {
+ throw InvalidStateException(
+ s"except state = $state but current state = $currentState",
+ actorId
+ )
+ }
+ }
+
+ def assertState(states: T*): Unit = {
+ if (!states.contains(currentState)) {
+ throw InvalidStateException(
+ s"except state in [${states.mkString(",")}] but current state = $currentState",
+ actorId
+ )
+ }
+ }
+
+ def conditionalTransitTo(currentState: T, targetState: T, callback: () => Unit): Unit = {
+ if (getCurrentState == currentState) {
+ transitTo(targetState)
+ callback()
+ }
+ }
+
+ def confirmState(state: T): Boolean = getCurrentState == state
+
+ def getCurrentState: T = currentState
+
+ def confirmState(states: T*): Boolean = states.contains(getCurrentState)
+
+ def transitTo(state: T): Unit = {
+ if (state == currentState) {
+ return
+ // throw InvalidTransitionException(s"current state is already $currentState")
+ }
+
+ if (!stateTransitionGraph.getOrElse(currentState, Set()).contains(state)) {
+ throw InvalidTransitionException(s"cannot transit from $currentState to $state", actorId)
+ }
+ currentState = state
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/statetransition/WorkerStateManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/statetransition/WorkerStateManager.scala
new file mode 100644
index 00000000000..4f5e92e0ae7
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/statetransition/WorkerStateManager.scala
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.statetransition
+
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState._
+
+// The following pattern is a good practice of enum in scala
+// We've always used this pattern in the codebase
+// https://nrinaudo.github.io/scala-best-practices/definitions/adt.html
+// https://nrinaudo.github.io/scala-best-practices/adts/product_with_serializable.html
+
+class WorkerStateManager(actorId: ActorVirtualIdentity, initialState: WorkerState = UNINITIALIZED)
+ extends StateManager[WorkerState](
+ actorId,
+ Map(
+ UNINITIALIZED -> Set(READY),
+ READY -> Set(PAUSED, RUNNING, COMPLETED),
+ RUNNING -> Set(PAUSED, COMPLETED),
+ PAUSED -> Set(RUNNING),
+ COMPLETED -> Set()
+ ),
+ initialState
+ ) {}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/EmptyRecordStorage.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/EmptyRecordStorage.scala
new file mode 100644
index 00000000000..0dbbd61b519
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/EmptyRecordStorage.scala
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.storage
+
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage.{
+ SequentialRecordReader,
+ SequentialRecordWriter
+}
+import org.apache.commons.io.input.NullInputStream
+import org.apache.hadoop.io.IOUtils.NullOutputStream
+
+import java.io.{DataInputStream, DataOutputStream}
+import scala.reflect.ClassTag
+
+class EmptyRecordStorage[T >: Null <: AnyRef: ClassTag] extends SequentialRecordStorage[T] {
+ override def getWriter(fileName: String): SequentialRecordWriter[T] = {
+ new SequentialRecordWriter(
+ new DataOutputStream(
+ new NullOutputStream()
+ )
+ )
+ }
+
+ override def getReader(fileName: String): SequentialRecordReader[T] = {
+ new SequentialRecordReader(() =>
+ new DataInputStream(
+ new NullInputStream()
+ )
+ )
+ }
+
+ override def deleteStorage(): Unit = {
+ // empty
+ }
+
+ override def containsFolder(folderName: String): Boolean = false
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/HDFSRecordStorage.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/HDFSRecordStorage.scala
new file mode 100644
index 00000000000..07d6bab0373
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/HDFSRecordStorage.scala
@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.storage
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage.{
+ SequentialRecordReader,
+ SequentialRecordWriter
+}
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileSystem, Path}
+
+import java.net.URI
+import scala.reflect.ClassTag
+
+class HDFSRecordStorage[T >: Null <: AnyRef: ClassTag](hdfsLogFolderURI: URI)
+ extends SequentialRecordStorage[T]
+ with LazyLogging {
+
+ // only support hdfs uris
+ assert(hdfsLogFolderURI.getScheme.toLowerCase == "hdfs")
+
+ private var fileSystem: FileSystem = _
+ private val fsConf = new Configuration()
+ // configuration for HDFS
+ fsConf.set("dfs.client.block.write.replace-datanode-on-failure.enable", "false")
+ fileSystem = FileSystem.get(hdfsLogFolderURI, fsConf)
+ fileSystem.setWriteChecksum(false)
+ fileSystem.setVerifyChecksum(false)
+
+ private val folderPath =
+ Path.mergePaths(fileSystem.getWorkingDirectory, new Path(hdfsLogFolderURI.getPath))
+
+ if (!fileSystem.exists(folderPath)) {
+ fileSystem.mkdirs(folderPath)
+ }
+
+ override def getWriter(fileName: String): SequentialRecordWriter[T] = {
+ new SequentialRecordWriter(fileSystem.create(folderPath.suffix("/" + fileName)))
+ }
+
+ override def getReader(fileName: String): SequentialRecordReader[T] = {
+ val path = folderPath.suffix("/" + fileName)
+ if (fileSystem.exists(path)) {
+ new SequentialRecordReader(() => fileSystem.open(path))
+ } else {
+ new EmptyRecordStorage[T]().getReader(fileName)
+ }
+ }
+
+ override def deleteStorage(): Unit = {
+ // delete the entire log folder if exists
+ if (fileSystem.exists(folderPath)) {
+ fileSystem.delete(folderPath, true)
+ }
+ }
+
+ override def containsFolder(folderName: String): Boolean = {
+ val path = folderPath.suffix("/" + folderName)
+ fileSystem.exists(path) && fileSystem.getFileStatus(path).isDirectory
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/SequentialRecordStorage.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/SequentialRecordStorage.scala
new file mode 100644
index 00000000000..f5664e6a49c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/SequentialRecordStorage.scala
@@ -0,0 +1,138 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.storage
+
+import com.esotericsoftware.kryo.io.{Input, Output}
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage.{
+ SequentialRecordReader,
+ SequentialRecordWriter
+}
+
+import java.io.{DataInputStream, DataOutputStream}
+import java.net.URI
+import scala.collection.mutable.ArrayBuffer
+import scala.reflect.{ClassTag, classTag}
+
+object SequentialRecordStorage {
+
+ // For debugging purpose only
+ def fetchAllRecords[T >: Null <: AnyRef](
+ storage: SequentialRecordStorage[T],
+ logFileName: String
+ ): Iterable[T] = {
+ val reader = storage.getReader(logFileName)
+ val recordIter = reader.mkRecordIterator()
+ val buffer = new ArrayBuffer[T]()
+ while (recordIter.hasNext) {
+ buffer.append(recordIter.next())
+ }
+ buffer
+ }
+
+ class SequentialRecordWriter[T >: Null <: AnyRef](outputStream: DataOutputStream) {
+ lazy val output = new Output(outputStream)
+ def writeRecord(obj: T): Unit = {
+ val bytes = AmberRuntime.serde.serialize(obj).get
+ output.writeInt(bytes.length)
+ output.write(bytes)
+ }
+ def flush(): Unit = {
+ output.flush()
+ }
+ def close(): Unit = {
+ output.close()
+ }
+ }
+
+ class SequentialRecordReader[T >: Null <: AnyRef: ClassTag](
+ inputStreamGen: () => DataInputStream
+ ) {
+ val clazz = classTag[T].runtimeClass.asInstanceOf[Class[T]]
+
+ def mkRecordIterator(): Iterator[T] = {
+ lazy val input = new Input(inputStreamGen())
+ new Iterator[T] {
+ var record: T = internalNext()
+ private def internalNext(): T = {
+ try {
+ val len = input.readInt()
+ val bytes = input.readBytes(len)
+ AmberRuntime.serde.deserialize(bytes, clazz).get
+ } catch {
+ case e: Throwable =>
+ input.close()
+ null
+ }
+ }
+ override def next(): T = {
+ val currentRecord = record
+ record = internalNext()
+ currentRecord
+ }
+ override def hasNext: Boolean = record != null
+ }
+ }
+ }
+
+ def getStorage[T >: Null <: AnyRef: ClassTag](
+ storageLocation: Option[URI]
+ ): SequentialRecordStorage[T] = {
+ storageLocation match {
+ case Some(location) =>
+ if (location.getScheme.toLowerCase == "hdfs") {
+ new HDFSRecordStorage(location) // hdfs lib supports r/w operations
+ } else {
+ new VFSRecordStorage(location)
+ }
+ case None => new EmptyRecordStorage()
+ }
+ }
+}
+
+/**
+ * Sequential record storage is designed to do read/write for sequential generic data. It represents
+ * a one-level folder (no nesting) which contains a list of files. Files are identified by a unique
+ * file name string.
+ *
+ * Key Features:
+ * - Allows for the sequential writing and reading of records of a generic type `T`.
+ * It utilizes Kryo serialization for efficient binary storage of records.
+ * - The class assumes a sequential access pattern to the data. It is not optimized for random
+ * access or querying specific records without reading sequentially.
+ * Usage:
+ * - To use `SequentialRecordStorage`, one must extend this abstract class and implement the
+ * methods for creating record readers and writers. Implementations can customize how and
+ * where the data is stored and retrieved.
+ * - The `SequentialRecordWriter` and `SequentialRecordReader` inner classes provide the
+ * functionality for writing to and reading from the storage.
+ *
+ * @tparam T The type of records that this storage system will handle.
+ */
+abstract class SequentialRecordStorage[T >: Null <: AnyRef] {
+
+ def getWriter(fileName: String): SequentialRecordWriter[T]
+
+ def getReader(fileName: String): SequentialRecordReader[T]
+
+ def deleteStorage(): Unit
+
+ def containsFolder(folderName: String): Boolean
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/VFSRecordStorage.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/VFSRecordStorage.scala
new file mode 100644
index 00000000000..a514360d6d6
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/storage/VFSRecordStorage.scala
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.storage
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage.{
+ SequentialRecordReader,
+ SequentialRecordWriter
+}
+import org.apache.commons.vfs2.{FileObject, FileSystemManager, VFS}
+
+import java.io.{DataInputStream, DataOutputStream}
+import java.net.URI
+import scala.reflect.ClassTag
+
+class VFSRecordStorage[T >: Null <: AnyRef: ClassTag](vfsLogFolderURI: URI)
+ extends SequentialRecordStorage[T]
+ with LazyLogging {
+
+ private val fs: FileSystemManager = VFS.getManager
+ private val folder: FileObject = fs.resolveFile(vfsLogFolderURI)
+
+ if (!folder.exists()) {
+ folder.createFolder()
+ }
+
+ override def getWriter(fileName: String): SequentialRecordStorage.SequentialRecordWriter[T] = {
+ val file = folder.resolveFile(fileName)
+ file.createFile()
+ val outputStream = file.getContent.getOutputStream
+ new SequentialRecordWriter(new DataOutputStream(outputStream))
+ }
+
+ override def getReader(fileName: String): SequentialRecordStorage.SequentialRecordReader[T] = {
+ new SequentialRecordReader(() => {
+ val inputStream = folder.resolveFile(fileName).getContent.getInputStream
+ new DataInputStream(inputStream)
+ })
+ }
+
+ override def deleteStorage(): Unit = {
+ folder.deleteAll()
+ }
+
+ override def containsFolder(folderName: String): Boolean = {
+ val fileObj = folder.getChild(folderName)
+ fileObj != null && fileObj.exists() && fileObj.isFolder
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/virtualidentity/util.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/virtualidentity/util.scala
new file mode 100644
index 00000000000..481aecab1dc
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/virtualidentity/util.scala
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.common.virtualidentity
+
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+
+object util {
+
+ lazy val CONTROLLER: ActorVirtualIdentity = ActorVirtualIdentity("CONTROLLER")
+ lazy val SELF: ActorVirtualIdentity = ActorVirtualIdentity("SELF")
+ lazy val CLIENT: ActorVirtualIdentity = ActorVirtualIdentity("CLIENT")
+}
diff --git a/amber/src/main/scala/org/apache/texera/amber/error/ErrorUtils.scala b/amber/src/main/scala/org/apache/texera/amber/error/ErrorUtils.scala
new file mode 100644
index 00000000000..af16044af29
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/amber/error/ErrorUtils.scala
@@ -0,0 +1,119 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.error
+
+import com.google.protobuf.timestamp.Timestamp
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.ConsoleMessage
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.ConsoleMessageType.ERROR
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ControlError, ErrorLanguage}
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+import java.time.Instant
+import scala.util.control.ControlThrowable
+
+object ErrorUtils {
+
+ /** A helper function for catching all throwable except some special scala internal throwable.
+ * reference: https://www.sumologic.com/blog/why-you-should-never-catch-throwable-in-scala/
+ *
+ * @param handler
+ * @tparam T
+ * @return
+ */
+ def safely[T](handler: PartialFunction[Throwable, T]): PartialFunction[Throwable, T] = {
+ case ex: ControlThrowable => throw ex
+ // case ex: OutOfMemoryError (Assorted other nasty exceptions you don't want to catch)
+ // If it's an exception they handle, pass it on
+ case ex: Throwable if handler.isDefinedAt(ex) => handler(ex)
+ // If they didn't handle it, rethrow automatically
+ }
+
+ def mkConsoleMessage(actorId: ActorVirtualIdentity, err: Throwable): ConsoleMessage = {
+ val source = if (err.getStackTrace.nonEmpty) {
+ "(" + err.getStackTrace.head.getFileName + ":" + err.getStackTrace.head.getLineNumber + ")"
+ } else {
+ "(Unknown Source)"
+ }
+ val title = err.toString
+ val message = err.getStackTrace.mkString("\n")
+ ConsoleMessage(actorId.name, Timestamp(Instant.now), ERROR, source, title, message)
+ }
+
+ def mkControlError(err: Throwable): ControlError = {
+ // Format each stack trace element with "at " prefix
+ val stacktrace = err.getStackTrace.map(element => s"at ${element}").mkString("\n")
+ if (err.getCause != null) {
+ ControlError(err.toString, err.getCause.toString, stacktrace, ErrorLanguage.SCALA)
+ } else {
+ ControlError(err.toString, "", stacktrace, ErrorLanguage.SCALA)
+ }
+ }
+
+ def reconstructThrowable(controlError: ControlError): Throwable = {
+ if (controlError.language == ErrorLanguage.PYTHON) {
+ return new Throwable(controlError.errorMessage)
+ } else {
+ val reconstructedThrowable = new Throwable(controlError.errorMessage)
+ if (controlError.errorDetails.nonEmpty) {
+ val causeThrowable = new Throwable(controlError.errorDetails)
+ reconstructedThrowable.initCause(causeThrowable)
+ }
+
+ val stackTracePattern = """\s*at\s+(.+)\((.*)\)""".r
+ val stackTraceElements = controlError.stackTrace.split("\n").flatMap { line =>
+ line match {
+ case stackTracePattern(className, location) =>
+ Some(new StackTraceElement(className, "", location, -1))
+ case _ => None
+ }
+ }
+ reconstructedThrowable.setStackTrace(stackTraceElements)
+ reconstructedThrowable
+ }
+ }
+
+ def getStackTraceWithAllCauses(err: Throwable, topLevel: Boolean = true): String = {
+ val header = if (topLevel) {
+ "Stack trace for developers: \n\n"
+ } else {
+ "\n\nCaused by:\n"
+ }
+ val message = header + err.toString + "\n" + err.getStackTrace.mkString("\n")
+ if (err.getCause != null) {
+ message + getStackTraceWithAllCauses(err.getCause, topLevel = false)
+ } else {
+ message
+ }
+ }
+
+ def getOperatorFromActorIdOpt(
+ actorIdOpt: Option[ActorVirtualIdentity]
+ ): (String, String) = {
+ var operatorId = "unknown operator"
+ var workerId = ""
+ if (actorIdOpt.isDefined) {
+ operatorId = VirtualIdentityUtils.getPhysicalOpId(actorIdOpt.get).logicalOpId.id
+ workerId = actorIdOpt.get.name
+ }
+ (operatorId, workerId)
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/ComputingUnitMaster.scala b/amber/src/main/scala/org/apache/texera/web/ComputingUnitMaster.scala
new file mode 100644
index 00000000000..41d8d3b5830
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/ComputingUnitMaster.scala
@@ -0,0 +1,307 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web
+
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import com.typesafe.scalalogging.LazyLogging
+import io.dropwizard.Configuration
+import io.dropwizard.configuration.{EnvironmentVariableSubstitutor, SubstitutingSourceProvider}
+import io.dropwizard.setup.{Bootstrap, Environment}
+import io.dropwizard.websockets.WebsocketBundle
+import org.apache.texera.amber.config.{ApplicationConfig, StorageConfig}
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.virtualidentity.ExecutionIdentity
+import org.apache.texera.amber.core.workflow.{PhysicalPlan, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.controller.ControllerConfig
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{
+ COMPLETED,
+ FAILED
+}
+import org.apache.texera.amber.engine.common.AmberRuntime.scheduleRecurringCallThroughActorSystem
+import org.apache.texera.amber.engine.common.Utils.maptoStatusCode
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.engine.common.{AmberRuntime, Utils}
+import org.apache.texera.amber.util.JSONUtils.objectMapper
+import org.apache.texera.amber.util.ObjectMapperUtils
+import org.apache.commons.jcs3.access.exception.InvalidArgumentException
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.tables.pojos.WorkflowExecutions
+import org.apache.texera.web.auth.JwtAuth.setupJwtAuth
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+import org.apache.texera.web.resource.{
+ SyncExecutionResource,
+ WebsocketPayloadSizeTuner,
+ WorkflowWebsocketResource
+}
+import org.apache.texera.web.service.ExecutionsMetadataPersistService
+import org.eclipse.jetty.server.session.SessionHandler
+import org.eclipse.jetty.servlet.FilterHolder
+import org.eclipse.jetty.websocket.server.WebSocketUpgradeFilter
+import org.apache.texera.web.resource.pythonvirtualenvironment.PveResource
+import org.apache.texera.web.resource.pythonvirtualenvironment.PveWebsocketResource
+
+import java.net.URI
+import java.time.Duration
+import scala.annotation.tailrec
+import scala.concurrent.duration.DurationInt
+
+object ComputingUnitMaster {
+
+ def createAmberRuntime(
+ workflowContext: WorkflowContext,
+ physicalPlan: PhysicalPlan,
+ conf: ControllerConfig,
+ errorHandler: Throwable => Unit
+ ): AmberClient = {
+ new AmberClient(
+ AmberRuntime.actorSystem,
+ workflowContext,
+ physicalPlan,
+ conf,
+ errorHandler
+ )
+ }
+
+ type OptionMap = Map[Symbol, Any]
+
+ def parseArgs(args: Array[String]): OptionMap = {
+ @tailrec
+ def nextOption(map: OptionMap, list: List[String]): OptionMap = {
+ list match {
+ case Nil => map
+ case "--cluster" :: value :: tail =>
+ nextOption(map ++ Map(Symbol("cluster") -> value.toBoolean), tail)
+ case option :: tail =>
+ throw new InvalidArgumentException("unknown command-line arg")
+ }
+ }
+
+ nextOption(Map(), args.toList)
+ }
+
+ def main(args: Array[String]): Unit = {
+ val argMap = parseArgs(args)
+
+ val clusterMode = argMap.get(Symbol("cluster")).asInstanceOf[Option[Boolean]].getOrElse(false)
+ // start actor system master node
+ AmberRuntime.startActorMaster(clusterMode)
+ // start web server
+ new ComputingUnitMaster().run(
+ "server",
+ Utils.amberHomePath
+ .resolve("src")
+ .resolve("main")
+ .resolve("resources")
+ .resolve("computing-unit-master-config.yml")
+ .toString
+ )
+ }
+}
+
+class ComputingUnitMaster extends io.dropwizard.Application[Configuration] with LazyLogging {
+
+ override def initialize(bootstrap: Bootstrap[Configuration]): Unit = {
+ // enable environment variable substitution in YAML config
+ bootstrap.setConfigurationSourceProvider(
+ new SubstitutingSourceProvider(
+ bootstrap.getConfigurationSourceProvider,
+ new EnvironmentVariableSubstitutor(false)
+ )
+ )
+ // add websocket bundle
+ bootstrap.addBundle(
+ new WebsocketBundle(
+ classOf[WorkflowWebsocketResource],
+ classOf[PveWebsocketResource]
+ )
+ )
+ // register scala module to dropwizard default object mapper
+ bootstrap.getObjectMapper.registerModule(DefaultScalaModule)
+ }
+
+ override def run(configuration: Configuration, environment: Environment): Unit = {
+ ObjectMapperUtils.warmupObjectMapperForOperatorsSerde()
+
+ SqlServer.initConnection(
+ StorageConfig.jdbcUrl,
+ StorageConfig.jdbcUsername,
+ StorageConfig.jdbcPassword
+ )
+
+ environment.jersey.setUrlPattern("/api/*")
+
+ val webSocketUpgradeFilter =
+ WebSocketUpgradeFilter.configureContext(environment.getApplicationContext)
+ webSocketUpgradeFilter.getFactory.getPolicy.setIdleTimeout(Duration.ofHours(1).toMillis)
+ environment.getApplicationContext.setAttribute(
+ classOf[WebSocketUpgradeFilter].getName,
+ webSocketUpgradeFilter
+ )
+
+ // register SessionHandler
+ environment.jersey.register(classOf[SessionHandler])
+ environment.servlets.setSessionHandler(new SessionHandler)
+
+ environment.jersey.register(classOf[PveResource])
+
+ setupJwtAuth(environment)
+
+ environment.jersey.register(
+ new io.dropwizard.auth.AuthValueFactoryProvider.Binder[SessionUser](classOf[SessionUser])
+ )
+ environment.jersey.register(
+ classOf[org.glassfish.jersey.server.filter.RolesAllowedDynamicFeature]
+ )
+ environment
+ .servlets()
+ .addServletListeners(
+ new WebsocketPayloadSizeTuner(ApplicationConfig.maxWorkflowWebsocketRequestPayloadSizeKb)
+ )
+
+ val timeToLive: Int = ApplicationConfig.sinkStorageTTLInSecs
+ if (ApplicationConfig.cleanupAllExecutionResults) {
+ // do one time cleanup of collections that were not closed gracefully before restart/crash
+ // retrieve all executions that were executing before the reboot.
+ val allExecutionsBeforeRestart: List[WorkflowExecutions] =
+ WorkflowExecutionsResource.getExpiredExecutionsWithResultOrLog(-1)
+ cleanExecutions(
+ allExecutionsBeforeRestart,
+ statusByte => {
+ if (statusByte != maptoStatusCode(COMPLETED)) {
+ maptoStatusCode(FAILED) // for incomplete executions, mark them as failed.
+ } else {
+ statusByte
+ }
+ }
+ )
+ }
+ scheduleRecurringCallThroughActorSystem(
+ 2.seconds,
+ ApplicationConfig.sinkStorageCleanUpCheckIntervalInSecs.seconds
+ ) {
+ recurringCheckExpiredResults(timeToLive)
+ }
+
+ environment.jersey.register(classOf[WorkflowExecutionsResource])
+ environment.jersey.register(classOf[SyncExecutionResource])
+
+ // Route request logs through SLF4J, controlled by TEXERA_SERVICE_LOG_LEVEL.
+ // TODO: replace with RequestLoggingFilter.register() from common/auth once Dropwizard is upgraded to 4.x
+ val requestLogger = org.slf4j.LoggerFactory.getLogger("org.eclipse.jetty.server.RequestLog")
+ environment.getApplicationContext.addFilter(
+ new FilterHolder(new javax.servlet.Filter {
+ override def init(filterConfig: javax.servlet.FilterConfig): Unit = {}
+ override def doFilter(
+ request: javax.servlet.ServletRequest,
+ response: javax.servlet.ServletResponse,
+ chain: javax.servlet.FilterChain
+ ): Unit = {
+ chain.doFilter(request, response)
+ if (requestLogger.isInfoEnabled) {
+ val req = request.asInstanceOf[javax.servlet.http.HttpServletRequest]
+ val resp = response.asInstanceOf[javax.servlet.http.HttpServletResponse]
+ requestLogger.info(
+ s"""${req.getRemoteAddr} - "${req.getMethod} ${req.getRequestURI} ${req.getProtocol}" ${resp.getStatus}"""
+ )
+ }
+ }
+ override def destroy(): Unit = {}
+ }),
+ "/*",
+ java.util.EnumSet.allOf(classOf[javax.servlet.DispatcherType])
+ )
+ }
+
+ /**
+ * This function drops the collections.
+ * MongoDB doesn't have an API of drop collection where collection name in (from a subquery), so the implementation is to retrieve
+ * the entire list of those documents that have expired, then loop the list to drop them one by one
+ */
+ private def cleanExecutions(
+ executions: List[WorkflowExecutions],
+ statusChangeFunc: Short => Short
+ ): Unit = {
+ // drop the collection and update the status to ABORTED
+ executions.foreach(execEntry => {
+ dropCollections(execEntry.getResult)
+ deleteReplayLog(execEntry.getLogLocation)
+ // then delete the pointer from mySQL
+ val executionIdentity = ExecutionIdentity(execEntry.getEid.longValue())
+ ExecutionsMetadataPersistService.tryUpdateExistingExecution(executionIdentity) { execution =>
+ execution.setResult("")
+ execution.setLogLocation(null)
+ execution.setStatus(statusChangeFunc(execution.getStatus))
+ }
+ })
+ }
+
+ private def dropCollections(result: String): Unit = {
+ if (result == null || result.isEmpty) {
+ return
+ }
+ // TODO: merge this logic to the server-side in-mem cleanup
+ // parse the JSON
+ try {
+ val node = objectMapper.readTree(result)
+ val collectionEntries = node.get("results")
+ // loop every collection and drop it
+ collectionEntries.forEach(collection => {
+ val storageType = collection.get("storageType").asText()
+ val collectionName = collection.get("storageKey").asText()
+ storageType match {
+ case DocumentFactory.ICEBERG =>
+ // rely on the server-side result cleanup logic.
+ }
+ })
+ } catch {
+ case e: Throwable =>
+ logger.warn("result collection cleanup failed.", e)
+ }
+ }
+
+ private def deleteReplayLog(logLocation: String): Unit = {
+ if (logLocation == null || logLocation.isEmpty) {
+ return
+ }
+ val uri = new URI(logLocation)
+ try {
+ val storage = SequentialRecordStorage.getStorage(Some(uri))
+ storage.deleteStorage()
+ } catch {
+ case throwable: Throwable =>
+ logger.warn(s"failed to delete log at $logLocation", throwable)
+ }
+ }
+
+ /**
+ * This function is called periodically and checks all expired collections and deletes them
+ */
+ private def recurringCheckExpiredResults(
+ timeToLive: Int
+ ): Unit = {
+ // retrieve all executions that are completed and their last update time goes beyond the ttl
+ val expiredResults: List[WorkflowExecutions] =
+ WorkflowExecutionsResource.getExpiredExecutionsWithResultOrLog(timeToLive)
+ // drop the collections and clean the logs
+ cleanExecutions(expiredResults, statusByte => statusByte)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/ComputingUnitWorker.scala b/amber/src/main/scala/org/apache/texera/web/ComputingUnitWorker.scala
new file mode 100644
index 00000000000..e17acb98f24
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/ComputingUnitWorker.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web
+
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.commons.jcs3.access.exception.InvalidArgumentException
+
+import scala.annotation.tailrec
+
+object ComputingUnitWorker {
+
+ type OptionMap = Map[Symbol, Any]
+
+ def parseArgs(args: Array[String]): OptionMap = {
+ @tailrec
+ def nextOption(map: OptionMap, list: List[String]): OptionMap = {
+ list match {
+ case Nil => map
+ case "--serverAddr" :: value :: tail =>
+ nextOption(map ++ Map(Symbol("serverAddr") -> value), tail)
+ case option :: tail =>
+ throw new InvalidArgumentException("unknown command-line arg")
+ }
+ }
+
+ nextOption(Map(), args.toList)
+ }
+
+ def main(args: Array[String]): Unit = {
+ val argMap = parseArgs(args)
+
+ // start actor system worker node
+ AmberRuntime.startActorWorker(argMap.get(Symbol("serverAddr")).asInstanceOf[Option[String]])
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/ServletAwareConfigurator.scala b/amber/src/main/scala/org/apache/texera/web/ServletAwareConfigurator.scala
new file mode 100644
index 00000000000..cb3628df5b3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/ServletAwareConfigurator.scala
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.http.client.utils.URLEncodedUtils
+import org.apache.texera.auth.JwtAuth.jwtConsumer
+import org.apache.texera.auth.util.HeaderField
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+
+import java.net.URI
+import java.nio.charset.Charset
+import javax.websocket.HandshakeResponse
+import javax.websocket.server.{HandshakeRequest, ServerEndpointConfig}
+import scala.jdk.CollectionConverters.{ListHasAsScala, _}
+
+/**
+ * This configurator extracts user identity from the HTTP handshake request
+ * and associates it with the ServerEndpointConfig, allowing it to be
+ * accessed by WebSocket connections.
+ */
+class ServletAwareConfigurator extends ServerEndpointConfig.Configurator with LazyLogging {
+
+ override def modifyHandshake(
+ config: ServerEndpointConfig,
+ request: HandshakeRequest,
+ response: HandshakeResponse
+ ): Unit = {
+ try {
+ val headers = request.getHeaders.asScala.view.mapValues(_.asScala.headOption).toMap
+ if (
+ headers.contains(HeaderField.UserComputingUnitAccess) &&
+ headers.contains(HeaderField.UserId) &&
+ headers.contains(HeaderField.UserName) &&
+ headers.contains(HeaderField.UserEmail)
+ ) {
+ // KUBERNETES MODE: Construct the User object from trusted headers
+ // coming from envoy and generated by access control service.
+
+ val userId = headers.get(HeaderField.UserId).flatten.map(_.toInt).get
+ val userName = headers.get(HeaderField.UserName).flatten.get
+ val userEmail = headers.get(HeaderField.UserEmail).flatten.get
+ val cuAccess = headers.get(HeaderField.UserComputingUnitAccess).flatten.getOrElse("")
+ config.getUserProperties.put(HeaderField.UserComputingUnitAccess, cuAccess)
+ logger.info(
+ s"User ID: $userId, User Name: $userName, User Email: $userEmail with CU Access: $cuAccess"
+ )
+
+ config.getUserProperties.put(
+ classOf[User].getName,
+ new User(
+ userId,
+ userName,
+ userEmail,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null
+ )
+ )
+ logger.debug(s"User created from headers: ID=$userId, Name=$userName")
+ } else {
+ // SINGLE NODE MODE: Construct the User object from JWT in query parameters.
+ val params =
+ URLEncodedUtils.parse(new URI("?" + request.getQueryString), Charset.defaultCharset())
+ config.getUserProperties.put(
+ HeaderField.UserComputingUnitAccess,
+ PrivilegeEnum.WRITE.name()
+ )
+ params.asScala
+ .map(pair => pair.getName -> pair.getValue)
+ .toMap
+ .get("access-token")
+ .map(token => {
+ val claims = jwtConsumer.process(token).getJwtClaims
+ config.getUserProperties.put(
+ classOf[User].getName,
+ new User(
+ claims.getClaimValue("userId").asInstanceOf[Long].toInt,
+ claims.getSubject,
+ String.valueOf(claims.getClaimValue("email").asInstanceOf[String]),
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null
+ )
+ )
+ })
+ }
+ } catch {
+ case e: Exception =>
+ logger.error("Failed to retrieve the User during websocket handshake", e)
+ }
+
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/SessionState.scala b/amber/src/main/scala/org/apache/texera/web/SessionState.scala
new file mode 100644
index 00000000000..b50df866afe
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/SessionState.scala
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web
+
+import io.reactivex.rxjava3.disposables.Disposable
+import org.apache.texera.amber.util.JSONUtils.objectMapper
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+import org.apache.texera.web.service.WorkflowService
+
+import javax.websocket.Session
+import scala.collection.mutable
+
+object SessionState {
+ private val sessionIdToSessionState = new mutable.HashMap[String, SessionState]()
+
+ def getState(sId: String): SessionState = {
+ sessionIdToSessionState(sId)
+ }
+
+ def setState(sId: String, state: SessionState): Unit = {
+ sessionIdToSessionState.put(sId, state)
+ }
+
+ def removeState(sId: String): Unit = {
+ sessionIdToSessionState(sId).unsubscribe()
+ sessionIdToSessionState.remove(sId)
+ }
+
+ def getAllSessionStates: Iterable[SessionState] = {
+ sessionIdToSessionState.values
+ }
+
+}
+
+class SessionState(session: Session) {
+ private var currentWorkflowState: Option[WorkflowService] = None
+ private var workflowSubscription = Disposable.empty()
+ private var executionSubscription = Disposable.empty()
+ private var userComputingUnitAccess: PrivilegeEnum = PrivilegeEnum.NONE
+
+ def send(msg: TexeraWebSocketEvent): Unit = {
+ session.getAsyncRemote.sendText(objectMapper.writeValueAsString(msg))
+ }
+
+ def getCurrentWorkflowState: Option[WorkflowService] = currentWorkflowState
+
+ def unsubscribe(): Unit = {
+ workflowSubscription.dispose()
+ executionSubscription.dispose()
+ if (currentWorkflowState.isDefined) {
+ currentWorkflowState.get.disconnect()
+ currentWorkflowState = None
+ }
+ }
+
+ def subscribe(workflowService: WorkflowService): Unit = {
+ unsubscribe()
+ currentWorkflowState = Some(workflowService)
+ workflowSubscription = workflowService.connect(evt =>
+ session.getAsyncRemote.sendText(objectMapper.writeValueAsString(evt))
+ )
+ executionSubscription = workflowService.connectToExecution(evt =>
+ session.getAsyncRemote.sendText(objectMapper.writeValueAsString(evt))
+ )
+
+ }
+
+ def setUserComputingUnitAccess(cuAccess: PrivilegeEnum): Unit = {
+ this.userComputingUnitAccess = cuAccess
+ }
+ def getUserComputingUnitAccess: PrivilegeEnum = this.userComputingUnitAccess
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/SubscriptionManager.scala b/amber/src/main/scala/org/apache/texera/web/SubscriptionManager.scala
new file mode 100644
index 00000000000..3671d1e4b6e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/SubscriptionManager.scala
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web
+
+import io.reactivex.rxjava3.disposables.Disposable
+
+import scala.collection.mutable
+
+trait SubscriptionManager {
+
+ private val subscriptions = mutable.ArrayBuffer[Disposable]()
+
+ def addSubscription(sub: Disposable): Unit = {
+ subscriptions.append(sub)
+ }
+
+ def unsubscribeAll(): Unit = {
+ subscriptions.foreach(_.dispose())
+ subscriptions.clear()
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/TexeraWebApplication.scala b/amber/src/main/scala/org/apache/texera/web/TexeraWebApplication.scala
new file mode 100644
index 00000000000..98b7c68c974
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/TexeraWebApplication.scala
@@ -0,0 +1,192 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web
+
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import com.github.dirkraft.dropwizard.fileassets.FileAssetsBundle
+import com.typesafe.scalalogging.LazyLogging
+import io.dropwizard.auth.AuthValueFactoryProvider
+import io.dropwizard.configuration.{EnvironmentVariableSubstitutor, SubstitutingSourceProvider}
+import io.dropwizard.setup.{Bootstrap, Environment}
+import io.dropwizard.websockets.WebsocketBundle
+import org.apache.texera.amber.config.StorageConfig
+import org.apache.texera.amber.engine.common.Utils
+import org.apache.texera.amber.util.ObjectMapperUtils
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.web.auth.JwtAuth.setupJwtAuth
+import org.apache.texera.web.resource._
+import org.apache.texera.web.resource.auth.{AuthResource, GoogleAuthResource}
+import org.apache.texera.web.resource.dashboard.DashboardResource
+import org.apache.texera.web.resource.dashboard.admin.execution.AdminExecutionResource
+import org.apache.texera.web.resource.dashboard.admin.settings.AdminSettingsResource
+import org.apache.texera.web.resource.dashboard.admin.user.AdminUserResource
+import org.apache.texera.web.resource.dashboard.hub.HubResource
+import org.apache.texera.web.resource.dashboard.user.UserResource
+import org.apache.texera.web.resource.dashboard.user.project.{
+ ProjectAccessResource,
+ ProjectResource,
+ PublicProjectResource
+}
+import org.apache.texera.web.resource.dashboard.user.quota.UserQuotaResource
+import org.apache.texera.web.resource.dashboard.user.workflow.{
+ WorkflowAccessResource,
+ WorkflowExecutionsResource,
+ WorkflowResource,
+ WorkflowVersionResource
+}
+import org.eclipse.jetty.server.session.SessionHandler
+import org.eclipse.jetty.servlet.{ErrorPageErrorHandler, FilterHolder}
+import org.eclipse.jetty.websocket.server.WebSocketUpgradeFilter
+import org.glassfish.jersey.server.filter.RolesAllowedDynamicFeature
+
+import java.time.Duration
+
+object TexeraWebApplication {
+
+ def main(args: Array[String]): Unit = {
+
+ // TODO: figure out a safety way of calling discardUncommittedChangesOfAllDatasets
+ // Currently in kubernetes, multiple pods calling this function can result into thread competition
+ // discardUncommittedChangesOfAllDatasets()
+
+ // start web server
+ new TexeraWebApplication().run(
+ "server",
+ Utils.amberHomePath
+ .resolve("src")
+ .resolve("main")
+ .resolve("resources")
+ .resolve("web-config.yml")
+ .toString
+ )
+ }
+}
+
+class TexeraWebApplication
+ extends io.dropwizard.Application[TexeraWebConfiguration]
+ with LazyLogging {
+
+ override def initialize(bootstrap: Bootstrap[TexeraWebConfiguration]): Unit = {
+ // enable environment variable substitution in YAML config
+ bootstrap.setConfigurationSourceProvider(
+ new SubstitutingSourceProvider(
+ bootstrap.getConfigurationSourceProvider,
+ new EnvironmentVariableSubstitutor(false)
+ )
+ )
+ // serve static frontend GUI files
+ bootstrap.addBundle(new FileAssetsBundle("../../frontend/dist", "/", "index.html"))
+ // add websocket bundle
+ bootstrap.addBundle(new WebsocketBundle(classOf[CollaborationResource]))
+ // register scala module to dropwizard default object mapper
+ bootstrap.getObjectMapper.registerModule(DefaultScalaModule)
+ }
+
+ override def run(configuration: TexeraWebConfiguration, environment: Environment): Unit = {
+ ObjectMapperUtils.warmupObjectMapperForOperatorsSerde()
+
+ // serve backend at /api
+ environment.jersey.setUrlPattern("/api/*")
+
+ SqlServer.initConnection(
+ StorageConfig.jdbcUrl,
+ StorageConfig.jdbcUsername,
+ StorageConfig.jdbcPassword
+ )
+
+ // redirect all 404 to index page, according to Angular routing requirements
+ val eph = new ErrorPageErrorHandler
+ eph.addErrorPage(404, "/")
+ environment.getApplicationContext.setErrorHandler(eph)
+
+ val webSocketUpgradeFilter =
+ WebSocketUpgradeFilter.configureContext(environment.getApplicationContext)
+ webSocketUpgradeFilter.getFactory.getPolicy.setIdleTimeout(Duration.ofHours(1).toMillis)
+ environment.getApplicationContext.setAttribute(
+ classOf[WebSocketUpgradeFilter].getName,
+ webSocketUpgradeFilter
+ )
+
+ // register SessionHandler
+ environment.jersey.register(classOf[SessionHandler])
+ environment.servlets.setSessionHandler(new SessionHandler)
+
+ environment.jersey.register(classOf[SystemMetadataResource])
+ // environment.jersey().register(classOf[MockKillWorkerResource])
+
+ environment.jersey.register(classOf[HealthCheckResource])
+
+ setupJwtAuth(environment)
+
+ environment.jersey.register(
+ new AuthValueFactoryProvider.Binder[SessionUser](classOf[SessionUser])
+ )
+ environment.jersey.register(classOf[RolesAllowedDynamicFeature])
+
+ environment.jersey.register(classOf[AuthResource])
+ environment.jersey.register(classOf[GoogleAuthResource])
+ environment.jersey.register(classOf[UserConfigResource])
+ environment.jersey.register(classOf[AdminUserResource])
+ environment.jersey.register(classOf[PublicProjectResource])
+ environment.jersey.register(classOf[WorkflowAccessResource])
+ environment.jersey.register(classOf[WorkflowResource])
+ environment.jersey.register(classOf[HubResource])
+ environment.jersey.register(classOf[UserResource])
+ environment.jersey.register(classOf[WorkflowVersionResource])
+ environment.jersey.register(classOf[ProjectResource])
+ environment.jersey.register(classOf[ProjectAccessResource])
+ environment.jersey.register(classOf[WorkflowExecutionsResource])
+ environment.jersey.register(classOf[DashboardResource])
+ environment.jersey.register(classOf[GmailResource])
+ environment.jersey.register(classOf[AdminExecutionResource])
+ environment.jersey.register(classOf[UserQuotaResource])
+ environment.jersey.register(classOf[AdminSettingsResource])
+ environment.jersey.register(classOf[AIAssistantResource])
+
+ AuthResource.createAdminUser()
+
+ // Route request logs through SLF4J, controlled by TEXERA_SERVICE_LOG_LEVEL.
+ // TODO: replace with RequestLoggingFilter.register() from common/auth once Dropwizard is upgraded to 4.x
+ val requestLogger = org.slf4j.LoggerFactory.getLogger("org.eclipse.jetty.server.RequestLog")
+ environment.getApplicationContext.addFilter(
+ new FilterHolder(new javax.servlet.Filter {
+ override def init(filterConfig: javax.servlet.FilterConfig): Unit = {}
+ override def doFilter(
+ request: javax.servlet.ServletRequest,
+ response: javax.servlet.ServletResponse,
+ chain: javax.servlet.FilterChain
+ ): Unit = {
+ chain.doFilter(request, response)
+ if (requestLogger.isInfoEnabled) {
+ val req = request.asInstanceOf[javax.servlet.http.HttpServletRequest]
+ val resp = response.asInstanceOf[javax.servlet.http.HttpServletResponse]
+ requestLogger.info(
+ s"""${req.getRemoteAddr} - "${req.getMethod} ${req.getRequestURI} ${req.getProtocol}" ${resp.getStatus}"""
+ )
+ }
+ }
+ override def destroy(): Unit = {}
+ }),
+ "/*",
+ java.util.EnumSet.allOf(classOf[javax.servlet.DispatcherType])
+ )
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/TexeraWebConfiguration.java b/amber/src/main/scala/org/apache/texera/web/TexeraWebConfiguration.java
new file mode 100644
index 00000000000..e7fcca9aa3c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/TexeraWebConfiguration.java
@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web;
+
+import io.dropwizard.Configuration;
+
+public class TexeraWebConfiguration extends Configuration {
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/WebsocketInput.scala b/amber/src/main/scala/org/apache/texera/web/WebsocketInput.scala
new file mode 100644
index 00000000000..95e7b885795
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/WebsocketInput.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web
+
+import io.reactivex.rxjava3.disposables.Disposable
+import io.reactivex.rxjava3.subjects.PublishSubject
+import org.apache.texera.web.model.websocket.request.TexeraWebSocketRequest
+
+import scala.reflect.{ClassTag, classTag}
+
+class WebsocketInput(errorHandler: Throwable => Unit) {
+ private val wsInput = PublishSubject.create[(TexeraWebSocketRequest, Option[Integer])]()
+
+ def subscribe[T <: TexeraWebSocketRequest: ClassTag](
+ callback: (T, Option[Integer]) => Unit
+ ): Disposable = {
+ wsInput.subscribe((evt: (TexeraWebSocketRequest, Option[Integer])) => {
+ evt._1 match {
+ case req: T if classTag[T].runtimeClass.isInstance(req) =>
+ try {
+ callback(req, evt._2)
+ } catch {
+ case throwable: Throwable =>
+ errorHandler(throwable)
+ }
+ case other =>
+ // skip this one because it doesn't match the type we want
+ }
+ })
+ }
+
+ def onNext(req: TexeraWebSocketRequest, uidOpt: Option[Integer]): Unit = {
+ wsInput.onNext((req, uidOpt))
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/WorkflowLifecycleManager.scala b/amber/src/main/scala/org/apache/texera/web/WorkflowLifecycleManager.scala
new file mode 100644
index 00000000000..6f2f456dda7
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/WorkflowLifecycleManager.scala
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web
+
+import org.apache.pekko.actor.Cancellable
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.RUNNING
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.executionruntimestate.ExecutionMetadataStore
+import org.apache.texera.web.storage.ExecutionStateStore
+
+import java.time.{LocalDateTime, Duration => JDuration}
+import scala.concurrent.duration.DurationInt
+
+class WorkflowLifecycleManager(id: String, cleanUpTimeout: Int, cleanUpCallback: () => Unit)
+ extends LazyLogging {
+ private var userCount = 0
+ private var cleanUpExecution: Cancellable = Cancellable.alreadyCancelled
+
+ private[this] def setCleanUpDeadline(status: WorkflowAggregatedState): Unit = {
+ synchronized {
+ if (userCount > 0 || status == RUNNING) {
+ cleanUpExecution.cancel()
+ logger.info(
+ s"[$id] workflow state clean up postponed. current user count = $userCount, workflow status = $status"
+ )
+ } else {
+ refreshDeadline()
+ }
+ }
+ }
+
+ private[this] def refreshDeadline(): Unit = {
+ if (cleanUpExecution.isCancelled || cleanUpExecution.cancel()) {
+ logger.info(
+ s"[$id] workflow state clean up will start at ${LocalDateTime.now().plus(JDuration.ofSeconds(cleanUpTimeout))}"
+ )
+ cleanUpExecution = AmberRuntime.scheduleCallThroughActorSystem(cleanUpTimeout.seconds) {
+ cleanUp()
+ }
+ }
+ }
+
+ private[this] def cleanUp(): Unit = {
+ synchronized {
+ if (userCount > 0) {
+ // do nothing
+ logger.info(s"[$id] workflow state clean up failed. current user count = $userCount")
+ } else {
+ cleanUpExecution.cancel()
+ cleanUpCallback()
+ logger.info(s"[$id] workflow state clean up completed.")
+ }
+ }
+ }
+
+ def increaseUserCount(): Unit = {
+ synchronized {
+ userCount += 1
+ cleanUpExecution.cancel()
+ logger.info(s"[$id] workflow state clean up postponed. current user count = $userCount")
+ }
+ }
+
+ def decreaseUserCount(currentWorkflowState: Option[WorkflowAggregatedState]): Unit = {
+ synchronized {
+ userCount -= 1
+ if (userCount == 0 && (currentWorkflowState.isEmpty || currentWorkflowState.get != RUNNING)) {
+ refreshDeadline()
+ } else {
+ logger.info(s"[$id] workflow state clean up postponed. current user count = $userCount")
+ }
+ }
+ }
+
+ def registerCleanUpOnStateChange(stateStore: ExecutionStateStore): Unit = {
+ cleanUpExecution.cancel()
+ stateStore.metadataStore.getStateObservable.subscribe { newState: ExecutionMetadataStore =>
+ setCleanUpDeadline(newState.state)
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/auth/GuestAuthFilter.scala b/amber/src/main/scala/org/apache/texera/web/auth/GuestAuthFilter.scala
new file mode 100644
index 00000000000..b7dda09489e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/auth/GuestAuthFilter.scala
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.auth
+
+import io.dropwizard.auth.AuthFilter
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.apache.texera.web.auth.GuestAuthFilter.GUEST
+
+import java.io.IOException
+import java.util.Optional
+import javax.annotation.{Nullable, Priority}
+import javax.ws.rs.Priorities
+import javax.ws.rs.container.{ContainerRequestContext, PreMatching}
+import javax.ws.rs.core.SecurityContext
+
+@PreMatching
+@Priority(Priorities.AUTHENTICATION) object GuestAuthFilter {
+ class Builder extends AuthFilter.AuthFilterBuilder[String, SessionUser, GuestAuthFilter] {
+ override protected def newInstance = new GuestAuthFilter
+ }
+
+ val GUEST: User =
+ new User(null, "guest", null, null, null, null, UserRoleEnum.REGULAR, null, null, null, null)
+}
+
+@PreMatching
+@Priority(Priorities.AUTHENTICATION) class GuestAuthFilter extends AuthFilter[String, SessionUser] {
+ @throws[IOException]
+ override def filter(requestContext: ContainerRequestContext): Unit =
+ authenticate(requestContext, "", "")
+
+ override protected def authenticate(
+ requestContext: ContainerRequestContext,
+ @Nullable credentials: String,
+ scheme: String
+ ): Boolean = {
+
+ val principal = Optional.of(new SessionUser(GUEST))
+ val securityContext = requestContext.getSecurityContext
+ val secure = securityContext != null && securityContext.isSecure
+ requestContext.setSecurityContext(new SecurityContext() {
+ override def getUserPrincipal: SessionUser = principal.get
+
+ override def isUserInRole(role: String): Boolean = authorizer.authorize(principal.get, role)
+
+ override def isSecure: Boolean = secure
+
+ override def getAuthenticationScheme: String = scheme
+ })
+ true
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/auth/JwtAuth.scala b/amber/src/main/scala/org/apache/texera/web/auth/JwtAuth.scala
new file mode 100644
index 00000000000..c1ade508869
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/auth/JwtAuth.scala
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.auth
+
+import com.github.toastshaman.dropwizard.auth.jwt.JwtAuthFilter
+import io.dropwizard.auth.AuthDynamicFeature
+import io.dropwizard.setup.Environment
+import org.apache.texera.auth.JwtAuth.jwtConsumer
+import org.apache.texera.auth.SessionUser
+
+// TODO: move this logic to Auth
+@Deprecated
+object JwtAuth {
+ def setupJwtAuth(environment: Environment): Unit = {
+ // register JWT Auth layer
+ environment.jersey.register(
+ new AuthDynamicFeature(
+ new JwtAuthFilter.Builder[SessionUser]()
+ .setJwtConsumer(jwtConsumer)
+ .setRealm("realm")
+ .setPrefix("Bearer")
+ .setAuthenticator(UserAuthenticator)
+ .setAuthorizer(UserRoleAuthorizer)
+ .buildAuthFilter()
+ )
+ )
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/auth/UserAuthenticator.scala b/amber/src/main/scala/org/apache/texera/web/auth/UserAuthenticator.scala
new file mode 100644
index 00000000000..e7fe67ca10a
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/auth/UserAuthenticator.scala
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.auth
+
+import com.typesafe.scalalogging.LazyLogging
+import io.dropwizard.auth.Authenticator
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.jose4j.jwt.consumer.JwtContext
+
+import java.time.OffsetDateTime
+import java.util.Optional
+
+object UserAuthenticator extends Authenticator[JwtContext, SessionUser] with LazyLogging {
+ override def authenticate(context: JwtContext): Optional[SessionUser] = {
+ // This method will be called once the token's signature has been verified,
+ // including the token secret and the expiration time
+ try {
+ val userName = context.getJwtClaims.getSubject
+ val email = context.getJwtClaims.getClaimValue("email").asInstanceOf[String]
+ val userId = context.getJwtClaims.getClaimValue("userId").asInstanceOf[Long].toInt
+ val role =
+ UserRoleEnum.valueOf(context.getJwtClaims.getClaimValue("role").asInstanceOf[String])
+ val googleId = context.getJwtClaims.getClaimValue("googleId").asInstanceOf[String]
+ val comment = context.getJwtClaims.getClaimValue("comment").asInstanceOf[String]
+ val accountCreation =
+ context.getJwtClaims.getClaimValue("accountCreation").asInstanceOf[OffsetDateTime]
+ val user =
+ new User(
+ userId,
+ userName,
+ email,
+ null,
+ googleId,
+ null,
+ role,
+ comment,
+ accountCreation,
+ null,
+ null
+ )
+ Optional.of(new SessionUser(user))
+ } catch {
+ case e: Exception =>
+ logger.error("Failed to authenticate the JwtContext", e)
+ Optional.empty()
+ }
+
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/auth/UserRoleAuthorizer.scala b/amber/src/main/scala/org/apache/texera/web/auth/UserRoleAuthorizer.scala
new file mode 100644
index 00000000000..f966caa8fe3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/auth/UserRoleAuthorizer.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.auth
+
+import io.dropwizard.auth.Authorizer
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+
+object UserRoleAuthorizer extends Authorizer[SessionUser] {
+ override def authorize(user: SessionUser, role: String): Boolean = {
+ user.isRoleOf(UserRoleEnum.valueOf(role))
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/event/CollabWebSocketEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/event/CollabWebSocketEvent.scala
new file mode 100644
index 00000000000..afc3a2df9fb
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/event/CollabWebSocketEvent.scala
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.event
+
+import com.fasterxml.jackson.annotation.JsonSubTypes.Type
+import com.fasterxml.jackson.annotation.{JsonSubTypes, JsonTypeInfo}
+import org.apache.texera.web.model.collab.response.HeartBeatResponse
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
+@JsonSubTypes(
+ Array(
+ new Type(value = classOf[CommandEvent]),
+ new Type(value = classOf[LockGrantedEvent]),
+ new Type(value = classOf[ReleaseLockEvent]),
+ new Type(value = classOf[LockRejectedEvent]),
+ new Type(value = classOf[RestoreVersionEvent]),
+ new Type(value = classOf[HeartBeatResponse]),
+ new Type(value = classOf[WorkflowAccessEvent])
+ )
+)
+trait CollabWebSocketEvent {}
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/event/CommandEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/event/CommandEvent.scala
new file mode 100644
index 00000000000..d42aa3e0819
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/event/CommandEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.event
+
+case class CommandEvent(commandMessage: String) extends CollabWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/event/LockGrantedEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/event/LockGrantedEvent.scala
new file mode 100644
index 00000000000..73930ba086b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/event/LockGrantedEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.event
+
+case class LockGrantedEvent() extends CollabWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/event/LockRejectedEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/event/LockRejectedEvent.scala
new file mode 100644
index 00000000000..227a1866404
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/event/LockRejectedEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.event
+
+case class LockRejectedEvent() extends CollabWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/event/ReleaseLockEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/event/ReleaseLockEvent.scala
new file mode 100644
index 00000000000..f7e4975013a
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/event/ReleaseLockEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.event
+
+case class ReleaseLockEvent() extends CollabWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/event/RestoreVersionEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/event/RestoreVersionEvent.scala
new file mode 100644
index 00000000000..2e2dec5db83
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/event/RestoreVersionEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.event
+
+case class RestoreVersionEvent() extends CollabWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/event/WorkflowAccessEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/event/WorkflowAccessEvent.scala
new file mode 100644
index 00000000000..5b4523337af
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/event/WorkflowAccessEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.event
+
+case class WorkflowAccessEvent(workflowReadonly: Boolean) extends CollabWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/request/AcquireLockRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/request/AcquireLockRequest.scala
new file mode 100644
index 00000000000..2bea78f53d2
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/request/AcquireLockRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.request
+
+case class AcquireLockRequest() extends CollabWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/request/CollabWebSocketRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/request/CollabWebSocketRequest.scala
new file mode 100644
index 00000000000..f8f1de5220c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/request/CollabWebSocketRequest.scala
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.request
+
+import com.fasterxml.jackson.annotation.JsonSubTypes.Type
+import com.fasterxml.jackson.annotation.{JsonSubTypes, JsonTypeInfo}
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
+@JsonSubTypes(
+ Array(
+ new Type(value = classOf[CommandRequest]),
+ new Type(value = classOf[AcquireLockRequest]),
+ new Type(value = classOf[TryLockRequest]),
+ new Type(value = classOf[RestoreVersionRequest]),
+ new Type(value = classOf[WIdRequest]),
+ new Type(value = classOf[HeartBeatRequest])
+ )
+)
+trait CollabWebSocketRequest {}
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/request/CommandRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/request/CommandRequest.scala
new file mode 100644
index 00000000000..798d38d4107
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/request/CommandRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.request
+
+case class CommandRequest(commandMessage: String) extends CollabWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/request/HeartBeatRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/request/HeartBeatRequest.scala
new file mode 100644
index 00000000000..ffeb81df76c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/request/HeartBeatRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.request
+
+case class HeartBeatRequest() extends CollabWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/request/RestoreVersionRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/request/RestoreVersionRequest.scala
new file mode 100644
index 00000000000..ce05116a3c4
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/request/RestoreVersionRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.request
+
+case class RestoreVersionRequest() extends CollabWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/request/TryLockRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/request/TryLockRequest.scala
new file mode 100644
index 00000000000..05c35ad2455
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/request/TryLockRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.request
+
+case class TryLockRequest() extends CollabWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/request/WIdRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/request/WIdRequest.scala
new file mode 100644
index 00000000000..c90912b6ecc
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/request/WIdRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.request
+
+case class WIdRequest(wId: Int) extends CollabWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/collab/response/HeartBeatResponse.scala b/amber/src/main/scala/org/apache/texera/web/model/collab/response/HeartBeatResponse.scala
new file mode 100644
index 00000000000..a0b20fbe517
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/collab/response/HeartBeatResponse.scala
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.collab.response
+
+import org.apache.texera.web.model.collab.event.CollabWebSocketEvent
+
+case class HeartBeatResponse() extends CollabWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/common/AccessEntry.scala b/amber/src/main/scala/org/apache/texera/web/model/common/AccessEntry.scala
new file mode 100644
index 00000000000..4cfacfbc84b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/common/AccessEntry.scala
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.common
+
+import org.jooq.EnumType
+
+case class AccessEntry(email: String, name: String, privilege: EnumType) {}
diff --git a/amber/src/main/scala/org/apache/texera/web/model/http/request/auth/UserLoginRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/http/request/auth/UserLoginRequest.scala
new file mode 100644
index 00000000000..a034c45cfaf
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/http/request/auth/UserLoginRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.http.request.auth
+
+case class UserLoginRequest(username: String, password: String)
diff --git a/amber/src/main/scala/org/apache/texera/web/model/http/request/auth/UserRegistrationRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/http/request/auth/UserRegistrationRequest.scala
new file mode 100644
index 00000000000..f28b27c87dd
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/http/request/auth/UserRegistrationRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.http.request.auth
+
+case class UserRegistrationRequest(username: String, password: String)
diff --git a/amber/src/main/scala/org/apache/texera/web/model/http/request/result/ResultExportRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/http/request/result/ResultExportRequest.scala
new file mode 100644
index 00000000000..56728a36762
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/http/request/result/ResultExportRequest.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.http.request.result
+
+import play.api.libs.json._
+
+case class OperatorExportInfo(
+ id: String,
+ outputType: String
+)
+
+object OperatorExportInfo {
+ implicit val fmt: OFormat[OperatorExportInfo] = Json.format[OperatorExportInfo]
+}
+
+case class ResultExportRequest(
+ exportType: String, // e.g. "csv", "google_sheet", "arrow", "data"
+ workflowId: Int,
+ workflowName: String,
+ operators: List[OperatorExportInfo],
+ datasetIds: List[Int],
+ rowIndex: Int, // used by "data" export
+ columnIndex: Int, // used by "data" export
+ filename: String, // optional filename override
+ // TODO: remove it once the lifecycle of result and compute are unbundled
+ computingUnitId: Int // the id of the computing unit
+)
+
+object ResultExportRequest {
+ implicit val fmt: OFormat[ResultExportRequest] = Json.format[ResultExportRequest]
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/model/http/response/SchemaPropagationResponse.scala b/amber/src/main/scala/org/apache/texera/web/model/http/response/SchemaPropagationResponse.scala
new file mode 100644
index 00000000000..7cdee2cc814
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/http/response/SchemaPropagationResponse.scala
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.http.response
+
+import org.apache.texera.amber.core.tuple.Attribute
+
+case class SchemaPropagationResponse(
+ code: Int,
+ result: Map[String, List[Option[List[Attribute]]]],
+ message: String
+)
diff --git a/amber/src/main/scala/org/apache/texera/web/model/http/response/TokenIssueResponse.scala b/amber/src/main/scala/org/apache/texera/web/model/http/response/TokenIssueResponse.scala
new file mode 100644
index 00000000000..11c894e0131
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/http/response/TokenIssueResponse.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.http.response
+
+case class TokenIssueResponse(accessToken: String)
diff --git a/amber/src/main/scala/org/apache/texera/web/model/http/response/result/ResultExportResponse.scala b/amber/src/main/scala/org/apache/texera/web/model/http/response/result/ResultExportResponse.scala
new file mode 100644
index 00000000000..53b4372bce9
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/http/response/result/ResultExportResponse.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.http.response.result
+
+case class ResultExportResponse(status: String, message: String)
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/CacheStatusUpdateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/CacheStatusUpdateEvent.scala
new file mode 100644
index 00000000000..cd1c3448bc0
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/CacheStatusUpdateEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+case class CacheStatusUpdateEvent(cacheStatusMap: Map[String, String]) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/ExecutionDurationUpdateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/ExecutionDurationUpdateEvent.scala
new file mode 100644
index 00000000000..8c5f9c368a5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/ExecutionDurationUpdateEvent.scala
@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+case class ExecutionDurationUpdateEvent(duration: Long, isRunning: Boolean)
+ extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/ExecutionStatusEnum.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/ExecutionStatusEnum.scala
new file mode 100644
index 00000000000..042f3ce1f39
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/ExecutionStatusEnum.scala
@@ -0,0 +1,18 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/OperatorStatisticsUpdateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/OperatorStatisticsUpdateEvent.scala
new file mode 100644
index 00000000000..d4aa6117c91
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/OperatorStatisticsUpdateEvent.scala
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+case class OperatorAggregatedMetrics(
+ operatorState: String,
+ aggregatedInputRowCount: Long,
+ aggregatedInputSize: Long,
+ inputPortMetrics: Map[String, Long],
+ aggregatedOutputRowCount: Long,
+ aggregatedOutputSize: Long,
+ outputPortMetrics: Map[String, Long],
+ numWorkers: Long,
+ aggregatedDataProcessingTime: Long,
+ aggregatedControlProcessingTime: Long,
+ aggregatedIdleTime: Long
+)
+
+case class OperatorStatisticsUpdateEvent(operatorStatistics: Map[String, OperatorAggregatedMetrics])
+ extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/PaginatedResultEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/PaginatedResultEvent.scala
new file mode 100644
index 00000000000..081d072180b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/PaginatedResultEvent.scala
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+import com.fasterxml.jackson.databind.node.ObjectNode
+import org.apache.texera.amber.core.tuple.Attribute
+import org.apache.texera.web.model.websocket.request.ResultPaginationRequest
+
+object PaginatedResultEvent {
+ def apply(
+ req: ResultPaginationRequest,
+ table: List[ObjectNode],
+ schema: List[Attribute]
+ ): PaginatedResultEvent = {
+ PaginatedResultEvent(req.requestID, req.operatorID, req.pageIndex, table, schema)
+ }
+}
+
+case class PaginatedResultEvent(
+ requestID: String,
+ operatorID: String,
+ pageIndex: Int,
+ table: List[ObjectNode],
+ schema: List[Attribute]
+) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/RegionStateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/RegionStateEvent.scala
new file mode 100644
index 00000000000..2e16ba9243e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/RegionStateEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+case class RegionStateEvent(id: Long, state: String) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/TexeraWebSocketEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/TexeraWebSocketEvent.scala
new file mode 100644
index 00000000000..da072c80ea5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/TexeraWebSocketEvent.scala
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+import com.fasterxml.jackson.annotation.JsonSubTypes.Type
+import com.fasterxml.jackson.annotation.{JsonSubTypes, JsonTypeInfo}
+import org.apache.texera.web.model.websocket.event.python.ConsoleUpdateEvent
+import org.apache.texera.web.model.websocket.response.python.PythonExpressionEvaluateResponse
+import org.apache.texera.web.model.websocket.response.{HeartBeatResponse, ModifyLogicResponse}
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
+@JsonSubTypes(
+ Array(
+ new Type(value = classOf[HeartBeatResponse]),
+ new Type(value = classOf[WorkflowErrorEvent]),
+ new Type(value = classOf[WorkflowStateEvent]),
+ new Type(value = classOf[OperatorStatisticsUpdateEvent]),
+ new Type(value = classOf[WebResultUpdateEvent]),
+ new Type(value = classOf[ConsoleUpdateEvent]),
+ new Type(value = classOf[CacheStatusUpdateEvent]),
+ new Type(value = classOf[PaginatedResultEvent]),
+ new Type(value = classOf[PythonExpressionEvaluateResponse]),
+ new Type(value = classOf[WorkerAssignmentUpdateEvent]),
+ new Type(value = classOf[ModifyLogicResponse])
+ )
+)
+trait TexeraWebSocketEvent {}
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WebResultUpdateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WebResultUpdateEvent.scala
new file mode 100644
index 00000000000..7e31be1ca7f
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WebResultUpdateEvent.scala
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+import org.apache.texera.web.service.ExecutionResultService.WebResultUpdate
+
+case class WebResultUpdateEvent(
+ updates: Map[String, WebResultUpdate],
+ tableStats: Map[String, Map[String, Map[String, Any]]]
+) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkerAssignmentUpdateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkerAssignmentUpdateEvent.scala
new file mode 100644
index 00000000000..0b3bce71f88
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkerAssignmentUpdateEvent.scala
@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+case class WorkerAssignmentUpdateEvent(operatorId: String, workerIds: Seq[String])
+ extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowAvailableResultEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowAvailableResultEvent.scala
new file mode 100644
index 00000000000..bc9e006cd39
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowAvailableResultEvent.scala
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+import org.apache.texera.web.model.websocket.event.WorkflowAvailableResultEvent.OperatorAvailableResult
+import org.apache.texera.web.service.ExecutionResultService.WebOutputMode
+
+object WorkflowAvailableResultEvent {
+ case class OperatorAvailableResult(
+ cacheValid: Boolean,
+ outputMode: WebOutputMode
+ )
+}
+
+case class WorkflowAvailableResultEvent(
+ availableOperators: Map[String, OperatorAvailableResult]
+) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowErrorEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowErrorEvent.scala
new file mode 100644
index 00000000000..5169a1f759e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowErrorEvent.scala
@@ -0,0 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+import org.apache.texera.amber.core.workflowruntimestate.WorkflowFatalError
+
+case class WorkflowErrorEvent(
+ fatalErrors: Seq[WorkflowFatalError]
+) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowStateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowStateEvent.scala
new file mode 100644
index 00000000000..3bad65234fe
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/WorkflowStateEvent.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event
+
+case class WorkflowStateEvent(state: String) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/event/python/ConsoleUpdateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/python/ConsoleUpdateEvent.scala
new file mode 100644
index 00000000000..4ec4732e358
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/event/python/ConsoleUpdateEvent.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.event.python
+
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.ConsoleMessage
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+
+object ConsoleUpdateEvent {}
+
+case class ConsoleUpdateEvent(
+ operatorId: String,
+ messages: Seq[ConsoleMessage]
+) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/EditingTimeCompilationRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/EditingTimeCompilationRequest.scala
new file mode 100644
index 00000000000..e15b441fcae
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/EditingTimeCompilationRequest.scala
@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+import org.apache.texera.amber.operator.LogicalOp
+import org.apache.texera.workflow.LogicalLink
+
+case class EditingTimeCompilationRequest(
+ operators: List[LogicalOp],
+ links: List[LogicalLink],
+ opsToViewResult: List[String],
+ opsToReuseResult: List[String]
+) extends TexeraWebSocketRequest {
+
+ def toLogicalPlanPojo: LogicalPlanPojo = {
+ LogicalPlanPojo(operators, links, opsToViewResult, opsToReuseResult)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/HeartBeatRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/HeartBeatRequest.scala
new file mode 100644
index 00000000000..712aa98773d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/HeartBeatRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+case class HeartBeatRequest() extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/ModifyLogicRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/ModifyLogicRequest.scala
new file mode 100644
index 00000000000..84ae56eecc6
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/ModifyLogicRequest.scala
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+import org.apache.texera.amber.operator.LogicalOp
+
+case class ModifyLogicRequest(operator: LogicalOp) extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/ResultPaginationRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/ResultPaginationRequest.scala
new file mode 100644
index 00000000000..4a3e0a58a3e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/ResultPaginationRequest.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+case class ResultPaginationRequest(
+ requestID: String,
+ operatorID: String,
+ pageIndex: Int,
+ pageSize: Int,
+ columnOffset: Int = 0,
+ columnLimit: Int = Int.MaxValue,
+ columnSearch: Option[String] = None
+) extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/RetryRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/RetryRequest.scala
new file mode 100644
index 00000000000..6a0bfba0391
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/RetryRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+case class RetryRequest(workers: Seq[String]) extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/SkipTupleRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/SkipTupleRequest.scala
new file mode 100644
index 00000000000..11f9a3e63f2
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/SkipTupleRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+case class SkipTupleRequest(workerIds: Array[String]) extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/TexeraWebSocketRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/TexeraWebSocketRequest.scala
new file mode 100644
index 00000000000..1a92514ff90
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/TexeraWebSocketRequest.scala
@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+import com.fasterxml.jackson.annotation.JsonSubTypes.Type
+import com.fasterxml.jackson.annotation.{JsonSubTypes, JsonTypeInfo}
+import org.apache.texera.web.model.websocket.request.python.{
+ DebugCommandRequest,
+ PythonExpressionEvaluateRequest
+}
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
+@JsonSubTypes(
+ Array(
+ new Type(value = classOf[EditingTimeCompilationRequest]),
+ new Type(value = classOf[HeartBeatRequest]),
+ new Type(value = classOf[ModifyLogicRequest]),
+ new Type(value = classOf[ResultPaginationRequest]),
+ new Type(value = classOf[RetryRequest]),
+ new Type(value = classOf[SkipTupleRequest]),
+ new Type(value = classOf[WorkflowExecuteRequest]),
+ new Type(value = classOf[WorkflowKillRequest]),
+ new Type(value = classOf[WorkflowPauseRequest]),
+ new Type(value = classOf[WorkflowResumeRequest]),
+ new Type(value = classOf[PythonExpressionEvaluateRequest]),
+ new Type(value = classOf[DebugCommandRequest]),
+ new Type(value = classOf[WorkflowCheckpointRequest])
+ )
+)
+trait TexeraWebSocketRequest {}
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowCheckpointRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowCheckpointRequest.scala
new file mode 100644
index 00000000000..79c82a25a40
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowCheckpointRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+case class WorkflowCheckpointRequest() extends TexeraWebSocketRequest()
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowExecuteRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowExecuteRequest.scala
new file mode 100644
index 00000000000..a346a1ec0a4
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowExecuteRequest.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize
+import org.apache.texera.amber.core.workflow.WorkflowSettings
+import org.apache.texera.amber.operator.LogicalOp
+import org.apache.texera.workflow.LogicalLink
+
+case class ReplayExecutionInfo(
+ @JsonDeserialize(contentAs = classOf[java.lang.Long])
+ eid: Long,
+ interaction: String
+)
+
+case class WorkflowExecuteRequest(
+ executionName: String,
+ engineVersion: String,
+ logicalPlan: LogicalPlanPojo,
+ replayFromExecution: Option[ReplayExecutionInfo], // contains execution Id, interaction Id.
+ workflowSettings: WorkflowSettings,
+ emailNotificationEnabled: Boolean,
+ computingUnitId: Int
+) extends TexeraWebSocketRequest
+
+case class LogicalPlanPojo(
+ operators: List[LogicalOp],
+ links: List[LogicalLink],
+ opsToViewResult: List[String],
+ opsToReuseResult: List[String]
+)
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowKillRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowKillRequest.scala
new file mode 100644
index 00000000000..da5adfe2942
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowKillRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+case class WorkflowKillRequest() extends TexeraWebSocketRequest()
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowPauseRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowPauseRequest.scala
new file mode 100644
index 00000000000..0eb23d6b21c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowPauseRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+case class WorkflowPauseRequest() extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowResumeRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowResumeRequest.scala
new file mode 100644
index 00000000000..669202616ed
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/WorkflowResumeRequest.scala
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request
+
+case class WorkflowResumeRequest() extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/python/DebugCommandRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/python/DebugCommandRequest.scala
new file mode 100644
index 00000000000..20dccf6f788
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/python/DebugCommandRequest.scala
@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request.python
+
+import org.apache.texera.web.model.websocket.request.TexeraWebSocketRequest
+
+case class DebugCommandRequest(operatorId: String, workerId: String, cmd: String)
+ extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/request/python/PythonExpressionEvaluateRequest.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/python/PythonExpressionEvaluateRequest.scala
new file mode 100644
index 00000000000..11086448c3c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/request/python/PythonExpressionEvaluateRequest.scala
@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.request.python
+
+import org.apache.texera.web.model.websocket.request.TexeraWebSocketRequest
+
+case class PythonExpressionEvaluateRequest(expression: String, operatorId: String)
+ extends TexeraWebSocketRequest
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/response/ClusterStatusUpdateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/ClusterStatusUpdateEvent.scala
new file mode 100644
index 00000000000..5e6b187e9cf
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/ClusterStatusUpdateEvent.scala
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.response
+
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+
+case class ClusterStatusUpdateEvent(numWorkers: Int) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/response/HeartBeatResponse.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/HeartBeatResponse.scala
new file mode 100644
index 00000000000..bf4beca397c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/HeartBeatResponse.scala
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.response
+
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+
+case class HeartBeatResponse() extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/response/ModifyLogicResponse.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/ModifyLogicResponse.scala
new file mode 100644
index 00000000000..d2270cfeff1
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/ModifyLogicResponse.scala
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.response
+
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+
+case class ModifyLogicResponse(
+ opId: String,
+ isValid: Boolean,
+ errorMessage: String
+) extends TexeraWebSocketEvent
+
+case class ModifyLogicCompletedEvent(
+ opIds: List[String]
+) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/response/RegionUpdateEvent.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/RegionUpdateEvent.scala
new file mode 100644
index 00000000000..9578b28b46d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/RegionUpdateEvent.scala
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.response
+
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+
+case class RegionUpdateEvent(regions: List[(Long, List[String])]) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/model/websocket/response/python/PythonExpressionEvaluateResponse.scala b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/python/PythonExpressionEvaluateResponse.scala
new file mode 100644
index 00000000000..f4f4a77ab4a
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/model/websocket/response/python/PythonExpressionEvaluateResponse.scala
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.model.websocket.response.python
+
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EvaluatedValue
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+
+case class PythonExpressionEvaluateResponse(
+ expression: String,
+ values: Seq[EvaluatedValue]
+) extends TexeraWebSocketEvent
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/CollaborationResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/CollaborationResource.scala
new file mode 100644
index 00000000000..4ac0e8769c9
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/CollaborationResource.scala
@@ -0,0 +1,210 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.util.JSONUtils
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.apache.texera.web.ServletAwareConfigurator
+import org.apache.texera.web.model.collab.event._
+import org.apache.texera.web.model.collab.request._
+import org.apache.texera.web.model.collab.response.HeartBeatResponse
+import org.apache.texera.web.resource.CollaborationResource._
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowAccessResource
+
+import javax.websocket.server.ServerEndpoint
+import javax.websocket.{OnClose, OnMessage, OnOpen, Session}
+import scala.collection.mutable
+import scala.jdk.CollectionConverters.MapHasAsScala
+
+object CollaborationResource {
+ final val sessionIdSessionMap = new mutable.HashMap[String, Session]()
+ final val sessionIdWIdMap = new mutable.HashMap[String, Int]()
+ final val sessionIdUIdMap = new mutable.HashMap[String, Int]()
+ final val wIdSessionIdsMap = new mutable.HashMap[Int, mutable.Set[String]]()
+ final val wIdLockHolderSessionIdMap = new mutable.HashMap[Int, String]()
+ final val DUMMY_WID = -1
+
+ private def checkIsReadOnly(wId: Int, uId: Int): Boolean = {
+ !WorkflowAccessResource.hasWriteAccess(Integer.valueOf(wId), Integer.valueOf(uId))
+ }
+}
+
+@ServerEndpoint(
+ value = "/wsapi/collab",
+ configurator = classOf[ServletAwareConfigurator]
+)
+class CollaborationResource extends LazyLogging {
+
+ final val objectMapper = JSONUtils.objectMapper
+
+ @OnMessage
+ def myOnMsg(senderSession: Session, message: String): Unit = {
+ val request = objectMapper.readValue(message, classOf[CollabWebSocketRequest])
+ val uidOpt = senderSession.getUserProperties.asScala
+ .get(classOf[User].getName)
+ .map(_.asInstanceOf[User].getUid)
+ val senderSessId = senderSession.getId
+ request match {
+ case wIdRequest: WIdRequest =>
+ val wId: Int = uidOpt match {
+ case Some(uId) =>
+ sessionIdUIdMap(senderSessId) = uId.intValue()
+ val wId = wIdRequest.wId
+ logger.info("New session from " + uId + " on workflow with workflowId: " + wId)
+ wId
+ case None =>
+ // use a fixed wid for reconnection
+ DUMMY_WID
+ }
+ sessionIdWIdMap(senderSessId) = wId
+ val sessionIdSet: mutable.Set[String] =
+ wIdSessionIdsMap.get(wId) match {
+ case Some(set) =>
+ set.union(Set(senderSessId))
+ case None =>
+ mutable.Set(senderSessId)
+ }
+ wIdSessionIdsMap(wId) = sessionIdSet
+
+ case commandRequest: CommandRequest =>
+ logger.debug("Received command message: " + commandRequest.commandMessage)
+ for (sessionId <- sessionIdSessionMap.keySet) {
+ // only send to other sessions, not the session that sent the message
+ val session = sessionIdSessionMap(sessionId)
+ val sessionWId = sessionIdWIdMap.get(sessionId)
+ val senderWId = sessionIdWIdMap.get(senderSessId)
+ if (
+ session != senderSession && sessionWId.isDefined && senderWId.isDefined && senderWId == sessionWId
+ ) {
+ send(session, CommandEvent(commandRequest.commandMessage))
+ logger.debug("Message propagated to workflow " + sessionWId.toString)
+ }
+ }
+ case heartbeat: HeartBeatRequest =>
+ send(senderSession, HeartBeatResponse())
+
+ case tryLock: TryLockRequest =>
+ val wId = sessionIdWIdMap(senderSessId)
+ if (wId == DUMMY_WID) {
+ send(senderSession, WorkflowAccessEvent(workflowReadonly = false))
+ send(senderSession, LockGrantedEvent())
+ } else {
+ val uId = sessionIdUIdMap(senderSessId)
+ if (checkIsReadOnly(wId, uId)) {
+ send(senderSession, LockRejectedEvent())
+ send(senderSession, WorkflowAccessEvent(workflowReadonly = true))
+ if (!wIdLockHolderSessionIdMap.keySet.contains(wId)) {
+ wIdLockHolderSessionIdMap(wId) = null
+ }
+ } else {
+ send(senderSession, WorkflowAccessEvent(workflowReadonly = false))
+ if (
+ !wIdLockHolderSessionIdMap.keySet.contains(wId) || wIdLockHolderSessionIdMap(
+ wId
+ ) == null || wIdLockHolderSessionIdMap(wId) == senderSessId
+ ) {
+ grantLock(senderSession, senderSessId, wId)
+ } else {
+ send(senderSession, LockRejectedEvent())
+ }
+ }
+ }
+
+ case acquireLock: AcquireLockRequest =>
+ try {
+ val senderSessId = senderSession.getId
+ val senderWid = sessionIdWIdMap(senderSessId)
+ if (wIdLockHolderSessionIdMap(senderWid) != senderSessId) {
+ val holderSessId = wIdLockHolderSessionIdMap(senderWid)
+ val holderSession = sessionIdSessionMap(holderSessId)
+ send(holderSession, ReleaseLockEvent())
+ send(senderSession, LockGrantedEvent())
+ wIdLockHolderSessionIdMap(senderWid) = senderSessId
+ logger.info("Session " + senderSessId + " has lock on " + senderWid)
+ } else {
+ send(senderSession, LockGrantedEvent())
+ }
+ } catch {
+ case exception: Exception =>
+ logger.error("Session " + senderSessId + " acquire lock failed.")
+ throw exception
+ }
+
+ case restoreVersion: RestoreVersionRequest =>
+ for (sessionId <- sessionIdSessionMap.keySet) {
+ // only send to other sessions, not the session that sent the message
+ val session = sessionIdSessionMap(sessionId)
+ val sessionStateId = sessionIdWIdMap.get(sessionId)
+ val senderStateId = sessionIdWIdMap.get(senderSession.getId)
+ if (
+ session != senderSession && sessionStateId.isDefined && senderStateId.isDefined && senderStateId == sessionStateId
+ ) {
+ send(session, RestoreVersionEvent())
+ logger.info("Reload propagated to workflow " + sessionStateId.toString)
+ }
+ }
+ }
+ }
+
+ @OnOpen
+ def myOnOpen(session: Session): Unit = {
+ sessionIdSessionMap += (session.getId -> session)
+ }
+
+ @OnClose
+ def myOnClose(senderSession: Session): Unit = {
+ val senderSessId = senderSession.getId
+ sessionIdSessionMap -= senderSessId
+ if (sessionIdWIdMap.contains(senderSessId)) {
+ val wId = sessionIdWIdMap(senderSessId)
+ if (wIdSessionIdsMap.contains(wId)) {
+ wIdSessionIdsMap(wId) -= senderSessId
+ if (
+ wIdLockHolderSessionIdMap.contains(wId) && wIdLockHolderSessionIdMap(wId) == senderSessId
+ ) {
+ wIdLockHolderSessionIdMap(wId) = null
+ val set = wIdSessionIdsMap(wId)
+ if (set.nonEmpty) {
+ var granted = false
+ for (sessId <- set) {
+ if (!checkIsReadOnly(wId, sessionIdUIdMap(sessId)) && !granted) {
+ grantLock(sessionIdSessionMap(sessId), sessId, wId)
+ granted = true
+ }
+ }
+ }
+ }
+ }
+ sessionIdWIdMap -= senderSessId
+ }
+ logger.info("Session " + senderSessId + " disconnected")
+ }
+
+ private def grantLock(session: Session, sessionId: String, wId: Int): Unit = {
+ wIdLockHolderSessionIdMap(wId) = sessionId
+ logger.info("Session " + sessionId + " has lock on " + wId + " now")
+ send(session, LockGrantedEvent())
+ }
+
+ private def send(session: Session, msg: CollabWebSocketEvent): Unit = {
+ session.getAsyncRemote.sendText(objectMapper.writeValueAsString(msg))
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/EmailTemplate.scala b/amber/src/main/scala/org/apache/texera/web/resource/EmailTemplate.scala
new file mode 100644
index 00000000000..d43ca1e5079
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/EmailTemplate.scala
@@ -0,0 +1,115 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import org.apache.texera.config.UserSystemConfig
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+
+/**
+ * EmailTemplate provides factory methods to generate email messages
+ * for different user notification scenarios.
+ */
+object EmailTemplate {
+
+ private val deployment: String =
+ UserSystemConfig.appDomain.map(_.replaceFirst("^https?://", "")).getOrElse("")
+
+ private val projectName: String =
+ UserSystemConfig.projectName
+
+ /**
+ * Creates an email message for user registration notifications.
+ * Depending on the 'toAdmin' flag, it either notifies an administrator
+ * of a pending account request or acknowledges receipt to the user.
+ *
+ * @param receiverEmail the email address of the receiver (admin or user)
+ * @param userEmail optional; the email address of the user requesting an account (only needed if toAdmin is true)
+ * @param toAdmin flag indicating whether the notification is for the admin (true) or the user (false)
+ * @return an EmailMessage ready to be sent
+ */
+ def userRegistrationNotification(
+ receiverEmail: String,
+ userEmail: Option[String],
+ affiliation: Option[String],
+ reason: Option[String],
+ toAdmin: Boolean
+ ): EmailMessage = {
+ if (toAdmin) {
+ val subject =
+ s"New Account Request Pending Approval${if (deployment.nonEmpty) s" for [$deployment]"
+ else ""}"
+ val content =
+ s"""
+ |Hello Admin,
+ |
+ |A new user has attempted to log in or register, but their account is not yet approved.
+ |Please review the account request for the following user:
+ |
+ |Email: ${userEmail.getOrElse("Unknown")}
+ |Affiliation: ${affiliation.filter(_.trim.nonEmpty).getOrElse("Not provided")}
+ |Reason: ${reason.filter(_.trim.nonEmpty).getOrElse("Not provided")}
+ |
+ |Visit the admin panel at: $deployment
+ |
+ |Thanks!
+ |""".stripMargin
+ EmailMessage(subject = subject, content = content, receiver = receiverEmail)
+ } else {
+ val subject =
+ s"Account Request Received${if (deployment.nonEmpty) s" for [$deployment]" else ""}"
+ val content =
+ s"""
+ |Hello,
+ |
+ |Thank you for submitting your account request.
+ |We have received your request and it is currently under review.
+ |You will be notified once your account has been approved.
+ |
+ |Thank you for your interest in $projectName!
+ |""".stripMargin
+ EmailMessage(subject = subject, content = content, receiver = receiverEmail)
+ }
+ }
+
+ /**
+ * Creates an email message to notify a user
+ * that their role has been updated.
+ *
+ * @param receiverEmail the user's email address
+ * @param newRole the new role assigned to the user
+ * @return an EmailMessage ready to be sent to the user
+ */
+ def createRoleChangeTemplate(receiverEmail: String, newRole: UserRoleEnum): EmailMessage = {
+ val subject =
+ s"Your Role Has Been Updated${if (deployment.nonEmpty) s" for [$deployment]" else ""}"
+ val content =
+ s"""
+ |Hello,
+ |
+ |Your user role has been updated to: $newRole.
+ |
+ |If you have any questions, please contact the administrator.
+ |
+ |Thank you for using $projectName!
+ |""".stripMargin
+
+ EmailMessage(subject = subject, content = content, receiver = receiverEmail)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/GmailResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/GmailResource.scala
new file mode 100644
index 00000000000..f06f1f92102
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/GmailResource.scala
@@ -0,0 +1,208 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.config.UserSystemConfig
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.UserDao
+import org.apache.texera.web.resource.EmailTemplate.userRegistrationNotification
+import org.apache.texera.web.resource.GmailResource.{isValidEmail, sendEmail, senderGmail, userDao}
+import org.slf4j.LoggerFactory
+
+import javax.annotation.security.RolesAllowed
+import javax.mail.internet.{InternetAddress, MimeMessage}
+import javax.mail.{Message, PasswordAuthentication, Session, Transport}
+import javax.ws.rs._
+import scala.util.{Failure, Success, Try}
+
+case class EmailMessage(
+ receiver: String,
+ subject: String,
+ content: String,
+ affiliation: Option[String] = None,
+ reason: Option[String] = None
+)
+
+object GmailResource {
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def userDao = new UserDao(context.configuration)
+
+ private lazy val senderGmail: String = UserSystemConfig.gmail
+ private val smtpProperties = Map(
+ "mail.smtp.host" -> "smtp.gmail.com",
+ "mail.smtp.port" -> "465",
+ "mail.smtp.auth" -> "true",
+ "mail.smtp.socketFactory.port" -> "465",
+ "mail.smtp.socketFactory.class" -> "javax.net.ssl.SSLSocketFactory"
+ )
+
+ private def createSession(): Session = {
+ Session.getInstance(
+ smtpProperties.foldLeft(new java.util.Properties) {
+ case (props, (key, value)) =>
+ props.put(key, value)
+ props
+ },
+ new javax.mail.Authenticator() {
+ override def getPasswordAuthentication: PasswordAuthentication =
+ new PasswordAuthentication(senderGmail, UserSystemConfig.smtpPassword)
+ }
+ )
+ }
+
+ private def createMimeMessage(
+ session: Session,
+ emailMessage: EmailMessage,
+ recipientEmail: String
+ ): MimeMessage = {
+ val email = new MimeMessage(session)
+ email.setFrom(new InternetAddress(senderGmail))
+ email.addRecipient(Message.RecipientType.TO, new InternetAddress(recipientEmail))
+ email.setSubject(emailMessage.subject)
+ email.setText(emailMessage.content)
+ email
+ }
+
+ def sendEmail(
+ emailMessage: EmailMessage,
+ recipientEmail: String
+ ): Either[String, Unit] = {
+ val logger = LoggerFactory.getLogger(this.getClass)
+
+ if (!isValidEmail(recipientEmail)) {
+ logger.warn(s"Attempted to send email to invalid address: $recipientEmail")
+ return Left("Invalid email format")
+ }
+
+ Try {
+ val session = createSession()
+ val email = createMimeMessage(session, withDomain(emailMessage), recipientEmail)
+ Transport.send(email)
+ } match {
+ case Success(_) => Right(())
+ case Failure(exception) => Left(s"Failed to send email: ${exception.getMessage}")
+ }
+ }
+
+ /**
+ * Validates whether a given email address has a basic correct format.
+ *
+ * This method uses a regular expression to ensure the email:
+ * - Has a valid local part containing letters, numbers, '+', '_', '.', or '-'
+ * - Contains a single '@' character separating the local part and domain
+ * - Has a valid domain containing letters, numbers, '.' or '-'
+ * - Ends with a domain suffix (e.g., '.com', '.net') that is at least two letters long
+ *
+ * @param email the email address to validate
+ * @return true if the email matches the expected format, false otherwise
+ */
+ private def isValidEmail(email: String): Boolean = {
+ val emailRegex = "^[A-Za-z0-9+_.-]+@[A-Za-z0-9.-]+\\.[A-Za-z]{2,}$".r
+ email != null && emailRegex.matches(email)
+ }
+
+ private def withDomain(message: EmailMessage): EmailMessage = {
+ val newContent = UserSystemConfig.appDomain match {
+ case Some(domain) =>
+ s"""${message.content}
+ |
+ |—
+ |Sent from: $domain
+ |""".stripMargin
+ case None => message.content
+ }
+
+ message.copy(content = newContent)
+ }
+}
+
+@Path("/gmail")
+class GmailResource {
+ @PUT
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/send")
+ def sendEmailRequest(emailMessage: EmailMessage, @Auth user: SessionUser): Unit = {
+ val recipientEmail = if (emailMessage.receiver.isEmpty) user.getEmail else emailMessage.receiver
+ sendEmail(emailMessage, recipientEmail)
+ }
+
+ @GET
+ @RolesAllowed(Array("ADMIN"))
+ @Path("/sender/email")
+ def getSenderEmail: String = senderGmail
+
+ @POST
+ @Path("/notify-unauthorized")
+ def notifyUnauthorizedUser(emailMessage: EmailMessage): Unit = {
+ val logger = LoggerFactory.getLogger(this.getClass)
+
+ if (!isValidEmail(emailMessage.receiver)) {
+ throw new ForbiddenException("Invalid email address.")
+ }
+
+ val adminUsers = userDao.fetchByRole(UserRoleEnum.ADMIN)
+ val adminUserIterator = adminUsers.iterator()
+
+ while (adminUserIterator.hasNext) {
+ val admin = adminUserIterator.next()
+ val adminEmail = admin.getEmail
+
+ try {
+ sendEmail(
+ userRegistrationNotification(
+ receiverEmail = adminEmail,
+ userEmail = Some(emailMessage.receiver),
+ affiliation = emailMessage.affiliation,
+ reason = emailMessage.reason,
+ toAdmin = true
+ ),
+ adminEmail
+ )
+ } catch {
+ case ex: Exception =>
+ logger.warn(s"Failed to send email to admin: $adminEmail. Error: ${ex.getMessage}")
+ }
+ }
+
+ try {
+ sendEmail(
+ userRegistrationNotification(
+ receiverEmail = emailMessage.receiver,
+ userEmail = None,
+ affiliation = None,
+ reason = None,
+ toAdmin = false
+ ),
+ emailMessage.receiver
+ )
+ } catch {
+ case ex: Exception =>
+ logger.warn(
+ s"Failed to send notification to user: ${emailMessage.receiver}. Error: ${ex.getMessage}"
+ )
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/HealthCheckResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/HealthCheckResource.scala
new file mode 100644
index 00000000000..0a182ecd4c8
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/HealthCheckResource.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import javax.ws.rs.core.MediaType
+import javax.ws.rs.{GET, Path, Produces}
+
+@Path("/healthcheck")
+@Produces(Array(MediaType.APPLICATION_JSON))
+class HealthCheckResource {
+ @GET
+ def healthCheck: Map[String, String] = Map("status" -> "ok")
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/MockKillWorkerResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/MockKillWorkerResource.scala
new file mode 100644
index 00000000000..d51556f53e8
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/MockKillWorkerResource.scala
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+//package org.apache.texera.web.resource
+//
+//import .KillAndRecover
+//import javax.ws.rs.core.MediaType
+//import javax.ws.rs.{POST, Path, Produces}
+//
+//@Path("/kill")
+//@Produces(Array(MediaType.APPLICATION_JSON))
+//class MockKillWorkerResource() {
+//
+// @POST
+// @Path("/worker") def mockKillWorker: Unit = {
+// WorkflowWebsocketResource.sessionJobs.foreach(p => {
+// val controller = p._2._2
+// Thread.sleep(1500)
+// controller ! KillAndRecover
+// })
+// }
+//
+//}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/SuccessExecutionResult.scala b/amber/src/main/scala/org/apache/texera/web/resource/SuccessExecutionResult.scala
new file mode 100644
index 00000000000..2318d054cab
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/SuccessExecutionResult.scala
@@ -0,0 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+case class SuccessExecutionResult(
+ resultID: String,
+ code: Integer = 0,
+ result: List[String] = List()
+)
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/SyncExecutionResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/SyncExecutionResource.scala
new file mode 100644
index 00000000000..d3047db5802
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/SyncExecutionResource.scala
@@ -0,0 +1,915 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import com.fasterxml.jackson.databind.node.ObjectNode
+import com.typesafe.scalalogging.LazyLogging
+import io.dropwizard.auth.Auth
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.operator.LogicalOp
+import org.apache.texera.amber.core.storage.model.VirtualDocument
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.{
+ ExecutionIdentity,
+ OperatorIdentity,
+ WorkflowIdentity
+}
+import org.apache.texera.amber.core.workflow.{PortIdentity, WorkflowContext, WorkflowSettings}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ ConsoleMessage,
+ ConsoleMessageType
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState._
+import org.apache.texera.amber.engine.common.executionruntimestate.{
+ ExecutionConsoleStore,
+ ExecutionMetadataStore,
+ ExecutionStatsStore
+}
+import io.reactivex.rxjava3.core.Observable
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables.OPERATOR_EXECUTIONS
+import org.apache.texera.web.model.websocket.request.{LogicalPlanPojo, WorkflowExecuteRequest}
+import org.apache.texera.workflow.{LogicalLink, WorkflowCompiler}
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+import org.apache.texera.web.service.{ExecutionResultService, WorkflowService}
+import org.apache.texera.web.storage.ExecutionStateStore.updateWorkflowState
+
+import java.net.URI
+import java.util.concurrent.TimeUnit
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+import com.fasterxml.jackson.databind.ObjectMapper
+
+case class SyncExecutionRequest(
+ executionName: String,
+ logicalPlan: LogicalPlanPojo,
+ workflowSettings: Option[WorkflowSettings],
+ targetOperatorIds: List[String],
+ timeoutSeconds: Int,
+ maxOperatorResultCharLimit: Int,
+ maxOperatorResultCellCharLimit: Int
+)
+
+case class ConsoleMessageInfo(
+ msgType: String,
+ title: String,
+ message: String
+)
+
+case class PortShape(
+ portIndex: Int,
+ rows: Long
+)
+
+case class OperatorInfo(
+ state: String,
+ inputTuples: Long,
+ outputTuples: Long,
+ inputPortShapes: Option[List[PortShape]],
+ resultMode: String, // "table" or "visualization"
+ result: Option[Any], // JSON array (List[ObjectNode])
+ totalRowCount: Option[Int],
+ displayedRows: Option[Int],
+ truncated: Option[Boolean],
+ consoleLogs: Option[List[ConsoleMessageInfo]],
+ error: Option[String],
+ warnings: Option[List[String]]
+)
+
+case class SyncExecutionResult(
+ success: Boolean,
+ state: String,
+ operators: Map[String, OperatorInfo],
+ compilationErrors: Option[Map[String, String]],
+ errors: Option[List[String]]
+)
+
+sealed trait TerminationReason
+case class TerminalStateReached(state: ExecutionMetadataStore) extends TerminationReason
+case class ConsoleErrorDetected(consoleState: ExecutionConsoleStore) extends TerminationReason
+case class TargetResultsReady(statsState: ExecutionStatsStore) extends TerminationReason
+
+@Path("/execution")
+@Consumes(Array(MediaType.APPLICATION_JSON))
+@Produces(Array(MediaType.APPLICATION_JSON))
+class SyncExecutionResource extends LazyLogging {
+
+ // Hard caps applied regardless of request — guard against runaway payloads.
+ private val MAX_OPERATOR_RESULT_CHARS = 100000
+ private val MAX_OPERATOR_RESULT_CELL_CHARS = 20000
+
+ @POST
+ @Path("/{wid}/{cuid}/run")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def executeWorkflowSync(
+ @PathParam("wid") workflowId: Long,
+ @PathParam("cuid") computingUnitId: Int,
+ request: SyncExecutionRequest,
+ @Auth user: SessionUser
+ ): SyncExecutionResult = {
+ val timeoutSeconds = request.timeoutSeconds
+
+ val maxOperatorResultCharLimit =
+ Math.min(request.maxOperatorResultCharLimit, MAX_OPERATOR_RESULT_CHARS)
+ val maxOperatorResultCellCharLimit =
+ Math.min(request.maxOperatorResultCellCharLimit, MAX_OPERATOR_RESULT_CELL_CHARS)
+
+ logger.info(
+ s"Starting sync execution for workflow $workflowId with limits: " +
+ s"maxOperatorResultCharLimit=${request.maxOperatorResultCharLimit} (capped to $maxOperatorResultCharLimit), " +
+ s"maxOperatorResultCellCharLimit=${request.maxOperatorResultCellCharLimit} (capped to $maxOperatorResultCellCharLimit)"
+ )
+
+ try {
+ val workflowService = WorkflowService.getOrCreate(
+ WorkflowIdentity(workflowId),
+ computingUnitId
+ )
+
+ shutdownPreviousExecution(workflowService)
+
+ // "Execute To" semantics: when a single target is given, run only its upstream sub-DAG.
+ val effectiveLogicalPlan =
+ computeSubDAGIfNeeded(request.logicalPlan, request.targetOperatorIds)
+
+ val executeRequest = WorkflowExecuteRequest(
+ executionName = request.executionName,
+ engineVersion = "1.0",
+ logicalPlan = effectiveLogicalPlan,
+ replayFromExecution = None,
+ workflowSettings = request.workflowSettings
+ .getOrElse(
+ WorkflowSettings(dataTransferBatchSize = ApplicationConfig.defaultDataTransferBatchSize)
+ ),
+ emailNotificationEnabled = false,
+ computingUnitId = computingUnitId
+ )
+
+ workflowService.initExecutionService(
+ executeRequest,
+ Some(user.getUser),
+ new URI(s"sync-execution://$workflowId")
+ )
+
+ val executionService = workflowService.executionService.getValue
+ if (executionService == null) {
+ return SyncExecutionResult(
+ success = false,
+ state = "Error",
+ operators = Map.empty,
+ compilationErrors = None,
+ errors = Some(List("Failed to initialize execution service"))
+ )
+ }
+
+ // Snapshot before subscribing — handles the race where a fast execution finishes
+ // before the Observable below sees any state change.
+ val currentState = executionService.executionStateStore.metadataStore.getState
+ val currentConsoleState = executionService.executionStateStore.consoleStore.getState
+ val currentStatsState = executionService.executionStateStore.statsStore.getState
+
+ // Multi-region operators (e.g., HashJoin: build region then probe region) report their
+ // aggregated logical state as COMPLETED for a brief window after the first region
+ // terminates and before the second region's workers are added to regionExecutions.
+ // Guard against firing during that window by also requiring every declared external
+ // input port to be present in the operator's input metrics — port-1 stats only appear
+ // once probe actually starts consuming, which closes the race.
+ val targetExpectedExternalInputs: Map[String, Int] = effectiveLogicalPlan.operators
+ .filter(op => request.targetOperatorIds.contains(op.operatorIdentifier.id))
+ .map(op => op.operatorIdentifier.id -> op.operatorInfo.inputPorts.count(!_.id.internal))
+ .toMap
+
+ // Require COMPLETED, not just "has output", so upstream operators finish flushing
+ // their data downstream before we tear the execution down.
+ def allTargetsCompleted(stats: ExecutionStatsStore): Boolean = {
+ request.targetOperatorIds.nonEmpty && request.targetOperatorIds.forall { opId =>
+ stats.operatorInfo.get(opId).exists { metrics =>
+ val externalInputPortsReporting =
+ metrics.operatorStatistics.inputMetrics.count(!_.portId.internal)
+ val expectedExternalInputs = targetExpectedExternalInputs.getOrElse(opId, 0)
+ metrics.operatorState == COMPLETED &&
+ externalInputPortsReporting >= expectedExternalInputs
+ }
+ }
+ }
+
+ val terminationReason: TerminationReason =
+ if (isTerminalState(currentState.state)) {
+ TerminalStateReached(currentState)
+ } else if (hasConsoleError(currentConsoleState)) {
+ ConsoleErrorDetected(currentConsoleState)
+ } else if (allTargetsCompleted(currentStatsState)) {
+ TargetResultsReady(currentStatsState)
+ } else {
+ val terminalStateObservable: Observable[TerminationReason] =
+ executionService.executionStateStore.metadataStore.getStateObservable
+ .filter((state: ExecutionMetadataStore) => isTerminalState(state.state))
+ .map[TerminationReason](state => TerminalStateReached(state))
+
+ val consoleErrorObservable: Observable[TerminationReason] =
+ executionService.executionStateStore.consoleStore.getStateObservable
+ .filter((consoleState: ExecutionConsoleStore) => hasConsoleError(consoleState))
+ .map[TerminationReason](consoleState => ConsoleErrorDetected(consoleState))
+
+ val targetResultsObservable: Observable[TerminationReason] =
+ executionService.executionStateStore.statsStore.getStateObservable
+ .filter((stats: ExecutionStatsStore) => allTargetsCompleted(stats))
+ .map[TerminationReason](stats => TargetResultsReady(stats))
+
+ try {
+ Observable
+ .amb(
+ java.util.Arrays.asList(
+ terminalStateObservable,
+ consoleErrorObservable,
+ targetResultsObservable
+ )
+ )
+ .firstOrError()
+ .timeout(timeoutSeconds.toLong, TimeUnit.SECONDS)
+ .blockingGet()
+ } catch {
+ case _: java.util.concurrent.TimeoutException =>
+ killExecution(executionService)
+ return SyncExecutionResult(
+ success = false,
+ state = "Killed",
+ operators = Map.empty,
+ compilationErrors = None,
+ errors = Some(List(s"Timeout after $timeoutSeconds seconds"))
+ )
+ case e: Exception =>
+ logger.error(s"Error waiting for execution: ${e.getMessage}", e)
+ return SyncExecutionResult(
+ success = false,
+ state = "Error",
+ operators = Map.empty,
+ compilationErrors = None,
+ errors = Some(List(e.getMessage))
+ )
+ }
+ }
+
+ val (finalState, terminatedByConsoleError, terminatedByTargetResults) =
+ terminationReason match {
+ case TerminalStateReached(state) =>
+ (state, false, false)
+ case ConsoleErrorDetected(_) =>
+ killExecution(executionService)
+ (executionService.executionStateStore.metadataStore.getState, true, false)
+ case TargetResultsReady(_) =>
+ // RegionExecutionCoordinator caches upstream results asynchronously after operators
+ // complete; sleep gives that caching a chance to finish before we shut down the client.
+ // TODO: replace with a synchronous signal from the engine.
+ Thread.sleep(500)
+ killExecution(executionService)
+ // Override to COMPLETED — we have everything we asked for, even though the engine
+ // sees this as a kill.
+ executionService.executionStateStore.metadataStore.updateState(metadataStore =>
+ updateWorkflowState(COMPLETED, metadataStore)
+ )
+ (executionService.executionStateStore.metadataStore.getState, false, true)
+ }
+
+ // Let the result writer flush before we read storage.
+ Thread.sleep(500)
+
+ // Console DB writes lag the in-memory store; pass the latter so error extraction
+ // can fall back when the row hasn't landed yet.
+ val inMemoryConsoleState = terminationReason match {
+ case ConsoleErrorDetected(consoleState) => Some(consoleState)
+ case _ => None
+ }
+
+ val executionId = executionService.workflowContext.executionId
+ val operatorInfos = collectOperatorInfos(
+ executionId,
+ executionService,
+ request.targetOperatorIds,
+ maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit,
+ inMemoryConsoleState
+ )
+
+ val fatalErrors = finalState.fatalErrors
+ .map(err => s"${err.`type`}: ${err.message}")
+ .toList
+
+ val hasOperatorConsoleError = operatorInfos.values.exists(_.error.isDefined)
+
+ val stateString =
+ if (terminatedByConsoleError) "Failed"
+ else if (terminatedByTargetResults) "Completed"
+ else stateToString(finalState.state)
+
+ val isSuccess = (finalState.state == COMPLETED || terminatedByTargetResults) &&
+ !hasOperatorConsoleError && !terminatedByConsoleError
+
+ SyncExecutionResult(
+ success = isSuccess,
+ state = stateString,
+ operators = operatorInfos,
+ compilationErrors = None,
+ errors = if (fatalErrors.nonEmpty) Some(fatalErrors) else None
+ )
+
+ } catch {
+ case e: Exception =>
+ logger.error(s"Sync execution error: ${e.getMessage}", e)
+ handleExecutionError(e)
+ }
+ }
+
+ private def shutdownPreviousExecution(workflowService: WorkflowService): Unit = {
+ try {
+ val previousEs = workflowService.executionService.getValue
+ if (previousEs != null && previousEs.client != null) {
+ logger.info(s"Shutting down previous execution client")
+ previousEs.client.shutdown()
+ }
+ } catch {
+ case e: Exception =>
+ logger.warn(s"Error shutting down previous execution client: ${e.getMessage}")
+ }
+ }
+
+ private def killExecution(
+ executionService: org.apache.texera.web.service.WorkflowExecutionService
+ ): Unit = {
+ try {
+ if (executionService.client != null) {
+ executionService.client.shutdown()
+ }
+ executionService.executionStateStore.statsStore.updateState(stats =>
+ stats.withEndTimeStamp(System.currentTimeMillis())
+ )
+ executionService.executionStateStore.metadataStore.updateState(metadataStore =>
+ updateWorkflowState(KILLED, metadataStore)
+ )
+ } catch {
+ case e: Exception =>
+ logger.warn(s"Error killing execution: ${e.getMessage}")
+ }
+ }
+
+ private def collectOperatorInfos(
+ executionId: ExecutionIdentity,
+ executionService: org.apache.texera.web.service.WorkflowExecutionService,
+ targetOperatorIds: List[String],
+ maxOperatorResultCharLimit: Int,
+ maxOperatorResultCellCharLimit: Int,
+ inMemoryConsoleState: Option[ExecutionConsoleStore] = None
+ ): Map[String, OperatorInfo] = {
+ val operatorInfos = mutable.Map[String, OperatorInfo]()
+
+ val statsState = executionService.executionStateStore.statsStore.getState
+ val operatorStats = statsState.operatorInfo
+
+ val baseTargetOps = if (targetOperatorIds.nonEmpty) {
+ targetOperatorIds
+ } else {
+ operatorStats.keys.toList
+ }
+
+ // Pull in any operator that logged a console error even if it isn't a target —
+ // otherwise the caller can't see why an upstream op failed.
+ val consoleErrorOps = inMemoryConsoleState
+ .map { consoleState =>
+ consoleState.operatorConsole.keys.toList
+ }
+ .getOrElse(List.empty)
+
+ val targetOps = (baseTargetOps ++ consoleErrorOps).distinct
+
+ for (opId <- targetOps) {
+ val stats = operatorStats.get(opId)
+ val (state, inputTuples, outputTuples): (String, Long, Long) = stats match {
+ case Some(s) =>
+ val inputCount = s.operatorStatistics.inputMetrics.map(_.tupleMetrics.count).sum
+ val outputCount = s.operatorStatistics.outputMetrics.map(_.tupleMetrics.count).sum
+ (stateToString(s.operatorState), inputCount, outputCount)
+ case None => ("Unknown", 0L, 0L)
+ }
+
+ val inputPortShapes: Option[List[PortShape]] = stats
+ .map { s =>
+ s.operatorStatistics.inputMetrics.map { pm =>
+ PortShape(pm.portId.id, pm.tupleMetrics.count)
+ }.toList
+ }
+ .filter(_.nonEmpty)
+
+ val (resultMode, result, totalRowCount, displayedRows, truncated) =
+ collectOperatorResult(
+ executionId,
+ opId,
+ maxOperatorResultCharLimit,
+ maxOperatorResultCellCharLimit
+ )
+
+ // DB is authoritative once written; fall back to in-memory state for in-flight runs
+ // where the console row hasn't been persisted yet.
+ val dbConsoleLogs = collectConsoleLogs(executionId, opId)
+ val consoleLogs = dbConsoleLogs.orElse {
+ inMemoryConsoleState.flatMap { consoleState =>
+ consoleState.operatorConsole
+ .get(opId)
+ .map { opConsole =>
+ opConsole.consoleMessages.map { msg =>
+ ConsoleMessageInfo(
+ msgType = msg.msgType.name,
+ title = msg.title,
+ message = msg.message
+ )
+ }.toList
+ }
+ .filter(_.nonEmpty)
+ }
+ }
+
+ // Python writes the full error text to `message`; Scala writes it to `title`
+ // (with a stack trace in `message`). Pick whichever is longer to avoid losing detail.
+ val errorMsg = consoleLogs.flatMap(
+ _.find(_.msgType == "ERROR").map { e =>
+ if (e.message.nonEmpty && e.message.length > e.title.length) e.message
+ else e.title
+ }
+ )
+
+ // Convention: PRINT messages prefixed with "WARNING: " surface as warnings.
+ val warningMsgs = consoleLogs
+ .map(_.filter(_.title.startsWith("WARNING: ")).map(_.title))
+ .filter(_.nonEmpty)
+
+ operatorInfos(opId) = OperatorInfo(
+ state = state,
+ inputTuples = inputTuples,
+ outputTuples = outputTuples,
+ inputPortShapes = inputPortShapes,
+ resultMode = resultMode,
+ result = result,
+ totalRowCount = totalRowCount,
+ displayedRows = displayedRows,
+ truncated = truncated,
+ consoleLogs = consoleLogs,
+ error = errorMsg,
+ warnings = warningMsgs
+ )
+ }
+
+ operatorInfos.toMap
+ }
+
+ private def handleExecutionError(e: Exception): SyncExecutionResult = {
+ val errorMsg = e.getMessage
+ val isCompilationError = errorMsg != null && (
+ errorMsg.contains("compilation") ||
+ errorMsg.contains("Compilation") ||
+ errorMsg.contains("operator") ||
+ errorMsg.contains("schema")
+ )
+
+ if (isCompilationError) {
+ SyncExecutionResult(
+ success = false,
+ state = "CompilationFailed",
+ operators = Map.empty,
+ compilationErrors = Some(Map("error" -> errorMsg)),
+ errors = Some(List(errorMsg))
+ )
+ } else {
+ SyncExecutionResult(
+ success = false,
+ state = "Error",
+ operators = Map.empty,
+ compilationErrors = None,
+ errors = Some(List(Option(e.getMessage).getOrElse("Unknown error")))
+ )
+ }
+ }
+
+ /**
+ * Symmetric truncation: fill half the char budget from the front of the result, keep a
+ * sliding-window of the most recent tuples for the back half. Returns a JSON array;
+ * serialization to table/toon format happens in agent-service.
+ */
+ private def collectOperatorResult(
+ executionId: ExecutionIdentity,
+ opId: String,
+ maxOperatorResultCharLimit: Int,
+ maxOperatorResultCellCharLimit: Int
+ ): (String, Option[Any], Option[Int], Option[Int], Option[Boolean]) = {
+ import com.fasterxml.jackson.databind.node.ObjectNode
+
+ try {
+ val storageUriOption = WorkflowExecutionsResource.getResultUriByLogicalPortId(
+ executionId,
+ OperatorIdentity(opId),
+ PortIdentity()
+ )
+
+ storageUriOption match {
+ case Some(storageUri) =>
+ val document = DocumentFactory
+ .openDocument(storageUri)
+ ._1
+ .asInstanceOf[VirtualDocument[Tuple]]
+
+ val totalCount = document.getCount.toInt
+ val mapper = new ObjectMapper()
+ val tupleIterator = document.get()
+
+ if (totalCount == 0 || !tupleIterator.hasNext) {
+ return (
+ "table",
+ Some(List.empty[ObjectNode].asJava),
+ Some(0),
+ Some(0),
+ Some(false)
+ )
+ }
+
+ // A single tuple with html-content / json-content is a visualization payload —
+ // the frontend renders it as an iframe rather than a table.
+ val firstTuple = tupleIterator.next()
+ if (totalCount == 1 && isVisualizationTuple(firstTuple)) {
+ val jsonResults =
+ ExecutionResultService.convertTuplesToJson(List(firstTuple), isVisualization = true)
+ jsonResults.foreach(
+ _.asInstanceOf[ObjectNode].put("__is_visualization__", true)
+ )
+ return (
+ "visualization",
+ Some(jsonResults),
+ Some(totalCount),
+ Some(1),
+ Some(false)
+ )
+ }
+
+ // __row_index__ preserves the original position so the frontend can show
+ // "row N" correctly after symmetric truncation drops the middle.
+ var rowIndex = 0
+ val firstJson = ExecutionResultService.convertTuplesToJson(List(firstTuple)).head
+ val truncatedFirst = truncateSingleTuple(firstJson, maxOperatorResultCellCharLimit)
+ truncatedFirst.put("__row_index__", rowIndex)
+ val firstSize = estimateTupleSize(truncatedFirst, mapper)
+
+ if (firstSize >= maxOperatorResultCharLimit) {
+ return (
+ "table",
+ Some(List(truncatedFirst).asJava),
+ Some(totalCount),
+ Some(1),
+ Some(true)
+ )
+ }
+
+ val halfLimit = maxOperatorResultCharLimit / 2
+ val truncationNoticeSize = 50 // reserved for the "...skipped..." marker
+
+ val frontTuples = mutable.ListBuffer[ObjectNode](truncatedFirst)
+ var frontSize = firstSize
+ var processedCount = 1
+
+ while (tupleIterator.hasNext && frontSize < halfLimit) {
+ val tuple = tupleIterator.next()
+ rowIndex += 1
+ processedCount += 1
+ val jsonTuple = ExecutionResultService.convertTuplesToJson(List(tuple)).head
+ val truncatedTuple = truncateSingleTuple(jsonTuple, maxOperatorResultCellCharLimit)
+ truncatedTuple.put("__row_index__", rowIndex)
+ val tupleSize = estimateTupleSize(truncatedTuple, mapper)
+
+ if (frontSize + tupleSize <= halfLimit) {
+ frontTuples += truncatedTuple
+ frontSize += tupleSize
+ } else {
+ // Front is full — switch to a sliding window for the back half.
+ val backBuffer = mutable.ArrayBuffer[(ObjectNode, Int)]()
+ backBuffer += ((truncatedTuple, tupleSize))
+ var backSize = tupleSize
+
+ while (tupleIterator.hasNext) {
+ val t = tupleIterator.next()
+ rowIndex += 1
+ processedCount += 1
+ val jt = ExecutionResultService.convertTuplesToJson(List(t)).head
+ val tt = truncateSingleTuple(jt, maxOperatorResultCellCharLimit)
+ tt.put("__row_index__", rowIndex)
+ val ts = estimateTupleSize(tt, mapper)
+
+ backBuffer += ((tt, ts))
+ backSize += ts
+
+ while (backSize > halfLimit - truncationNoticeSize && backBuffer.size > 1) {
+ val (_, removedSize) = backBuffer.remove(0)
+ backSize -= removedSize
+ }
+ }
+
+ val backTuples = backBuffer.map(_._1).toList
+ val allTuples = frontTuples.toList ++ backTuples
+ val skippedRows = totalCount - allTuples.size
+
+ return (
+ "table",
+ Some(allTuples.asJava),
+ Some(totalCount),
+ Some(allTuples.size),
+ Some(skippedRows > 0)
+ )
+ }
+ }
+
+ if (tupleIterator.hasNext) {
+ val backBuffer = mutable.ArrayBuffer[(ObjectNode, Int)]()
+ var backSize = 0
+
+ while (tupleIterator.hasNext) {
+ val t = tupleIterator.next()
+ rowIndex += 1
+ processedCount += 1
+ val jt = ExecutionResultService.convertTuplesToJson(List(t)).head
+ val tt = truncateSingleTuple(jt, maxOperatorResultCellCharLimit)
+ tt.put("__row_index__", rowIndex)
+ val ts = estimateTupleSize(tt, mapper)
+
+ backBuffer += ((tt, ts))
+ backSize += ts
+
+ while (backSize > halfLimit - truncationNoticeSize && backBuffer.size > 1) {
+ val (_, removedSize) = backBuffer.remove(0)
+ backSize -= removedSize
+ }
+ }
+
+ val backTuples = backBuffer.map(_._1).toList
+ val allTuples = frontTuples.toList ++ backTuples
+ val skippedRows = totalCount - allTuples.size
+
+ (
+ "table",
+ Some(allTuples.asJava),
+ Some(totalCount),
+ Some(allTuples.size),
+ Some(skippedRows > 0)
+ )
+ } else {
+ (
+ "table",
+ Some(frontTuples.toList.asJava),
+ Some(totalCount),
+ Some(frontTuples.size),
+ Some(false)
+ )
+ }
+
+ case None =>
+ ("table", None, None, None, None)
+ }
+ } catch {
+ case e: Exception =>
+ logger.warn(s"Error collecting result for operator $opId: ${e.getMessage}", e)
+ ("table", None, None, None, None)
+ }
+ }
+
+ private def truncateSingleTuple(
+ tuple: ObjectNode,
+ maxCellChars: Int
+ ): ObjectNode = {
+ import com.fasterxml.jackson.databind.ObjectMapper
+ import com.fasterxml.jackson.databind.node.TextNode
+
+ val mapper = new ObjectMapper()
+ val truncatedTuple = mapper.createObjectNode()
+ val fieldNames = tuple.fieldNames()
+
+ while (fieldNames.hasNext) {
+ val fieldName = fieldNames.next()
+ val fieldValue = tuple.get(fieldName)
+ if (fieldValue.isTextual) {
+ val text = fieldValue.asText()
+ if (text.length > maxCellChars) {
+ val truncatedText = symmetricTruncateCellValue(text, maxCellChars)
+ truncatedTuple.set(fieldName, new TextNode(truncatedText))
+ } else {
+ truncatedTuple.set(fieldName, fieldValue)
+ }
+ } else {
+ truncatedTuple.set(fieldName, fieldValue)
+ }
+ }
+ truncatedTuple
+ }
+
+ private def estimateTupleSize(
+ tuple: ObjectNode,
+ mapper: ObjectMapper
+ ): Int = {
+ mapper.writeValueAsString(tuple).length + 1 // +1 for the array separator
+ }
+
+ private def symmetricTruncateCellValue(text: String, maxChars: Int): String = {
+ if (text.length <= maxChars) {
+ text
+ } else {
+ val notice = "...[truncated]..."
+ val availableChars = maxChars - notice.length
+ if (availableChars <= 0) {
+ text.substring(0, maxChars)
+ } else {
+ val halfChars = availableChars / 2
+ text.substring(0, halfChars) + notice + text.substring(text.length - halfChars)
+ }
+ }
+ }
+
+ private def isVisualizationTuple(tuple: Tuple): Boolean = {
+ try {
+ val schema = tuple.getSchema
+ val fieldNames = schema.getAttributes.map(_.getName)
+ fieldNames.exists(name => name == "html-content" || name == "json-content")
+ } catch {
+ case _: Exception => false
+ }
+ }
+
+ private def collectConsoleLogs(
+ executionId: ExecutionIdentity,
+ opId: String
+ ): Option[List[ConsoleMessageInfo]] = {
+ try {
+ val uriOption = getConsoleMessageUri(executionId, OperatorIdentity(opId))
+
+ uriOption.flatMap { uri =>
+ val document = DocumentFactory
+ .openDocument(uri)
+ ._1
+ .asInstanceOf[VirtualDocument[Tuple]]
+
+ val messages = document.get().toList.flatMap { tuple =>
+ try {
+ val protoString = tuple.getField[String](0)
+ val msg = ConsoleMessage.fromAscii(protoString)
+ Some(
+ ConsoleMessageInfo(
+ msgType = msg.msgType.name,
+ title = msg.title,
+ message = msg.message
+ )
+ )
+ } catch {
+ case _: Exception => None
+ }
+ }
+
+ if (messages.nonEmpty) Some(messages) else None
+ }
+ } catch {
+ case _: Exception => None
+ }
+ }
+
+ private def getConsoleMessageUri(
+ eid: ExecutionIdentity,
+ opId: OperatorIdentity
+ ): Option[URI] = {
+ val context = SqlServer.getInstance().createDSLContext()
+ Option(
+ context
+ .select(OPERATOR_EXECUTIONS.CONSOLE_MESSAGES_URI)
+ .from(OPERATOR_EXECUTIONS)
+ .where(OPERATOR_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .and(OPERATOR_EXECUTIONS.OPERATOR_ID.eq(opId.id))
+ .fetchOneInto(classOf[String])
+ ).filter(uri => uri != null && uri.nonEmpty)
+ .map(s => URI.create(s))
+ }
+
+ private def isTerminalState(state: WorkflowAggregatedState): Boolean = {
+ state match {
+ case COMPLETED | FAILED | KILLED | TERMINATED => true
+ case _ => false
+ }
+ }
+
+ private def hasConsoleError(consoleState: ExecutionConsoleStore): Boolean = {
+ consoleState.operatorConsole.values.exists { opConsole =>
+ opConsole.consoleMessages.exists(_.msgType == ConsoleMessageType.ERROR)
+ }
+ }
+
+ private def stateToString(state: WorkflowAggregatedState): String = {
+ state match {
+ case UNINITIALIZED => "Uninitialized"
+ case READY => "Ready"
+ case RUNNING => "Running"
+ case PAUSING => "Pausing"
+ case PAUSED => "Paused"
+ case RESUMING => "Resuming"
+ case COMPLETED => "Completed"
+ case FAILED => "Failed"
+ case KILLED => "Killed"
+ case TERMINATED => "Terminated"
+ case _ => "Unknown"
+ }
+ }
+
+ private def computeSubDAGIfNeeded(
+ logicalPlan: LogicalPlanPojo,
+ targetOperatorIds: List[String]
+ ): LogicalPlanPojo = {
+ if (targetOperatorIds.length != 1) {
+ return logicalPlan
+ }
+
+ val targetOpId = targetOperatorIds.head
+ val operatorMap: Map[String, LogicalOp] =
+ logicalPlan.operators.map(op => op.operatorIdentifier.id -> op).toMap
+
+ if (!operatorMap.contains(targetOpId)) {
+ logger.warn(s"Target operator $targetOpId not found in logical plan, using full DAG")
+ return logicalPlan
+ }
+
+ val incomingLinks: Map[String, List[LogicalLink]] =
+ logicalPlan.links.groupBy(_.toOpId.id)
+
+ val visited = mutable.Set[String]()
+ val subDagOperators = mutable.ListBuffer[LogicalOp]()
+ val subDagLinks = mutable.ListBuffer[LogicalLink]()
+
+ def dfs(currentOpId: String): Unit = {
+ if (visited.contains(currentOpId)) return
+ visited.add(currentOpId)
+
+ operatorMap.get(currentOpId).foreach { op =>
+ subDagOperators += op
+ incomingLinks.getOrElse(currentOpId, List.empty).foreach { link =>
+ subDagLinks += link
+ dfs(link.fromOpId.id)
+ }
+ }
+ }
+
+ dfs(targetOpId)
+
+ LogicalPlanPojo(
+ operators = subDagOperators.toList,
+ links = subDagLinks.toList,
+ opsToViewResult = targetOperatorIds.filter(id => visited.contains(id)),
+ opsToReuseResult = logicalPlan.opsToReuseResult.filter(id => visited.contains(id))
+ )
+ }
+
+ // Returns operator-id -> error message; empty map means compilation succeeded.
+ private def validateWorkflow(
+ workflowId: Long,
+ logicalPlan: LogicalPlanPojo
+ ): Map[String, String] = {
+ try {
+ val tempContext = new WorkflowContext(WorkflowIdentity(workflowId))
+ val compiler = new WorkflowCompiler(tempContext)
+ compiler.compile(logicalPlan)
+ Map.empty
+ } catch {
+ case e: Exception =>
+ val errorMsg = Option(e.getMessage).getOrElse("Compilation failed")
+ val operatorIdPattern = """operator[- ]?(\S+)""".r
+ val operatorId = operatorIdPattern
+ .findFirstMatchIn(errorMsg.toLowerCase)
+ .map(_.group(1))
+ .getOrElse("workflow")
+ Map(operatorId -> errorMsg)
+ }
+ }
+
+ @GET
+ @Path("/health")
+ def healthCheck: Map[String, String] = Map("status" -> "ok")
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/SystemMetadataResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/SystemMetadataResource.scala
new file mode 100644
index 00000000000..a17cb736a8c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/SystemMetadataResource.scala
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import org.apache.texera.amber.operator.metadata.{AllOperatorMetadata, OperatorMetadataGenerator}
+
+import javax.ws.rs.core.MediaType
+import javax.ws.rs.{GET, Path, Produces}
+
+@Path("/resources")
+@Produces(Array(MediaType.APPLICATION_JSON))
+class SystemMetadataResource {
+
+ @GET
+ @Path("/operator-metadata")
+ def getOperatorMetadata: AllOperatorMetadata = {
+ OperatorMetadataGenerator.allOperatorMetadata
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/UserConfigResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/UserConfigResource.scala
new file mode 100644
index 00000000000..3e36adf9346
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/UserConfigResource.scala
@@ -0,0 +1,182 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables.USER_CONFIG
+import org.apache.texera.dao.jooq.generated.tables.daos.UserConfigDao
+import org.apache.texera.dao.jooq.generated.tables.pojos.{User, UserConfig}
+
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core._
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+/**
+ * This class handles requests to read and write the user dictionary,
+ * an abstract collection of (key, value) string pairs that is unique for each user
+ * This is accomplished using a mysql table called user_dictionary.
+ * The details of user_dictionary can be found in /sql/texera_ddl.sql
+ */
+@Path("/user/config")
+@RolesAllowed(Array("REGULAR", "ADMIN"))
+@Consumes(Array(MediaType.TEXT_PLAIN))
+class UserConfigResource {
+ private def userDictionaryDao =
+ new UserConfigDao(
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ .configuration
+ )
+
+ @GET
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getAllDict(@Auth sessionUser: SessionUser): Map[String, String] = {
+ val user = sessionUser.getUser
+ getDict(user)
+ }
+
+ /**
+ * This method retrieves all of a user's dictionary entries in
+ * the user_dictionary table as a json object
+ */
+ private def getDict(user: User): Map[String, String] = {
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ .select()
+ .from(USER_CONFIG)
+ .where(USER_CONFIG.UID.eq(user.getUid))
+ .fetchInto(classOf[UserConfig])
+ .asScala
+ .map { entry => (entry.getKey, entry.getValue) }
+ .toMap
+ }
+
+ @GET
+ @Produces(Array(MediaType.TEXT_PLAIN))
+ @Path("/{key}")
+ def getEntry(@PathParam("key") key: String, @Auth sessionUser: SessionUser): String = {
+ val user = sessionUser.getUser
+
+ if (key == null || key.trim.isEmpty) {
+ throw new BadRequestException("key cannot be null or empty")
+ }
+ if (!dictEntryExists(user, key)) {
+ null
+ } else {
+ getValueByKey(user, key)
+ }
+ }
+
+ /**
+ * This method retrieves a value from the user_dictionary table
+ * given a user's uid and key. each tuple (uid, key) is a primary key
+ * in user_dictionary, and should uniquely identify one value
+ *
+ * @return String or null if entry doesn't exist
+ */
+ private def getValueByKey(user: User, key: String): String = {
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ .fetchOne(
+ USER_CONFIG,
+ USER_CONFIG.UID.eq(user.getUid).and(USER_CONFIG.KEY.eq(key))
+ )
+ .getValue
+ }
+
+ /**
+ * This method creates or updates an entry in the current in-session user's dictionary based on
+ * the "key" and "value" attributes of the PostRequest
+ */
+ @PUT
+ @Path("/{key}")
+ def setEntry(
+ @PathParam("key") key: String,
+ value: String,
+ @Auth sessionUser: SessionUser
+ ): Unit = {
+ val user = sessionUser.getUser
+ if (key == null || key.trim.isEmpty) {
+ throw new BadRequestException("key cannot be null or empty")
+ }
+ if (dictEntryExists(user, key)) {
+ userDictionaryDao.update(new UserConfig(user.getUid, key, value))
+ } else {
+ userDictionaryDao.insert(new UserConfig(user.getUid, key, value))
+ }
+ }
+
+ /**
+ * This method checks if a given entry exists
+ * each tuple (uid, key) is a primary key in user_dictionary,
+ * and should uniquely identify one value
+ */
+ private def dictEntryExists(user: User, key: String): Boolean = {
+ userDictionaryDao.existsById(
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ .newRecord(USER_CONFIG.UID, USER_CONFIG.KEY)
+ .values(user.getUid, key)
+ )
+ }
+
+ /**
+ * This method deletes a key-value pair from the current in-session user's dictionary based on
+ * the "key" attribute of the DeleteRequest
+ *
+ * @return
+ * 401 unauthorized -
+ * 400 bad request -
+ * 422 Unprocessable Entity - payload: "no such entry" (if no entry exists for provided key)
+ */
+ @DELETE
+ @Path("/{key}")
+ def deleteEntry(@PathParam("key") key: String, @Auth sessionUser: SessionUser): Unit = {
+ val user = sessionUser.getUser
+ if (key == null || key.trim.isEmpty) {
+ throw new BadRequestException("key cannot be null or empty")
+ }
+ if (dictEntryExists(user, key)) {
+ deleteDictEntry(user, key)
+ }
+ }
+
+ /**
+ * This method deletes a single entry
+ * each tuple (uid, key) is a primary key in user_dictionary,
+ * and should uniquely identify one value
+ */
+ private def deleteDictEntry(user: User, key: String): Unit = {
+ userDictionaryDao.deleteById(
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ .newRecord(USER_CONFIG.UID, USER_CONFIG.KEY)
+ .values(user.getUid, key)
+ )
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/WebsocketPayloadSizeTuner.scala b/amber/src/main/scala/org/apache/texera/web/resource/WebsocketPayloadSizeTuner.scala
new file mode 100644
index 00000000000..92c7ad87f00
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/WebsocketPayloadSizeTuner.scala
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import javax.servlet.{ServletContextEvent, ServletContextListener}
+import javax.websocket.server.ServerContainer
+
+class WebsocketPayloadSizeTuner(maxKB: Int) // by default, 64 KiB
+ extends ServletContextListener {
+
+ override def contextInitialized(sce: ServletContextEvent): Unit = {
+ val container = sce.getServletContext
+ .getAttribute(classOf[ServerContainer].getName)
+ .asInstanceOf[ServerContainer]
+
+ container.setDefaultMaxTextMessageBufferSize(maxKB * 1024)
+ container.setDefaultMaxBinaryMessageBufferSize(maxKB * 1024)
+ }
+
+ override def contextDestroyed(sce: ServletContextEvent): Unit = {}
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/WorkflowWebsocketResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/WorkflowWebsocketResource.scala
new file mode 100644
index 00000000000..3a02a9a4c21
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/WorkflowWebsocketResource.scala
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import com.google.protobuf.timestamp.Timestamp
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.clustering.ClusterListener
+import org.apache.texera.amber.core.virtualidentity.WorkflowIdentity
+import org.apache.texera.amber.core.workflowruntimestate.FatalErrorType.COMPILATION_ERROR
+import org.apache.texera.amber.core.workflowruntimestate.WorkflowFatalError
+import org.apache.texera.amber.error.ErrorUtils.getStackTraceWithAllCauses
+import org.apache.texera.amber.util.JSONUtils.objectMapper
+import org.apache.texera.auth.util.HeaderField
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.apache.texera.web.model.websocket.event.{WorkflowErrorEvent, WorkflowStateEvent}
+import org.apache.texera.web.model.websocket.request._
+import org.apache.texera.web.model.websocket.response._
+import org.apache.texera.web.service.WorkflowService
+import org.apache.texera.web.{ServletAwareConfigurator, SessionState}
+
+import java.time.Instant
+import javax.websocket._
+import javax.websocket.server.ServerEndpoint
+import scala.jdk.CollectionConverters.MapHasAsScala
+
+@ServerEndpoint(
+ value = "/wsapi/workflow-websocket",
+ configurator = classOf[ServletAwareConfigurator]
+)
+class WorkflowWebsocketResource extends LazyLogging {
+
+ @OnOpen
+ def myOnOpen(session: Session, config: EndpointConfig): Unit = {
+ val sessionState = new SessionState(session)
+ SessionState.setState(session.getId, sessionState)
+ val wid = session.getRequestParameterMap.get("wid").get(0).toLong
+ val cuid = session.getRequestParameterMap.get("cuid").get(0).toInt
+ val cuAccessEnum: PrivilegeEnum = PrivilegeEnum.valueOf(
+ session.getUserProperties
+ .get(HeaderField.UserComputingUnitAccess)
+ .asInstanceOf[String]
+ )
+
+ sessionState.setUserComputingUnitAccess(cuAccessEnum)
+ logger.info(
+ s"Websocket connection opened for workflow $wid with computing unit $cuid and access $cuAccessEnum"
+ )
+ // hack to refresh frontend run button state
+ sessionState.send(WorkflowStateEvent("Uninitialized"))
+ val workflowState =
+ WorkflowService.getOrCreate(WorkflowIdentity(wid), cuid)
+ sessionState.subscribe(workflowState)
+ sessionState.send(ClusterStatusUpdateEvent(ClusterListener.numWorkerNodesInCluster))
+ }
+
+ @OnClose
+ def myOnClose(session: Session, cr: CloseReason): Unit = {
+ SessionState.removeState(session.getId)
+ }
+
+ @OnMessage
+ def myOnMsg(session: Session, message: String): Unit = {
+ val request = objectMapper.readValue(message, classOf[TexeraWebSocketRequest])
+ val userOpt = session.getUserProperties.asScala
+ .get(classOf[User].getName)
+ .map(_.asInstanceOf[User])
+ val uidOpt = userOpt.map(_.getUid)
+
+ val sessionState = SessionState.getState(session.getId)
+ val workflowStateOpt = sessionState.getCurrentWorkflowState
+ val executionStateOpt = workflowStateOpt.flatMap(x => Option(x.executionService.getValue))
+ try {
+ request match {
+ case heartbeat: HeartBeatRequest =>
+ sessionState.send(HeartBeatResponse())
+ case paginationRequest: ResultPaginationRequest =>
+ workflowStateOpt.foreach(state =>
+ sessionState.send(state.resultService.handleResultPagination(paginationRequest))
+ )
+ case modifyLogicRequest: ModifyLogicRequest =>
+ if (workflowStateOpt.isDefined) {
+ val executionService = workflowStateOpt.get.executionService.getValue
+ val modifyLogicResponse =
+ executionService.executionReconfigurationService.modifyOperatorLogic(
+ modifyLogicRequest
+ )
+ sessionState.send(modifyLogicResponse)
+ }
+ case workflowExecuteRequest: WorkflowExecuteRequest =>
+ if (sessionState.getUserComputingUnitAccess != PrivilegeEnum.WRITE) {
+ throw new IllegalStateException("User does not have write access to the computing unit")
+ }
+ workflowStateOpt match {
+ case Some(workflow) =>
+ sessionState.send(WorkflowStateEvent("Initializing"))
+ synchronized {
+ workflow.initExecutionService(
+ workflowExecuteRequest,
+ userOpt,
+ session.getRequestURI
+ )
+ }
+ case None => throw new IllegalStateException("workflow is not initialized")
+ }
+ case other =>
+ workflowStateOpt.map(_.executionService.getValue) match {
+ case Some(value) => value.wsInput.onNext(other, uidOpt)
+ case None => throw new IllegalStateException("workflow execution is not initialized")
+ }
+ }
+ } catch {
+ case err: Exception =>
+ logger.error("error occurred in websocket", err)
+ val errEvt = WorkflowFatalError(
+ COMPILATION_ERROR,
+ Timestamp(Instant.now),
+ err.toString,
+ getStackTraceWithAllCauses(err),
+ "unknown operator"
+ )
+ if (executionStateOpt.isDefined) {
+ executionStateOpt.get.executionStateStore.metadataStore.updateState { metadataStore =>
+ metadataStore
+ .withFatalErrors(metadataStore.fatalErrors.filter(e => e.`type` != COMPILATION_ERROR))
+ .addFatalErrors(errEvt)
+ }
+ } else {
+ sessionState.send(
+ WorkflowErrorEvent(
+ Seq(errEvt)
+ )
+ )
+ }
+ throw err
+ }
+
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/AiAssistantManager.scala b/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/AiAssistantManager.scala
new file mode 100644
index 00000000000..374525ff3ff
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/AiAssistantManager.scala
@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.aiassistant
+
+import com.typesafe.config.Config
+import org.apache.texera.amber.config.ApplicationConfig
+
+import java.net.{HttpURLConnection, URL}
+
+object AiAssistantManager {
+ // Optionally retrieve the configuration
+ private val aiAssistantConfigOpt: Option[Config] = ApplicationConfig.aiAssistantConfig
+ private val noAssistant: String = "NoAiAssistant"
+ // Public variables, accessible from outside the object
+ var accountKey: String = _
+ var sharedUrl: String = _
+
+ // Initialize accountKey and sharedUrl if the configuration is present
+ aiAssistantConfigOpt.foreach { aiAssistantConfig =>
+ accountKey = aiAssistantConfig.getString("ai-service-key")
+ sharedUrl = aiAssistantConfig.getString("ai-service-url")
+ }
+
+ val validAIAssistant: String = aiAssistantConfigOpt match {
+ case Some(aiAssistantConfig) =>
+ val assistantType: String = aiAssistantConfig.getString("assistant")
+ assistantType match {
+ case "none" => noAssistant
+ case "openai" => initOpenAI()
+ case _ => noAssistant
+ }
+ case None =>
+ noAssistant
+ }
+
+ private def initOpenAI(): String = {
+ var connection: HttpURLConnection = null
+ try {
+ val url = new URL(s"${sharedUrl}/models")
+ connection = url.openConnection().asInstanceOf[HttpURLConnection]
+ connection.setRequestMethod("GET")
+ connection.setRequestProperty(
+ "Authorization",
+ s"Bearer ${accountKey.trim.replaceAll("^\"|\"$", "")}"
+ )
+ val responseCode = connection.getResponseCode
+ if (responseCode == 200) {
+ "OpenAI"
+ } else {
+ noAssistant
+ }
+ } catch {
+ case e: Exception =>
+ noAssistant
+ } finally {
+ if (connection != null) {
+ connection.disconnect()
+ }
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/AiAssistantResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/AiAssistantResource.scala
new file mode 100644
index 00000000000..3da8de87e81
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/AiAssistantResource.scala
@@ -0,0 +1,218 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import io.dropwizard.auth.Auth
+import kong.unirest.Unirest
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.web.resource.aiassistant.AiAssistantManager
+import play.api.libs.json._
+
+import java.nio.file.Paths
+import java.util.Base64
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.{MediaType, Response}
+import scala.sys.process._
+
+case class AIAssistantRequest(code: String, lineNumber: Int, allcode: String)
+
+case class LocateUnannotatedRequest(selectedCode: String, startLine: Int)
+
+case class UnannotatedArgument(
+ name: String,
+ startLine: Int,
+ startColumn: Int,
+ endLine: Int,
+ endColumn: Int
+)
+
+object UnannotatedArgument {
+ implicit val format: Format[UnannotatedArgument] = Json.format[UnannotatedArgument]
+}
+
+@Path("/aiassistant")
+class AIAssistantResource {
+ val objectMapper = new ObjectMapper()
+ objectMapper.registerModule(DefaultScalaModule)
+ final private lazy val isEnabled = AiAssistantManager.validAIAssistant
+
+ @GET
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/isenabled")
+ def isAIAssistantEnable: String = isEnabled
+
+ /**
+ * A way to send prompts to open ai
+ *
+ * @param prompt The input prompt for the OpenAI model.
+ * @param user The authenticated session user.
+ * @return A response containing the generated comment from OpenAI or an error message.
+ */
+ @POST
+ @Path("/openai")
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ def sendPromptToOpenAIApi(prompt: String, @Auth user: SessionUser): Response = {
+ // Prepare the final prompt by escaping necessary characters
+ // Escape backslashes and double quotes in the prompt to prevent breaking the JSON format
+ val finalPrompt = prompt.replace("\\", "\\\\").replace("\"", "\\\"")
+
+ // Create the JSON request body
+ val requestBody =
+ s"""
+ |{
+ | "model": "gpt-4o",
+ | "messages": [{"role": "user", "content": "$finalPrompt"}],
+ | "max_tokens": 1000
+ |}
+ """.stripMargin
+
+ try {
+ // Send the request to the OpenAI API using Unirest
+ val response = Unirest
+ .post("https://api.openai.com/v1/chat/completions")
+ .header("Authorization", s"Bearer ${AiAssistantManager.accountKey}")
+ .header("Content-Type", "application/json")
+ .body(requestBody)
+ .asJson()
+
+ // Return the response from the API
+ Response.status(response.getStatus).entity(response.getBody.toString).build()
+ } catch {
+ // Handle exceptions and return an error response
+ case e: Exception =>
+ e.printStackTrace()
+ Response
+ .status(Response.Status.INTERNAL_SERVER_ERROR)
+ .entity("Error occur when requesting the OpenAI API")
+ .build()
+ }
+ }
+
+ /**
+ * To get the type annotation suggestion from OpenAI
+ */
+ @POST
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/annotationresult")
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ def getAnnotation(
+ request: AIAssistantRequest,
+ @Auth user: SessionUser
+ ): Response = {
+ val finalPrompt = generatePrompt(request.code, request.lineNumber, request.allcode)
+ val requestBodyJson = Json.obj(
+ "model" -> "gpt-4",
+ "messages" -> Json.arr(
+ Json.obj(
+ "role" -> "user",
+ "content" -> finalPrompt
+ )
+ ),
+ "max_tokens" -> 15
+ )
+
+ val response = Unirest
+ .post(s"${AiAssistantManager.sharedUrl}/chat/completions")
+ .header("Authorization", s"Bearer ${AiAssistantManager.accountKey}")
+ .header("Content-Type", "application/json")
+ .body(Json.stringify(requestBodyJson))
+ .asString()
+ if (response.getStatus >= 400) {
+ throw new RuntimeException(s"getAnnotation error: ${response.getStatus}: ${response.getBody}")
+ }
+ Response.status(response.getStatus).entity(response.getBody).build()
+ }
+
+ // Helper function to get the type annotation
+ def generatePrompt(code: String, lineNumber: Int, allcode: String): String = {
+ s"""
+ |Your task is to analyze the given Python code and provide only the type annotation as stated in the instructions.
+ |Instructions:
+ |- The provided code will only be one of the 2 situations below:
+ |- First situation: The input is not start with "def". If the provided code only contains variable, output the result in the format ":type".
+ |- Second situation: The input is start with "def". If the provided code starts with "def" (a longer line than just a variable, indicative of a function or method), output the result in the format " -> type".
+ |- The type should only be one word, such as "str", "int", etc.
+ |Examples:
+ |- First situation:
+ | - Provided code is "name", then the output may be : str
+ | - Provided code is "age", then the output may be : int
+ | - Provided code is "data", then the output may be : Tuple[int, str]
+ | - Provided code is "new_user", then the output may be : User
+ | - A special case: provided code is "self" and the context is something like "def __init__(self, username :str , age :int)", if the user requires the type annotation for the first parameter "self", then you should generate nothing.
+ |- Second situation: (actual output depends on the complete code content)
+ | - Provided code is "process_data(data: List[Tuple[int, str]], config: Dict[str, Union[int, str]])", then the output may be -> Optional[str]
+ | - Provided code is "def add(a: int, b: int)", then the output may be -> int
+ |Counterexamples:
+ | - Provided code is "def __init__(self, username: str, age: int)" and you generate the result:
+ | The result is The provided code is "def __init__(self, username: str, age: int)", so it fits the second situation, which means the result should be in " -> type" format. However, the __init__ method in Python doesn't return anything or in other words, it implicitly returns None. Hence the correct type hint would be: -> None.
+ |Details:
+ |- Provided code: $code
+ |- Line number of the provided code in the complete code context: $lineNumber
+ |- Complete code context: $allcode
+ |Important: (you must follow!!)
+ |- For the first situation: you must return strictly according to the format ": type", without adding any extra characters. No need for an explanation, just the result : type is enough!
+ |- For the second situation: you return strictly according to the format " -> type", without adding any extra characters. No need for an explanation, just the result -> type is enough!
+ """.stripMargin
+ }
+
+ @POST
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/annotate-argument")
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ def locateUnannotated(request: LocateUnannotatedRequest, @Auth user: SessionUser): Response = {
+ // Encoding the code to transmit multi-line code as a single command-line argument
+ val encodedCode = Base64.getEncoder.encodeToString(request.selectedCode.getBytes("UTF-8"))
+ val pythonScriptPath =
+ Paths
+ .get(
+ "src",
+ "main",
+ "scala",
+ "edu",
+ "uci",
+ "ics",
+ "texera",
+ "web",
+ "resource",
+ "aiassistant",
+ "type_annotation_visitor.py"
+ )
+ .toString
+
+ try {
+ val command = s"""python $pythonScriptPath "$encodedCode" ${request.startLine}"""
+ val result = command.!!
+ val parsedResult = objectMapper.readValue(result, classOf[List[List[Any]]]).map {
+ case List(name: String, startLine: Int, startColumn: Int, endLine: Int, endColumn: Int) =>
+ UnannotatedArgument(name, startLine, startColumn, endLine, endColumn)
+ case _ =>
+ throw new RuntimeException("Unexpected format in Python script result")
+ }
+ Response.ok(Json.obj("result" -> Json.toJson(parsedResult))).build()
+ } catch {
+ case e: Exception =>
+ e.printStackTrace()
+ Response.status(500).entity(s"Error executing the Python code: ${e.getMessage}").build()
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/test_type_annotation_visitor.py b/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/test_type_annotation_visitor.py
new file mode 100644
index 00000000000..280b49b2d4d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/test_type_annotation_visitor.py
@@ -0,0 +1,158 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from type_annotation_visitor import find_untyped_variables
+
+
+class TestFunctionsAndMethods:
+
+ @pytest.fixture
+ def global_functions_code(self):
+ """This is the test for global function"""
+ return """def global_function(a, b=2, /, c=3, *, d, e=5, **kwargs):
+ pass
+
+def global_function_no_return(a, b):
+ return a + b
+
+def global_function_with_return(a: int, b: int) -> int:
+ return a + b
+"""
+
+ def test_global_functions(self, global_functions_code):
+ expected_result = [
+ ["c", 1, 32, 1, 33],
+ ["a", 1, 21, 1, 22],
+ ["b", 1, 24, 1, 25],
+ ["d", 1, 40, 1, 41],
+ ["e", 1, 43, 1, 44],
+ ["kwargs", 1, 50, 1, 56],
+ ["a", 4, 31, 4, 32],
+ ["b", 4, 34, 4, 35]
+
+ ]
+ untyped_vars = find_untyped_variables(global_functions_code, 1)
+ assert untyped_vars == expected_result
+
+ @pytest.fixture
+ def class_methods_code(self):
+ """This is the test for class methods and static methods"""
+ return """class MyClass:
+ def instance_method_no_annotation(self, x, y):
+ pass
+
+ @staticmethod
+ def static_method(a, b, /, c=3, *, d, **kwargs):
+ pass
+
+ @staticmethod
+ def static_method_with_annotation(a: int, b: int, /, *, c: int = 5) -> int:
+ return a + b + c
+
+ @classmethod
+ def class_method(cls, value, /, *, option=True):
+ pass
+
+ @classmethod
+ def class_method_with_annotation(cls, value: str, /, *, flag: bool = False) -> str:
+ return value.upper()
+"""
+
+ def test_class_methods(self, class_methods_code):
+ expected_result = [
+ ["x", 2, 45, 2, 46],
+ ["y", 2, 48, 2, 49],
+ ["c", 6, 32, 6, 33],
+ ["a", 6, 23, 6, 24],
+ ["b", 6, 26, 6, 27],
+ ["d", 6, 40, 6, 41],
+ ["kwargs", 6, 45, 6, 51],
+ ["value", 14, 27, 14, 32],
+ ["option", 14, 40, 14, 46]
+ ]
+ untyped_vars = find_untyped_variables(class_methods_code, 1)
+ assert untyped_vars == expected_result
+
+ @pytest.fixture
+ def lambda_code(self):
+ """This is the test for lambda function"""
+ return """lambda_function = lambda x, y, /, z=0, *, w=1: x + y + z + w
+lambda_function_with_annotation = lambda x: x * 2
+"""
+
+ def test_lambda_functions(self, lambda_code):
+ with pytest.raises(ValueError) as exc_info:
+ find_untyped_variables(lambda_code, 1)
+ assert "Lambda functions do not support type annotation" in str(exc_info.value)
+
+ @pytest.fixture
+ def comprehensive_functions_code(self):
+ """This is the test for comprehensive function"""
+ return """def default_args_function(a, b=2, /, c=3, *, d=4):
+ pass
+
+def args_kwargs_function(*args, **kwargs):
+ pass
+
+def function_with_return_annotation(a: int, b: int, /, *, c: int = 0) -> int:
+ return a + b + c
+
+def function_without_return_annotation(a, b):
+ return a + b
+"""
+
+ def test_comprehensive(self, comprehensive_functions_code):
+ expected_result = [
+ ["c", 1, 38, 1, 39],
+ ["a", 1, 27, 1, 28],
+ ["b", 1, 30, 1, 31],
+ ["d", 1, 46, 1, 47],
+ ["args", 4, 27, 4, 31],
+ ["kwargs", 4, 35, 4, 41],
+ ["a", 10, 40, 10, 41],
+ ["b", 10, 43, 10, 44]
+ ]
+ untyped_vars = find_untyped_variables(comprehensive_functions_code, 1)
+ assert untyped_vars == expected_result
+
+ @pytest.fixture
+ def multi_line_function_code(self):
+ """This is the test for multi-line function"""
+ return """def multi_line_function(
+ a,
+ b: int = 10,
+ /,
+ c: str = "hello",
+ *,
+ d,
+ e=20,
+ **kwargs
+):
+ pass
+"""
+
+ def test_multi_lines_argument(self, multi_line_function_code):
+ expected_result = [
+ ["a", 2, 5, 2, 6],
+ ["d", 7, 5, 7, 6],
+ ["e", 8, 5, 8, 6],
+ ["kwargs", 9, 7, 9, 13]
+ ]
+ untyped_vars = find_untyped_variables(multi_line_function_code, 1)
+ assert untyped_vars == expected_result
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/type_annotation_visitor.py b/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/type_annotation_visitor.py
new file mode 100644
index 00000000000..9edf4f95bd3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/aiassistant/type_annotation_visitor.py
@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import ast
+import base64
+import json
+import sys
+
+
+class ParentNodeVisitor(ast.NodeVisitor):
+ def __init__(self):
+ self.parent = None
+
+ def generic_visit(self, node):
+ node.parent = self.parent
+ previous_parent = self.parent
+ self.parent = node
+ super().generic_visit(node)
+ self.parent = previous_parent
+
+
+class TypeAnnotationVisitor(ast.NodeVisitor):
+ def __init__(self, start_line_offset=0):
+ self.untyped_args = []
+ self.start_line_offset = start_line_offset
+
+ def visit(self, node):
+ if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+ self.process_function(node)
+ elif isinstance(node, ast.Lambda):
+ raise ValueError("Lambda functions do not support type annotation")
+ self.generic_visit(node)
+
+ def process_function(self, node):
+ # Boolean to determine if it's a global function or a method
+ is_method = isinstance(node.parent, ast.ClassDef)
+ # Boolean to determine if it's a static method
+ is_staticmethod = False
+ if is_method and hasattr(node, 'decorator_list'):
+ for decorator in node.decorator_list:
+ if isinstance(decorator, ast.Name) and decorator.id == 'staticmethod':
+ is_staticmethod = True
+ elif isinstance(decorator, ast.Attribute) and decorator.attr == 'staticmethod':
+ is_staticmethod = True
+ args = node.args
+
+ all_args = []
+ all_args.extend(args.args)
+ # Positional-only
+ all_args.extend(args.posonlyargs)
+ # Keyword-only
+ all_args.extend(args.kwonlyargs)
+ # *args
+ if args.vararg:
+ all_args.append(args.vararg)
+ # **kwargs
+ if args.kwarg:
+ all_args.append(args.kwarg)
+
+ start_index = 0
+ # Skip the "self" or "cls"
+ if is_method and not is_staticmethod:
+ start_index = 1
+ for i, arg in enumerate(all_args[start_index:]):
+ if not arg.annotation:
+ self.add_untyped_arg(arg)
+
+ def add_untyped_arg(self, arg):
+ start_line = arg.lineno + self.start_line_offset - 1
+ start_col = arg.col_offset + 1
+ end_line = start_line
+ end_col = start_col + len(arg.arg)
+ self.untyped_args.append([arg.arg, start_line, start_col, end_line, end_col])
+
+
+def find_untyped_variables(source_code, start_line):
+ tree = ast.parse(source_code)
+ ParentNodeVisitor().visit(tree)
+ visitor = TypeAnnotationVisitor(start_line_offset=start_line)
+ visitor.visit(tree)
+ return visitor.untyped_args
+
+
+if __name__ == "__main__":
+ encoded_code = sys.argv[1]
+ start_line = int(sys.argv[2])
+ # Encoding the code to transmit multi-line code as a single command-line argument before, so we need to decode it here
+ source_code = base64.b64decode(encoded_code).decode('utf-8')
+ untyped_variables = find_untyped_variables(source_code, start_line)
+ print(json.dumps(untyped_variables))
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/auth/AuthResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/auth/AuthResource.scala
new file mode 100644
index 00000000000..0f99da681d5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/auth/AuthResource.scala
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.auth
+
+import org.apache.texera.auth.JwtAuth.{TOKEN_EXPIRE_TIME_IN_MINUTES, jwtClaims, jwtToken}
+import org.apache.texera.config.UserSystemConfig
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables.USER
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.UserDao
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.apache.texera.web.model.http.request.auth.{UserLoginRequest, UserRegistrationRequest}
+import org.apache.texera.web.model.http.response.TokenIssueResponse
+import org.apache.texera.web.resource.auth.AuthResource._
+import org.jasypt.util.password.StrongPasswordEncryptor
+
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+
+object AuthResource {
+
+ private def userDao =
+ new UserDao(
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ .configuration
+ )
+
+ /**
+ * Retrieve exactly one User from databases with the given username and password.
+ * The password is used to validate against the hashed password stored in the db.
+ *
+ * @param name String
+ * @param password String, plain text password
+ * @return
+ */
+ def retrieveUserByUsernameAndPassword(name: String, password: String): Option[User] = {
+ if (password == null) return None
+ if (name == null) return None
+ Option(
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ .select()
+ .from(USER)
+ .where(USER.NAME.eq(name))
+ .fetchOneInto(classOf[User])
+ ).filter(user => new StrongPasswordEncryptor().checkPassword(password, user.getPassword))
+ }
+
+ def createAdminUser(): Unit = {
+ val adminUsername = UserSystemConfig.adminUsername
+ val adminPassword = UserSystemConfig.adminPassword
+
+ if (adminUsername.trim.nonEmpty && adminPassword.trim.nonEmpty) {
+ val existingUser = userDao.fetchByName(adminUsername)
+ if (existingUser.isEmpty) {
+ val user = new User
+ user.setName(adminUsername)
+ user.setEmail(adminUsername)
+ user.setRole(UserRoleEnum.ADMIN)
+ user.setPassword(new StrongPasswordEncryptor().encryptPassword(adminPassword))
+ userDao.insert(user)
+ }
+ }
+ }
+}
+
+@Path("/auth/")
+@Consumes(Array(MediaType.APPLICATION_JSON))
+@Produces(Array(MediaType.APPLICATION_JSON))
+class AuthResource {
+
+ @POST
+ @Path("/login")
+ def login(request: UserLoginRequest): TokenIssueResponse = {
+ retrieveUserByUsernameAndPassword(request.username, request.password) match {
+ case Some(user) =>
+ TokenIssueResponse(jwtToken(jwtClaims(user, TOKEN_EXPIRE_TIME_IN_MINUTES)))
+ case None => throw new NotAuthorizedException("Login credentials are incorrect.")
+ }
+ }
+
+ @POST
+ @Path("/register")
+ def register(request: UserRegistrationRequest): TokenIssueResponse = {
+ val username = request.username
+ if (username == null) throw new NotAcceptableException("Username cannot be null.")
+ if (username.trim.isEmpty) throw new NotAcceptableException("Username cannot be empty.")
+ userDao.fetchByName(username).size() match {
+ case 0 =>
+ val user = new User
+ user.setName(username)
+ user.setEmail(username)
+ user.setRole(UserRoleEnum.RESTRICTED)
+ // hash the plain text password
+ user.setPassword(new StrongPasswordEncryptor().encryptPassword(request.password))
+ userDao.insert(user)
+ TokenIssueResponse(jwtToken(jwtClaims(user, TOKEN_EXPIRE_TIME_IN_MINUTES)))
+ case _ =>
+ // the username exists already
+ throw new NotAcceptableException("Username exists already.")
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/auth/GoogleAuthResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/auth/GoogleAuthResource.scala
new file mode 100644
index 00000000000..2f99b9c1bd3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/auth/GoogleAuthResource.scala
@@ -0,0 +1,116 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.auth
+
+import com.google.api.client.googleapis.auth.oauth2.GoogleIdTokenVerifier
+import com.google.api.client.http.javanet.NetHttpTransport
+import com.google.api.client.json.gson.GsonFactory
+import org.apache.texera.auth.JwtAuth.{TOKEN_EXPIRE_TIME_IN_MINUTES, jwtClaims, jwtToken}
+import org.apache.texera.config.UserSystemConfig
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.UserDao
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.apache.texera.web.model.http.response.TokenIssueResponse
+import org.apache.texera.web.resource.auth.GoogleAuthResource.userDao
+
+import java.util.Collections
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+
+object GoogleAuthResource {
+ private def userDao =
+ new UserDao(
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ .configuration
+ )
+}
+
+@Path("/auth/google")
+class GoogleAuthResource {
+ final private lazy val clientId = UserSystemConfig.googleClientId
+
+ @GET
+ @Path("/clientid")
+ def getClientId: String = clientId
+
+ @POST
+ @Consumes(Array(MediaType.TEXT_PLAIN))
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @Path("/login")
+ def login(credential: String): TokenIssueResponse = {
+ val idToken =
+ new GoogleIdTokenVerifier.Builder(new NetHttpTransport, GsonFactory.getDefaultInstance)
+ .setAudience(
+ Collections.singletonList(clientId)
+ )
+ .build()
+ .verify(credential)
+ if (idToken != null) {
+ val payload = idToken.getPayload
+ val googleId = payload.getSubject
+ val googleName = payload.get("name").asInstanceOf[String]
+ val googleEmail = payload.getEmail
+ val googleAvatar = Option(payload.get("picture").asInstanceOf[String])
+ .flatMap(_.split("/").lastOption)
+ .getOrElse("")
+ val user = Option(userDao.fetchOneByGoogleId(googleId)) match {
+ case Some(user) =>
+ if (user.getName != googleName) {
+ user.setName(googleName)
+ userDao.update(user)
+ }
+ if (user.getEmail != googleEmail) {
+ user.setEmail(googleEmail)
+ userDao.update(user)
+ }
+ if (user.getGoogleAvatar != googleAvatar) {
+ user.setGoogleAvatar(googleAvatar)
+ userDao.update(user)
+ }
+ user
+ case None =>
+ Option(userDao.fetchOneByEmail(googleEmail)) match {
+ case Some(user) =>
+ if (user.getName != googleName) {
+ user.setName(googleName)
+ }
+ user.setGoogleId(googleId)
+ user.setGoogleAvatar(googleAvatar)
+ userDao.update(user)
+ user
+ case None =>
+ // create a new user with googleId
+ val user = new User
+ user.setName(googleName)
+ user.setEmail(googleEmail)
+ user.setGoogleId(googleId)
+ user.setRole(UserRoleEnum.INACTIVE)
+ user.setGoogleAvatar(googleAvatar)
+ userDao.insert(user)
+ user
+ }
+ }
+ TokenIssueResponse(jwtToken(jwtClaims(user, TOKEN_EXPIRE_TIME_IN_MINUTES)))
+ } else throw new NotAuthorizedException("Login credentials are incorrect.")
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/DashboardResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/DashboardResource.scala
new file mode 100644
index 00000000000..704219fc2d3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/DashboardResource.scala
@@ -0,0 +1,240 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.tables.pojos._
+import org.apache.texera.web.resource.dashboard.DashboardResource._
+import org.apache.texera.web.resource.dashboard.SearchQueryBuilder.{ALL_RESOURCE_TYPE, context}
+import org.apache.texera.web.resource.dashboard.user.dataset.DatasetResource.DashboardDataset
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource.DashboardWorkflow
+import org.jooq.{Field, OrderField}
+
+import java.util
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+import scala.jdk.CollectionConverters._
+
+object DashboardResource {
+ case class DashboardClickableFileEntry(
+ resourceType: String,
+ workflow: Option[DashboardWorkflow] = None,
+ project: Option[Project] = None,
+ dataset: Option[DashboardDataset] = None
+ )
+
+ case class UserInfo(userId: Integer, userName: String, googleAvatar: Option[String])
+
+ case class DashboardSearchResult(
+ results: List[DashboardClickableFileEntry],
+ more: Boolean,
+ hasMismatch: Boolean = false
+ )
+
+ /*
+ The following class describe the available params from the frontend for full text search.
+ * @param user The authenticated user performing the search.
+ * @param keywords A list of search keywords. The API will return resources that match any of these keywords.
+ * @param resourceType The type of the resources to include in the search results. Acceptable values are "workflow", "project", "file" and "" (for all types).
+ * @param creationStartDate The start of the date range for the creation time filter. It should be provided in 'yyyy-MM-dd' format.
+ * @param creationEndDate The end of the date range for the creation time filter. It should be provided in 'yyyy-MM-dd' format.
+ * @param modifiedStartDate The start of the date range for the modification time filter. It should be provided in 'yyyy-MM-dd' format.
+ * @param modifiedEndDate The end of the date range for the modification time filter. It should be provided in 'yyyy-MM-dd' format.
+ * @param owners A list of owner names to include in the search results.
+ * @param workflowIDs A list of workflow IDs to include in the search results.
+ * @param operators A list of operators to include in the search results.
+ * @param projectIds A list of project IDs to include in the search results.
+ * @param offset The number of initial results to skip. This is useful for implementing pagination.
+ * @param count The maximum number of results to return.
+ * @param orderBy The order in which to sort the results. Acceptable values are 'NameAsc', 'NameDesc', 'CreateTimeDesc', and 'EditTimeDesc'.
+ */
+ case class SearchQueryParams(
+ @QueryParam("query") keywords: java.util.List[String] = new util.ArrayList[String](),
+ @QueryParam("resourceType") @DefaultValue("") resourceType: String = ALL_RESOURCE_TYPE,
+ @QueryParam("createDateStart") @DefaultValue("") creationStartDate: String = "",
+ @QueryParam("createDateEnd") @DefaultValue("") creationEndDate: String = "",
+ @QueryParam("modifiedDateStart") @DefaultValue("") modifiedStartDate: String = "",
+ @QueryParam("modifiedDateEnd") @DefaultValue("") modifiedEndDate: String = "",
+ @QueryParam("owner") owners: java.util.List[String] = new util.ArrayList(),
+ @QueryParam("id") workflowIDs: java.util.List[Integer] = new util.ArrayList(),
+ @QueryParam("operator") operators: java.util.List[String] = new util.ArrayList(),
+ @QueryParam("projectId") projectIds: java.util.List[Integer] = new util.ArrayList(),
+ @QueryParam("datasetId") datasetIds: java.util.List[Integer] = new util.ArrayList(),
+ @QueryParam("start") @DefaultValue("0") offset: Int = 0,
+ @QueryParam("count") @DefaultValue("20") count: Int = 20,
+ @QueryParam("orderBy") @DefaultValue("EditTimeDesc") orderBy: String = "EditTimeDesc"
+ )
+
+ // Construct query for workflows
+
+ def searchAllResources(
+ @Auth user: SessionUser,
+ @BeanParam params: SearchQueryParams,
+ includePublic: Boolean = false
+ ): DashboardSearchResult = {
+ val uid = user.getUid
+ val query = params.resourceType match {
+ case SearchQueryBuilder.WORKFLOW_RESOURCE_TYPE =>
+ WorkflowSearchQueryBuilder.constructQuery(uid, params, includePublic)
+ case SearchQueryBuilder.PROJECT_RESOURCE_TYPE =>
+ ProjectSearchQueryBuilder.constructQuery(uid, params, includePublic)
+ case SearchQueryBuilder.DATASET_RESOURCE_TYPE =>
+ DatasetSearchQueryBuilder.constructQuery(uid, params, includePublic)
+ case SearchQueryBuilder.ALL_RESOURCE_TYPE =>
+ val q1 = WorkflowSearchQueryBuilder.constructQuery(uid, params, includePublic)
+ val q3 = ProjectSearchQueryBuilder.constructQuery(uid, params, includePublic)
+ val q4 = DatasetSearchQueryBuilder.constructQuery(uid, params, includePublic)
+ q1.unionAll(q3).unionAll(q4)
+ case _ => throw new IllegalArgumentException(s"Unknown resource type: ${params.resourceType}")
+ }
+
+ val finalQuery =
+ query.orderBy(getOrderFields(params): _*).offset(params.offset).limit(params.count + 1)
+ val queryResult = finalQuery.fetch()
+
+ val allEntries = queryResult.asScala.toList
+ .take(params.count)
+ .map(record => {
+ val resourceType = record.get("resourceType", classOf[String])
+ resourceType match {
+ case SearchQueryBuilder.WORKFLOW_RESOURCE_TYPE =>
+ WorkflowSearchQueryBuilder.toEntry(uid, record)
+ case SearchQueryBuilder.PROJECT_RESOURCE_TYPE =>
+ ProjectSearchQueryBuilder.toEntry(uid, record)
+ case SearchQueryBuilder.DATASET_RESOURCE_TYPE =>
+ DatasetSearchQueryBuilder.toEntry(uid, record)
+ }
+ })
+
+ val entries = allEntries.filter(_ != null)
+ val hasMismatch =
+ params.resourceType match {
+ case SearchQueryBuilder.DATASET_RESOURCE_TYPE | SearchQueryBuilder.ALL_RESOURCE_TYPE =>
+ allEntries.exists(_ == null)
+ case _ =>
+ false
+ }
+
+ DashboardSearchResult(
+ results = entries,
+ more = queryResult.size() > params.count,
+ hasMismatch = hasMismatch
+ )
+ }
+
+ def getOrderFields(
+ searchQueryParams: SearchQueryParams
+ ): List[OrderField[_]] = {
+ // Regex pattern to extract column name and order direction
+ val pattern = "(Name|CreateTime|EditTime)(Asc|Desc)".r
+
+ searchQueryParams.orderBy match {
+ case pattern(column, order) =>
+ val field = getColumnField(column)
+ field match {
+ case Some(value) =>
+ List(order match {
+ case "Asc" => value.asc()
+ case "Desc" => value.desc()
+ })
+ case None => List()
+ }
+ case _ => List() // Default case if the orderBy string doesn't match the pattern
+ }
+ }
+
+ // Helper method to map column names to actual database fields based on resource type
+ private def getColumnField(columnName: String): Option[Field[_]] = {
+ Option(columnName match {
+ case "Name" => UnifiedResourceSchema.resourceNameField
+ case "CreateTime" => UnifiedResourceSchema.resourceCreationTimeField
+ case "EditTime" => UnifiedResourceSchema.resourceLastModifiedTimeField
+ case _ => null // Default case for unmatched resource types or column names
+ })
+ }
+
+}
+
+@Produces(Array(MediaType.APPLICATION_JSON))
+@Path("/dashboard")
+class DashboardResource {
+
+ /**
+ * This method performs a full-text search across all resources - workflows, projects, and files -
+ * that match the specified keywords.
+ * It supports advanced filters such as resource type, creation and modification dates, owner,
+ * workflow IDs, operators, project IDs and allows to specify the number of results and their ordering.
+ *
+ * This method utilizes MySQL Boolean Full-Text Searches
+ * reference: https://dev.mysql.com/doc/refman/8.0/en/fulltext-boolean.html
+ *
+ * @return A DashboardSearchResult object containing a list of DashboardClickableFileEntry objects that match the search criteria, and a boolean indicating whether more results are available.
+ */
+ @GET
+ @Path("/search")
+ def searchAllResourcesCall(
+ @Auth user: SessionUser,
+ @BeanParam params: SearchQueryParams,
+ @QueryParam("includePublic") includePublic: Boolean = false
+ ): DashboardSearchResult = {
+ DashboardResource.searchAllResources(user, params, includePublic = includePublic)
+ }
+
+ @GET
+ @Path("/publicSearch")
+ def searchAllPublicResourceCall(
+ @BeanParam params: SearchQueryParams,
+ @QueryParam("includePublic ") includePublic: Boolean = true
+ ): DashboardSearchResult = {
+ DashboardResource.searchAllResources(
+ new SessionUser(new User()),
+ params,
+ includePublic = includePublic
+ )
+ }
+
+ @GET
+ @Path("/resultsOwnersInfo")
+ def resultsOwnersInfo(
+ @QueryParam("userIds") userIds: util.List[Integer]
+ ): util.Map[Integer, UserInfo] = {
+ val scalaUserIds: Set[Integer] = userIds.asScala.toSet
+
+ val records = context
+ .select(USER.UID, USER.NAME, USER.GOOGLE_AVATAR)
+ .from(USER)
+ .where(USER.UID.in(scalaUserIds.asJava))
+ .fetch()
+
+ val userIdToInfoMap = records.asScala
+ .map { record =>
+ val userId = record.get(USER.UID)
+ val userName = record.get(USER.NAME)
+ val googleAvatar = Option(record.get(USER.GOOGLE_AVATAR))
+ userId -> UserInfo(userId, userName, googleAvatar)
+ }
+ .toMap
+ .asJava
+
+ userIdToInfoMap
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/DatasetSearchQueryBuilder.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/DatasetSearchQueryBuilder.scala
new file mode 100644
index 00000000000..89fe805d58c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/DatasetSearchQueryBuilder.scala
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.storage.util.LakeFSStorageClient
+import org.apache.texera.dao.jooq.generated.Tables.{DATASET, DATASET_USER_ACCESS}
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.User.USER
+import org.apache.texera.dao.jooq.generated.tables.pojos.{Dataset, User}
+import org.apache.texera.web.resource.dashboard.DashboardResource.DashboardClickableFileEntry
+import org.apache.texera.web.resource.dashboard.FulltextSearchQueryUtils.{
+ getContainsFilter,
+ getDateFilter,
+ getFullTextSearchFilter
+}
+import org.apache.texera.web.resource.dashboard.user.dataset.DatasetResource.DashboardDataset
+import org.jooq.impl.DSL
+import org.jooq.{Condition, GroupField, Record, TableLike}
+
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+object DatasetSearchQueryBuilder extends SearchQueryBuilder with LazyLogging {
+ override protected val mappedResourceSchema: UnifiedResourceSchema = UnifiedResourceSchema(
+ resourceType = DSL.inline(SearchQueryBuilder.DATASET_RESOURCE_TYPE),
+ name = DATASET.NAME,
+ description = DATASET.DESCRIPTION,
+ creationTime = DATASET.CREATION_TIME,
+ ownerId = DATASET.OWNER_UID,
+ did = DATASET.DID,
+ repositoryName = DATASET.REPOSITORY_NAME,
+ isDatasetPublic = DATASET.IS_PUBLIC,
+ isDatasetDownloadable = DATASET.IS_DOWNLOADABLE,
+ datasetUserAccess = DATASET_USER_ACCESS.PRIVILEGE
+ )
+
+ /*
+ * constructs the FROM clause for querying datasets with specific access controls.
+ *
+ * Parameter:
+ * - uid: Integer - Represents the unique identifier of the current user.
+ * - uid is 'null' if the user is not logged in or performing a public search.
+ * - Otherwise, `uid` holds the identifier for the logged-in user.
+ * - includePublic - Boolean - Specifies whether to include public datasets in the result.
+ */
+ override protected def constructFromClause(
+ uid: Integer,
+ params: DashboardResource.SearchQueryParams,
+ includePublic: Boolean = false
+ ): TableLike[_] = {
+ val baseJoin = DATASET
+ .leftJoin(DATASET_USER_ACCESS)
+ .on(DATASET_USER_ACCESS.DID.eq(DATASET.DID))
+ .leftJoin(USER)
+ .on(USER.UID.eq(DATASET.OWNER_UID))
+
+ // Default condition starts as true, ensuring all datasets are selected initially.
+ var condition: Condition = DSL.trueCondition()
+
+ if (uid == null) {
+ // If `uid` is null, the user is not logged in or performing a public search
+ // We only select datasets marked as public
+ condition = DATASET.IS_PUBLIC.eq(true)
+ } else {
+ // When `uid` is present, we add a condition to only include datasets with direct user access.
+ val userAccessCondition = DATASET_USER_ACCESS.UID.eq(uid)
+
+ if (includePublic) {
+ // If `includePublic` is true, we extend visibility to public datasets as well.
+ condition = userAccessCondition.or(DATASET.IS_PUBLIC.eq(true))
+ } else {
+ condition = userAccessCondition
+ }
+ }
+ baseJoin.where(condition)
+ }
+
+ override protected def constructWhereClause(
+ uid: Integer,
+ params: DashboardResource.SearchQueryParams
+ ): Condition = {
+ val splitKeywords = params.keywords.asScala
+ .flatMap(_.split("[+\\-()<>~*@\"]"))
+ .filter(_.nonEmpty)
+ .toSeq
+
+ getDateFilter(
+ params.creationStartDate,
+ params.creationEndDate,
+ DATASET.CREATION_TIME
+ )
+ .and(getContainsFilter(params.datasetIds, DATASET.DID))
+ .and(
+ getFullTextSearchFilter(splitKeywords, List(DATASET.NAME, DATASET.DESCRIPTION))
+ )
+ }
+
+ override protected def getGroupByFields: Seq[GroupField] = {
+ Seq.empty
+ }
+
+ override protected def toEntryImpl(
+ uid: Integer,
+ record: Record
+ ): DashboardResource.DashboardClickableFileEntry = {
+ val dataset = record.into(DATASET).into(classOf[Dataset])
+ val owner = record.into(USER).into(classOf[User])
+ var size = 0L
+
+ try {
+ size = LakeFSStorageClient.retrieveRepositorySize(dataset.getRepositoryName)
+ } catch {
+ case e: io.lakefs.clients.sdk.ApiException =>
+ // Treat all LakeFS ApiException as mismatch (repository not found, being deleted, or any fatal error)
+ logger.error(
+ s"LakeFS ApiException for dataset repository '${dataset.getRepositoryName}': ${e.getMessage}",
+ e
+ )
+ return null
+ }
+
+ val dd = DashboardDataset(
+ dataset,
+ owner.getEmail,
+ record.get(DATASET_USER_ACCESS.PRIVILEGE, classOf[PrivilegeEnum]),
+ dataset.getOwnerUid == uid,
+ size
+ )
+ DashboardClickableFileEntry(
+ resourceType = SearchQueryBuilder.DATASET_RESOURCE_TYPE,
+ dataset = Some(dd)
+ )
+ }
+}
+
+class DatasetSearchQueryBuilder {}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/FulltextSearchQueryUtils.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/FulltextSearchQueryUtils.scala
new file mode 100644
index 00000000000..2901b28fe0b
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/FulltextSearchQueryUtils.scala
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard
+
+import org.jooq.impl.DSL.{condition, noCondition}
+import org.jooq.{Condition, Field}
+
+import java.sql.Timestamp
+import java.text.{ParseException, SimpleDateFormat}
+import java.util.concurrent.TimeUnit
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+object FulltextSearchQueryUtils {
+
+ var usePgroonga: Boolean = true // only override by tests
+
+ def getFullTextSearchFilter(
+ keywords: Seq[String],
+ fields: List[Field[String]]
+ ): Condition = {
+ // If no target columns, skip fulltext search
+ if (fields.isEmpty) {
+ return noCondition()
+ }
+ // Filter out empty keywords and trim
+ val trimmedKeywords = keywords.filter(_.nonEmpty).map(_.trim)
+ // If no keywords, skip fulltext search
+ if (trimmedKeywords.isEmpty) {
+ return noCondition()
+ }
+ // Concatenate the fields into a single expression.
+ val combinedFields = fields
+ .map(f => s"COALESCE($f, '')") // convert null values to empty string
+ .mkString(" || ' ' || ")
+ if (usePgroonga) {
+ // Combine all keywords (AND) into a single PGroonga
+ // fuzzy search condition with a fixed threshold
+ val fuzzySearchCondition =
+ s"($combinedFields) &@~ pgroonga_condition('${trimmedKeywords.mkString(" ")}', fuzzy_max_distance_ratio => 0.34)"
+ // Return the condition
+ condition(fuzzySearchCondition, trimmedKeywords.mkString(" "))
+ } else {
+ // Only invoked by tests that uses embedded DB
+ trimmedKeywords.foldLeft(noCondition()) { (acc, keyword) =>
+ val words = keyword.split("\\s+").filter(_.nonEmpty)
+ val tsQuery = words.mkString(" & ")
+ val conditionExpr =
+ s"to_tsvector('english', $combinedFields) @@ to_tsquery('english', '$tsQuery')"
+ acc.and(condition(conditionExpr, keyword))
+ }
+ }
+ }
+
+ /**
+ * Generates a filter condition for querying based on whether a specified field contains any of the given values.
+ *
+ * This method converts a Java list of values into a Scala set to ensure uniqueness, and then iterates over each unique value,
+ * constructing a filter condition that checks if the specified field equals any of those values. The resulting condition
+ * is a disjunction (`OR`) of all these equality conditions, which can be used in database queries to find records where
+ * the field matches any of the provided values.
+ *
+ * @tparam T The type of the elements in the `values` list and the type of the field being compared.
+ * @param values A Java list of values to be checked against the field. The list is converted to a Scala set to remove duplicates.
+ * @param field The field to be checked for containing any of the values in the `values` list. This is typically a field in a database table.
+ * @return A `Condition` that represents the disjunction of equality checks between the field and each unique value in the input list.
+ * This condition can be used as part of a query to select records where the field matches any of the specified values.
+ */
+ def getContainsFilter[T](values: java.util.List[T], field: Field[T]): Condition = {
+ val valueSet = values.asScala.toSet
+ var filterForOneField: Condition = noCondition()
+ for (value <- valueSet) {
+ filterForOneField = filterForOneField.or(field.eq(value))
+ }
+ filterForOneField
+ }
+
+ /**
+ * Returns a date filter condition for the specified date range and date type.
+ *
+ * @param startDate A string representing the start date of the filter range in "yyyy-MM-dd" format.
+ * If empty, the default value "1970-01-01" will be used.
+ * @param endDate A string representing the end date of the filter range in "yyyy-MM-dd" format.
+ * If empty, the default value "9999-12-31" will be used.
+ * @param fieldToFilterOn the field for applying the start and end dates.
+ * @return A Condition object that can be used to filter workflows based on the date range and type.
+ */
+ @throws[ParseException]
+ def getDateFilter(
+ startDate: String,
+ endDate: String,
+ fieldToFilterOn: Field[Timestamp]
+ ): Condition = {
+ if (startDate.nonEmpty || endDate.nonEmpty) {
+ val start = if (startDate.nonEmpty) startDate else "1970-01-01"
+ val end = if (endDate.nonEmpty) endDate else "9999-12-31"
+ val dateFormat = new SimpleDateFormat("yyyy-MM-dd")
+
+ val startTimestamp = new Timestamp(dateFormat.parse(start).getTime)
+ val endTimestamp =
+ if (end == "9999-12-31") {
+ new Timestamp(dateFormat.parse(end).getTime)
+ } else {
+ new Timestamp(
+ dateFormat.parse(end).getTime + TimeUnit.DAYS.toMillis(1) - 1
+ )
+ }
+ fieldToFilterOn.between(startTimestamp, endTimestamp)
+ } else {
+ noCondition()
+ }
+ }
+
+ /**
+ * Helper function to retrieve the operators filter.
+ * Applies a filter based on the specified operators.
+ *
+ * @param operators The list of operators to filter by.
+ * @return The operators filter.
+ */
+ def getOperatorsFilter(
+ operators: java.util.List[String],
+ field: Field[String]
+ ): Condition = {
+ // Convert to a Set to avoid duplicates
+ val operatorSet = operators.asScala.toSet
+ // Start with a "no condition" (logical TRUE) so we can accumulate
+ var fieldFilter = noCondition()
+
+ // For each operator, build the substring pattern
+ operatorSet.foreach { operator =>
+ // e.g. => % "operatorType":"someOperator" %
+ val searchKey = s"""%"operatorType":"$operator"%"""
+
+ // Use jOOQ's likeIgnoreCase for case-insensitive matching
+ val cond = field.likeIgnoreCase(searchKey)
+
+ // Accumulate with OR
+ fieldFilter = fieldFilter.or(cond)
+ }
+
+ fieldFilter
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/ProjectSearchQueryBuilder.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/ProjectSearchQueryBuilder.scala
new file mode 100644
index 00000000000..5e247f87172
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/ProjectSearchQueryBuilder.scala
@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard
+
+import org.apache.texera.dao.jooq.generated.Tables.{PROJECT, PROJECT_USER_ACCESS}
+import org.apache.texera.dao.jooq.generated.tables.pojos.Project
+import org.apache.texera.web.resource.dashboard.DashboardResource.DashboardClickableFileEntry
+import org.apache.texera.web.resource.dashboard.FulltextSearchQueryUtils.{
+ getContainsFilter,
+ getDateFilter,
+ getFullTextSearchFilter
+}
+import org.jooq.impl.DSL
+import org.jooq.{Condition, GroupField, Record, TableLike}
+
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+object ProjectSearchQueryBuilder extends SearchQueryBuilder {
+
+ override val mappedResourceSchema: UnifiedResourceSchema = UnifiedResourceSchema(
+ resourceType = DSL.inline(SearchQueryBuilder.PROJECT_RESOURCE_TYPE),
+ name = PROJECT.NAME,
+ description = PROJECT.DESCRIPTION,
+ creationTime = PROJECT.CREATION_TIME,
+ lastModifiedTime = PROJECT.CREATION_TIME,
+ pid = PROJECT.PID,
+ ownerId = PROJECT.OWNER_ID,
+ projectColor = PROJECT.COLOR
+ )
+
+ override protected def constructFromClause(
+ uid: Integer,
+ params: DashboardResource.SearchQueryParams,
+ includePublic: Boolean = false
+ ): TableLike[_] = {
+ PROJECT
+ .leftJoin(PROJECT_USER_ACCESS)
+ .on(PROJECT_USER_ACCESS.PID.eq(PROJECT.PID))
+ .where(PROJECT_USER_ACCESS.UID.eq(uid))
+ }
+
+ override protected def constructWhereClause(
+ uid: Integer,
+ params: DashboardResource.SearchQueryParams
+ ): Condition = {
+ val splitKeywords = params.keywords.asScala
+ .flatMap(_.split("[+\\-()<>~*@\"]"))
+ .filter(_.nonEmpty)
+ .toSeq
+
+ getDateFilter(
+ params.creationStartDate,
+ params.creationEndDate,
+ PROJECT.CREATION_TIME
+ )
+ .and(getContainsFilter(params.projectIds, PROJECT.PID))
+ .and(
+ getFullTextSearchFilter(splitKeywords, List(PROJECT.NAME, PROJECT.DESCRIPTION))
+ )
+ }
+
+ override protected def getGroupByFields: Seq[GroupField] = Seq.empty
+
+ override def toEntryImpl(
+ uid: Integer,
+ record: Record
+ ): DashboardResource.DashboardClickableFileEntry = {
+ val dp = record.into(PROJECT).into(classOf[Project])
+ DashboardClickableFileEntry(SearchQueryBuilder.PROJECT_RESOURCE_TYPE, project = Some(dp))
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/SearchQueryBuilder.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/SearchQueryBuilder.scala
new file mode 100644
index 00000000000..e19755e6116
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/SearchQueryBuilder.scala
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard
+
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.web.resource.dashboard.DashboardResource.{
+ DashboardClickableFileEntry,
+ SearchQueryParams
+}
+import org.apache.texera.web.resource.dashboard.SearchQueryBuilder.context
+import org.jooq._
+
+object SearchQueryBuilder {
+
+ def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ val FILE_RESOURCE_TYPE = "file"
+ val WORKFLOW_RESOURCE_TYPE = "workflow"
+ val PROJECT_RESOURCE_TYPE = "project"
+ val DATASET_RESOURCE_TYPE = "dataset"
+ val ALL_RESOURCE_TYPE = ""
+}
+
+trait SearchQueryBuilder {
+
+ protected val mappedResourceSchema: UnifiedResourceSchema
+
+ protected def constructFromClause(
+ uid: Integer,
+ params: SearchQueryParams,
+ includePublic: Boolean = false
+ ): TableLike[_]
+
+ protected def constructWhereClause(uid: Integer, params: SearchQueryParams): Condition
+
+ protected def getGroupByFields: Seq[GroupField] = Seq.empty
+
+ protected def toEntryImpl(uid: Integer, record: Record): DashboardClickableFileEntry
+
+ private def translateRecord(record: Record): Record = mappedResourceSchema.translateRecord(record)
+
+ def toEntry(uid: Integer, record: Record): DashboardClickableFileEntry = {
+ toEntryImpl(uid, translateRecord(record))
+ }
+
+ final def constructQuery(
+ uid: Integer,
+ params: SearchQueryParams,
+ includePublic: Boolean
+ ): SelectHavingStep[Record] = {
+ val query: SelectGroupByStep[Record] = context
+ .selectDistinct(mappedResourceSchema.allFields: _*)
+ .from(constructFromClause(uid, params, includePublic))
+ .where(constructWhereClause(uid, params))
+ val groupByFields = getGroupByFields
+ if (groupByFields.nonEmpty) {
+ query.groupBy(groupByFields: _*)
+ } else {
+ query
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/UnifiedResourceSchema.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/UnifiedResourceSchema.scala
new file mode 100644
index 00000000000..dbcf1926407
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/UnifiedResourceSchema.scala
@@ -0,0 +1,170 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard
+
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.web.resource.dashboard.UnifiedResourceSchema.context
+import org.jooq.impl.DSL
+import org.jooq.{Field, Record}
+
+import java.sql.Timestamp
+import scala.collection.mutable
+
+object UnifiedResourceSchema {
+
+ // Define alias strings
+ private val resourceTypeAlias = "resourceType"
+ private val resourceNameAlias = "resourceName"
+ private val resourceDescriptionAlias = "resourceDescription"
+ private val resourceCreationTimeAlias = "resourceCreationTime"
+ private val resourceOwnerIdAlias = "resourceOwnerId"
+ private val resourceLastModifiedTimeAlias = "resourceLastModifiedTime"
+
+ // Use the alias variables to create fields
+ val resourceTypeField: Field[_] = DSL.field(DSL.name(resourceTypeAlias))
+ val resourceNameField: Field[_] = DSL.field(DSL.name(resourceNameAlias))
+ val resourceDescriptionField: Field[_] = DSL.field(DSL.name(resourceDescriptionAlias))
+ val resourceCreationTimeField: Field[_] = DSL.field(DSL.name(resourceCreationTimeAlias))
+ val resourceOwnerIdField: Field[_] = DSL.field(DSL.name(resourceOwnerIdAlias))
+ val resourceLastModifiedTimeField: Field[_] = DSL.field(DSL.name(resourceLastModifiedTimeAlias))
+
+ def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+
+ def apply(
+ resourceType: Field[String] = DSL.inline(""),
+ name: Field[String] = DSL.inline(""),
+ description: Field[String] = DSL.inline(""),
+ creationTime: Field[Timestamp] = DSL.cast(null, classOf[Timestamp]),
+ lastModifiedTime: Field[Timestamp] = DSL.cast(null, classOf[Timestamp]),
+ ownerId: Field[Integer] = DSL.cast(null, classOf[Integer]),
+ wid: Field[Integer] = DSL.cast(null, classOf[Integer]),
+ workflowUserAccess: Field[PrivilegeEnum] = DSL.castNull(classOf[PrivilegeEnum]),
+ projectsOfWorkflow: Field[String] = DSL.inline(""),
+ uid: Field[Integer] = DSL.cast(null, classOf[Integer]),
+ userName: Field[String] = DSL.inline(""),
+ userEmail: Field[String] = DSL.inline(""),
+ pid: Field[Integer] = DSL.cast(null, classOf[Integer]),
+ projectOwnerId: Field[Integer] = DSL.cast(null, classOf[Integer]),
+ projectColor: Field[String] = DSL.inline(""),
+ did: Field[Integer] = DSL.cast(null, classOf[Integer]),
+ datasetStoragePath: Field[String] = DSL.cast(null, classOf[String]),
+ repositoryName: Field[String] = DSL.inline(""),
+ isDatasetPublic: Field[java.lang.Boolean] = DSL.cast(null, classOf[java.lang.Boolean]),
+ isDatasetDownloadable: Field[java.lang.Boolean] = DSL.cast(null, classOf[java.lang.Boolean]),
+ datasetUserAccess: Field[PrivilegeEnum] = DSL.castNull(classOf[PrivilegeEnum])
+ ): UnifiedResourceSchema = {
+ new UnifiedResourceSchema(
+ Seq(
+ resourceType -> resourceType.as(resourceTypeAlias),
+ name -> name.as(resourceNameAlias),
+ description -> description.as(resourceDescriptionAlias),
+ creationTime -> creationTime.as(resourceCreationTimeAlias),
+ lastModifiedTime -> lastModifiedTime.as(resourceLastModifiedTimeAlias),
+ ownerId -> ownerId.as(resourceOwnerIdAlias),
+ wid -> wid.as("wid"),
+ workflowUserAccess -> workflowUserAccess.as("workflow_privilege"),
+ projectsOfWorkflow -> projectsOfWorkflow.as("projects"),
+ uid -> uid.as("uid"),
+ userName -> userName.as("userName"),
+ userEmail -> userEmail.as("email"),
+ pid -> pid.as("pid"),
+ projectOwnerId -> projectOwnerId.as("owner_uid"),
+ projectColor -> projectColor.as("color"),
+ did -> did.as("did"),
+ datasetStoragePath -> datasetStoragePath.as("dataset_storage_path"),
+ repositoryName -> repositoryName.as("repository_name"),
+ isDatasetPublic -> isDatasetPublic.as("is_dataset_public"),
+ isDatasetDownloadable -> isDatasetDownloadable.as("is_dataset_downloadable"),
+ datasetUserAccess -> datasetUserAccess.as("user_dataset_access")
+ )
+ )
+ }
+}
+
+/**
+ * Refer to /sql/texera_ddl.sql to understand what each attribute is
+ *
+ * Attributes common across all resource types:
+ * - `resourceType`: The type of the resource (e.g., project, workflow, file) as a `String`.
+ * - `name`: The name of the resource as a `String`.
+ * - `description`: A textual description of the resource as a `String`.
+ * - `creationTime`: The timestamp when the resource was created, as a `Timestamp`.
+ * - `lastModifiedTime`: The timestamp of the last modification to the resource, as a `Timestamp` (applicable to workflows).
+ * - `ownerId`: The identifier of the resource's owner, as an `Integer`.
+ *
+ * Attributes specific to workflows:
+ * - `wid`: Workflow ID, as an `Integer`.
+ * - `workflowUserAccess`: Access privileges associated with the workflow, as a `PrivilegeEnum`.
+ * - `projectsOfWorkflow`: IDs of projects associated with the workflow, concatenated as a `String`.
+ * - `uid`: User ID associated with the workflow, as an `Integer`.
+ * - `userName`: Name of the user associated with the workflow, as a `String`.
+ * - `userEmail`: Email of the user associated with the workflow, as a `String`.
+ *
+ * Attributes specific to projects:
+ * - `pid`: Project ID, as an `Integer`.
+ * - `projectOwnerId`: ID of the project owner, as an `Integer`.
+ * - `projectColor`: Color associated with the project, as a `String`.
+ *
+ * Attributes specific to files:
+ * - `fid`: File ID, as an `Integer`.
+ * - `fileUploadTime`: Timestamp when the file was uploaded, as a `Timestamp`.
+ * - `filePath`: Path of the file, as a `String`.
+ * - `fileSize`: Size of the file, as an `Integer`.
+ * - `fileUserAccess`: Access privileges for the file, as a `UserFileAccessPrivilege`.
+ *
+ * Attributes specific to datasets:
+ * - `did`: Dataset ID, as an `Integer`.
+ * - `datasetStoragePath`: The storage path of the dataset, as a `String`.
+ * - `repositoryName`: The name of the repository where the dataset is stored, as a `String`.
+ * - `isDatasetPublic`: Indicates if the dataset is public, as a `Boolean`.
+ * - `isDatasetDownloadable`: Indicates if the dataset is downloadable, as a `Boolean`.
+ * - `datasetUserAccess`: Access privileges for the dataset, as a `PrivilegeEnum`
+ */
+class UnifiedResourceSchema private (
+ fieldMappingSeq: Seq[(Field[_], Field[_])]
+) {
+ val allFields: Seq[Field[_]] = fieldMappingSeq.map(_._2)
+
+ private val translatedFieldSet: Seq[(Field[_], Field[_])] = {
+ val addedFields = new mutable.HashSet[Field[_]]()
+ val output = new mutable.ArrayBuffer[(Field[_], Field[_])]()
+ fieldMappingSeq.foreach {
+ case (original, translated) =>
+ if (!addedFields.contains(original)) {
+ addedFields.add(original)
+ output.addOne((original, translated))
+ }
+ }
+ output.toSeq
+ }
+
+ def translateRecord(record: Record): Record = {
+ val ret = context.newRecord(translatedFieldSet.map(_._1): _*)
+ translatedFieldSet.foreach {
+ case (original, translated) =>
+ ret.set(original.asInstanceOf[org.jooq.Field[Any]], record.get(translated))
+ }
+ ret
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/WorkflowSearchQueryBuilder.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/WorkflowSearchQueryBuilder.scala
new file mode 100644
index 00000000000..cfa653316d2
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/WorkflowSearchQueryBuilder.scala
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard
+
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.tables.pojos.Workflow
+import org.apache.texera.web.resource.dashboard.DashboardResource.DashboardClickableFileEntry
+import org.apache.texera.web.resource.dashboard.FulltextSearchQueryUtils._
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource.DashboardWorkflow
+import org.jooq.impl.DSL
+import org.jooq.impl.DSL.groupConcatDistinct
+import org.jooq.{Condition, GroupField, Record, TableLike}
+
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+object WorkflowSearchQueryBuilder extends SearchQueryBuilder {
+
+ override val mappedResourceSchema: UnifiedResourceSchema = {
+ UnifiedResourceSchema(
+ resourceType = DSL.inline(SearchQueryBuilder.WORKFLOW_RESOURCE_TYPE),
+ name = WORKFLOW.NAME,
+ description = WORKFLOW.DESCRIPTION,
+ creationTime = WORKFLOW.CREATION_TIME,
+ wid = WORKFLOW.WID,
+ lastModifiedTime = WORKFLOW.LAST_MODIFIED_TIME,
+ workflowUserAccess = WORKFLOW_USER_ACCESS.PRIVILEGE,
+ uid = WORKFLOW_OF_USER.UID,
+ ownerId = WORKFLOW_OF_USER.UID,
+ userName = USER.NAME,
+ projectsOfWorkflow = groupConcatDistinct(WORKFLOW_OF_PROJECT.PID)
+ )
+ }
+
+ override protected def constructFromClause(
+ uid: Integer,
+ params: DashboardResource.SearchQueryParams,
+ includePublic: Boolean = false
+ ): TableLike[_] = {
+ val baseQuery = WORKFLOW
+ .leftJoin(WORKFLOW_USER_ACCESS)
+ .on(WORKFLOW_USER_ACCESS.WID.eq(WORKFLOW.WID))
+ .leftJoin(WORKFLOW_OF_USER)
+ .on(WORKFLOW_OF_USER.WID.eq(WORKFLOW.WID))
+ .leftJoin(USER)
+ .on(USER.UID.eq(WORKFLOW_OF_USER.UID))
+ .leftJoin(WORKFLOW_OF_PROJECT)
+ .on(WORKFLOW_OF_PROJECT.WID.eq(WORKFLOW.WID))
+ .leftJoin(PROJECT_USER_ACCESS)
+ .on(PROJECT_USER_ACCESS.PID.eq(WORKFLOW_OF_PROJECT.PID))
+
+ var condition: Condition = DSL.trueCondition()
+ if (uid == null) {
+ condition = WORKFLOW.IS_PUBLIC.eq(true)
+ } else {
+ val privateAccessCondition =
+ WORKFLOW_USER_ACCESS.UID.eq(uid).or(PROJECT_USER_ACCESS.UID.eq(uid))
+ if (includePublic) {
+ condition = privateAccessCondition.or(WORKFLOW.IS_PUBLIC.eq(true))
+ } else {
+ condition = privateAccessCondition
+ }
+ }
+
+ baseQuery.where(condition)
+ }
+
+ override protected def constructWhereClause(
+ uid: Integer,
+ params: DashboardResource.SearchQueryParams
+ ): Condition = {
+ val splitKeywords = params.keywords.asScala
+ .flatMap(_.split("[+\\-()<>~*@\"]"))
+ .filter(_.nonEmpty)
+ .toSeq
+ getDateFilter(
+ params.creationStartDate,
+ params.creationEndDate,
+ WORKFLOW.CREATION_TIME
+ )
+ // Apply lastModified_time date filter
+ .and(
+ getDateFilter(
+ params.modifiedStartDate,
+ params.modifiedEndDate,
+ WORKFLOW.LAST_MODIFIED_TIME
+ )
+ )
+ // Apply workflowID filter
+ .and(getContainsFilter(params.workflowIDs, WORKFLOW.WID))
+ // Apply owner filter
+ .and(getContainsFilter(params.owners, USER.EMAIL))
+ // Apply operators filter
+ .and(getOperatorsFilter(params.operators, WORKFLOW.CONTENT))
+ // Apply projectId filter
+ .and(getContainsFilter(params.projectIds, WORKFLOW_OF_PROJECT.PID))
+ // Apply fulltext search filter
+ .and(
+ getFullTextSearchFilter(
+ splitKeywords,
+ List(WORKFLOW.NAME, WORKFLOW.DESCRIPTION, WORKFLOW.CONTENT)
+ )
+ )
+ }
+
+ override protected def getGroupByFields: Seq[GroupField] = {
+ Seq(
+ WORKFLOW.NAME,
+ WORKFLOW.DESCRIPTION,
+ WORKFLOW.CREATION_TIME,
+ WORKFLOW.WID,
+ WORKFLOW.LAST_MODIFIED_TIME,
+ WORKFLOW_USER_ACCESS.PRIVILEGE,
+ WORKFLOW_OF_USER.UID,
+ USER.NAME
+ )
+ }
+
+ override def toEntryImpl(
+ uid: Integer,
+ record: Record
+ ): DashboardResource.DashboardClickableFileEntry = {
+ val pidField = groupConcatDistinct(WORKFLOW_OF_PROJECT.PID)
+ val dw = DashboardWorkflow(
+ record.into(WORKFLOW_OF_USER).getUid.eq(uid),
+ record
+ .get(WORKFLOW_USER_ACCESS.PRIVILEGE)
+ .toString,
+ record.into(USER).getName,
+ record.into(WORKFLOW).into(classOf[Workflow]),
+ if (record.get(pidField) == null) {
+ List[Integer]()
+ } else {
+ record
+ .get(pidField)
+ .asInstanceOf[String]
+ .split(',')
+ .map(number => Integer.valueOf(number))
+ .toList
+ },
+ record.into(USER).getUid
+ )
+ DashboardClickableFileEntry(SearchQueryBuilder.WORKFLOW_RESOURCE_TYPE, workflow = Some(dw))
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/execution/AdminExecutionResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/execution/AdminExecutionResource.scala
new file mode 100644
index 00000000000..e9b890cd1d9
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/execution/AdminExecutionResource.scala
@@ -0,0 +1,224 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.admin.execution
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.web.resource.dashboard.admin.execution.AdminExecutionResource._
+import org.jooq.impl.DSL
+
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+import scala.jdk.CollectionConverters._
+
+/**
+ * This file handles various request related to saved-executions.
+ */
+
+object AdminExecutionResource {
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+
+ case class dashboardExecution(
+ workflowName: String,
+ workflowId: Integer,
+ userName: String,
+ userId: Integer,
+ executionId: Integer,
+ executionStatus: String,
+ executionTime: Double,
+ executionName: String,
+ startTime: Long,
+ endTime: Long,
+ access: Boolean
+ )
+
+ def mapToName(code: Short): String = {
+ code match {
+ case 0 => "READY"
+ case 1 => "RUNNING"
+ case 2 => "PAUSED"
+ case 3 => "COMPLETED"
+ case 4 => "FAILED"
+ case 5 => "KILLED"
+ case _ => "UNKNOWN" // or throw an exception, depends on your needs
+ }
+ }
+
+ def mapToStatus(status: String): Int = {
+ status match {
+ case "READY" => 0
+ case "RUNNING" => 1
+ case "PAUSED" => 2
+ case "COMPLETED" => 3
+ case "FAILED" => 4
+ case "KILLED" => 5
+ case _ => -1 // or throw an exception, depends on your needs
+ }
+ }
+
+ val sortFieldMapping = Map(
+ "workflow_name" -> WORKFLOW.NAME,
+ "execution_name" -> WORKFLOW_EXECUTIONS.NAME,
+ "initiator" -> USER.NAME,
+ "end_time" -> WORKFLOW_EXECUTIONS.LAST_UPDATE_TIME
+ )
+
+}
+
+@Produces(Array(MediaType.APPLICATION_JSON))
+@Path("/admin/execution")
+@RolesAllowed(Array("ADMIN"))
+class AdminExecutionResource {
+
+ @GET
+ @Path("/totalWorkflow")
+ @Produces()
+ def getTotalWorkflows: Int = {
+ context
+ .select(
+ DSL.countDistinct(WORKFLOW.WID)
+ )
+ .from(WORKFLOW_EXECUTIONS)
+ .join(WORKFLOW_VERSION)
+ .on(WORKFLOW_EXECUTIONS.VID.eq(WORKFLOW_VERSION.VID))
+ .join(USER)
+ .on(WORKFLOW_EXECUTIONS.UID.eq(USER.UID))
+ .join(WORKFLOW)
+ .on(WORKFLOW.WID.eq(WORKFLOW_VERSION.WID))
+ .fetchOne(0, classOf[Int])
+ }
+
+ /**
+ * This method retrieves latest execution of each workflow for specified page.
+ * The returned executions are sorted and filtered according to the parameters.
+ */
+ @GET
+ @Path("/executionList/{pageSize}/{pageIndex}/{sortField}/{sortDirection}")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def listWorkflows(
+ @Auth current_user: SessionUser,
+ @PathParam("pageSize") page_size: Int = 20,
+ @PathParam("pageIndex") page_index: Int = 0,
+ @PathParam("sortField") sortField: String = "end_time",
+ @PathParam("sortDirection") sortDirection: String = "desc",
+ @QueryParam("filter") filter: java.util.List[String]
+ ): List[dashboardExecution] = {
+ val filter_status = filter.asScala.map(mapToStatus).toSeq.filter(_ != -1).asJava
+
+ // Base query that retrieves latest execution info for each workflow without sorting and filtering.
+ // Only retrieving executions in current page according to pageSize and pageIndex parameters.
+ val executions_base_query = context
+ .select(
+ WORKFLOW_EXECUTIONS.UID,
+ USER.NAME,
+ WORKFLOW_VERSION.WID,
+ WORKFLOW.NAME,
+ WORKFLOW_EXECUTIONS.EID,
+ WORKFLOW_EXECUTIONS.STARTING_TIME,
+ WORKFLOW_EXECUTIONS.LAST_UPDATE_TIME,
+ WORKFLOW_EXECUTIONS.STATUS,
+ WORKFLOW_EXECUTIONS.NAME
+ )
+ .from(WORKFLOW_EXECUTIONS)
+ .join(WORKFLOW_VERSION)
+ .on(WORKFLOW_EXECUTIONS.VID.eq(WORKFLOW_VERSION.VID))
+ .join(USER)
+ .on(WORKFLOW_EXECUTIONS.UID.eq(USER.UID))
+ .join(WORKFLOW)
+ .on(WORKFLOW.WID.eq(WORKFLOW_VERSION.WID))
+ .naturalJoin(
+ context
+ .select(
+ DSL.max(WORKFLOW_EXECUTIONS.EID).as("eid")
+ )
+ .from(WORKFLOW_EXECUTIONS)
+ .join(WORKFLOW_VERSION)
+ .on(WORKFLOW_VERSION.VID.eq(WORKFLOW_EXECUTIONS.VID))
+ .groupBy(WORKFLOW_VERSION.WID)
+ )
+
+ // Apply filter if the status are not empty.
+ val executions_apply_filter = if (!filter_status.isEmpty) {
+ executions_base_query.where(WORKFLOW_EXECUTIONS.STATUS.in(filter_status))
+ } else {
+ executions_base_query
+ }
+
+ // Apply sorting if user specified.
+ var executions_apply_order =
+ executions_apply_filter.limit(page_size).offset(page_index * page_size)
+ if (sortField != "NO_SORTING") {
+ executions_apply_order = executions_apply_filter
+ .orderBy(
+ if (sortDirection == "desc") sortFieldMapping.getOrElse(sortField, WORKFLOW.NAME).desc()
+ else sortFieldMapping.getOrElse(sortField, WORKFLOW.NAME).asc()
+ )
+ .limit(page_size)
+ .offset(page_index * page_size)
+ }
+
+ val executions = executions_apply_order.fetch()
+
+ // Retrieve the id of each workflow that the user has access to.
+ val availableWorkflowIds = context
+ .select(WORKFLOW_USER_ACCESS.WID)
+ .from(WORKFLOW_USER_ACCESS)
+ .where(WORKFLOW_USER_ACCESS.UID.eq(current_user.getUid))
+ .fetchInto(classOf[Integer])
+
+ // Calculate the statistics needed for each execution.
+ executions
+ .map(workflowRecord => {
+ val startingTime =
+ workflowRecord.get(WORKFLOW_EXECUTIONS.STARTING_TIME).getTime
+
+ var lastUpdateTime: Long = 0
+ if (workflowRecord.get(WORKFLOW_EXECUTIONS.LAST_UPDATE_TIME) == null) {
+ lastUpdateTime = 0
+ } else {
+ lastUpdateTime = workflowRecord.get(WORKFLOW_EXECUTIONS.LAST_UPDATE_TIME).getTime
+ }
+
+ val timeDifferenceSeconds = (lastUpdateTime - startingTime) / 1000.0
+ val hasAccess = availableWorkflowIds.contains(workflowRecord.get(WORKFLOW_VERSION.WID))
+ dashboardExecution(
+ workflowRecord.get(WORKFLOW.NAME),
+ workflowRecord.get(WORKFLOW_VERSION.WID),
+ workflowRecord.get(USER.NAME),
+ workflowRecord.get(WORKFLOW_EXECUTIONS.UID),
+ workflowRecord.get(WORKFLOW_EXECUTIONS.EID),
+ mapToName(workflowRecord.get(WORKFLOW_EXECUTIONS.STATUS)),
+ timeDifferenceSeconds,
+ workflowRecord.get(WORKFLOW_EXECUTIONS.NAME),
+ startingTime,
+ lastUpdateTime,
+ hasAccess
+ )
+ })
+ .asScala
+ .toList
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/settings/AdminSettingsResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/settings/AdminSettingsResource.scala
new file mode 100644
index 00000000000..d98ede5610e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/settings/AdminSettingsResource.scala
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.admin.settings
+
+import com.fasterxml.jackson.annotation.JsonProperty
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.config.DefaultsConfig
+import org.apache.texera.dao.SqlServer
+import org.jooq.impl.DSL
+
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.{MediaType, Response}
+
+case class AdminSettingsPojo(
+ @JsonProperty("key") settingKey: String,
+ @JsonProperty("value") settingValue: String
+)
+
+@Path("/admin/settings")
+@Produces(Array(MediaType.APPLICATION_JSON))
+class AdminSettingsResource {
+
+ private def ctx = SqlServer.getInstance().createDSLContext()
+ private val siteSettings = DSL.table("site_settings")
+ private val key = DSL.field("key", classOf[String])
+ private val value = DSL.field("value", classOf[String])
+ private val updatedBy = DSL.field("updated_by", classOf[String])
+
+ @GET
+ @Path("{key}")
+ def getSetting(@PathParam("key") keyParam: String): AdminSettingsPojo = {
+ ctx
+ .select(key, value)
+ .from(siteSettings)
+ .where(key.eq(keyParam))
+ .fetchOneInto(classOf[AdminSettingsPojo])
+ }
+
+ @PUT
+ @Path("{key}")
+ @RolesAllowed(Array("ADMIN"))
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ def updateSetting(
+ @Auth currentUser: SessionUser,
+ @PathParam("key") keyParam: String,
+ setting: AdminSettingsPojo
+ ): Response = {
+ if (setting.settingValue != null && keyParam.nonEmpty) {
+ upsertSetting(keyParam, setting.settingValue, currentUser.getName)
+ }
+ Response.ok().build()
+ }
+
+ /**
+ * Resets the specified configuration key to its default value defined in default.conf.
+ */
+ @POST
+ @Path("/reset/{key}")
+ @RolesAllowed(Array("ADMIN"))
+ def resetSetting(
+ @Auth currentUser: SessionUser,
+ @PathParam("key") keyParam: String
+ ): Response = {
+ DefaultsConfig.allDefaults.get(keyParam) match {
+ case Some(defaultValue) =>
+ upsertSetting(keyParam, defaultValue, currentUser.getName)
+ Response.ok().build()
+ case None =>
+ Response
+ .status(Response.Status.NOT_FOUND)
+ .entity(s"No default for key '$keyParam'")
+ .build()
+ }
+ }
+
+ private def upsertSetting(keyParam: String, valueParam: String, userName: String): Unit = {
+ ctx
+ .insertInto(siteSettings)
+ .set(key, keyParam)
+ .set(value, valueParam)
+ .set(updatedBy, userName)
+ .onConflict(key)
+ .doUpdate()
+ .set(value, valueParam)
+ .set(DSL.field("updated_by", classOf[String]), userName)
+ .set(DSL.field("updated_at", classOf[java.sql.Timestamp]), DSL.currentTimestamp())
+ .execute()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/user/AdminUserResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/user/AdminUserResource.scala
new file mode 100644
index 00000000000..cd5ead915df
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/admin/user/AdminUserResource.scala
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.admin.user
+
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.User.USER
+import org.apache.texera.dao.jooq.generated.tables.UserLastActiveTime.USER_LAST_ACTIVE_TIME
+import org.apache.texera.dao.jooq.generated.tables.daos.UserDao
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.apache.texera.web.resource.EmailTemplate.createRoleChangeTemplate
+import org.apache.texera.web.resource.GmailResource.sendEmail
+import org.apache.texera.web.resource.dashboard.admin.user.AdminUserResource.userDao
+import org.apache.texera.web.resource.dashboard.user.quota.UserQuotaResource._
+import org.jasypt.util.password.StrongPasswordEncryptor
+
+import java.util
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.{MediaType, Response}
+
+case class UserInfo(
+ uid: Int,
+ name: String,
+ email: String,
+ googleId: String,
+ role: UserRoleEnum,
+ googleAvatar: String,
+ comment: String,
+ lastLogin: java.time.OffsetDateTime, // will be null if never logged in
+ accountCreation: java.time.OffsetDateTime,
+ affiliation: String,
+ joiningReason: String
+)
+
+object AdminUserResource {
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def userDao = new UserDao(context.configuration)
+}
+
+@Path("/admin/user")
+@RolesAllowed(Array("ADMIN"))
+class AdminUserResource {
+
+ /**
+ * This method returns the list of users
+ *
+ * @return a list of UserInfo
+ */
+ @GET
+ @Path("/list")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def list(): util.List[UserInfo] = {
+ AdminUserResource.context
+ .select(
+ USER.UID,
+ USER.NAME,
+ USER.EMAIL,
+ USER.GOOGLE_ID,
+ USER.ROLE,
+ USER.GOOGLE_AVATAR,
+ USER.COMMENT,
+ USER_LAST_ACTIVE_TIME.LAST_ACTIVE_TIME,
+ USER.ACCOUNT_CREATION_TIME,
+ USER.AFFILIATION,
+ USER.JOINING_REASON
+ )
+ .from(USER)
+ .leftJoin(USER_LAST_ACTIVE_TIME)
+ .on(USER.UID.eq(USER_LAST_ACTIVE_TIME.UID))
+ .fetchInto(classOf[UserInfo])
+ }
+
+ @PUT
+ @Path("/update")
+ def updateUser(user: User): Unit = {
+ val existingUser = userDao.fetchOneByEmail(user.getEmail)
+ if (existingUser != null && existingUser.getUid != user.getUid) {
+ throw new WebApplicationException("Email already exists", Response.Status.CONFLICT)
+ }
+ val updatedUser = userDao.fetchOneByUid(user.getUid)
+ val roleChanged = updatedUser.getRole != user.getRole
+ updatedUser.setName(user.getName)
+ updatedUser.setEmail(user.getEmail)
+ updatedUser.setRole(user.getRole)
+ updatedUser.setComment(user.getComment)
+ userDao.update(updatedUser)
+
+ if (roleChanged)
+ sendEmail(
+ createRoleChangeTemplate(receiverEmail = updatedUser.getEmail, newRole = user.getRole),
+ updatedUser.getEmail
+ )
+ }
+
+ @POST
+ @Path("/add")
+ def addUser(): Unit = {
+ val random = System.currentTimeMillis().toString
+ val newUser = new User
+ newUser.setName("User" + random)
+ newUser.setPassword(new StrongPasswordEncryptor().encryptPassword(random))
+ newUser.setRole(UserRoleEnum.INACTIVE)
+ userDao.insert(newUser)
+ }
+
+ @GET
+ @Path("/created_workflows")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getCreatedWorkflow(@QueryParam("user_id") user_id: Integer): List[Workflow] = {
+ getUserCreatedWorkflow(user_id)
+ }
+
+ @GET
+ @Path("/access_workflows")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getAccessedWorkflow(@QueryParam("user_id") user_id: Integer): util.List[Integer] = {
+ getUserAccessedWorkflow(user_id)
+ }
+
+ @GET
+ @Path("/user_quota_size")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getUserQuota(@QueryParam("user_id") user_id: Integer): Array[QuotaStorage] = {
+ getUserQuotaSize(user_id)
+ }
+
+ @DELETE
+ @Path("/deleteCollection/{eid}")
+ def deleteCollection(@PathParam("eid") eid: Integer): Unit = {
+ deleteExecutionCollection(eid)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/ActionType.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/ActionType.scala
new file mode 100644
index 00000000000..ce82001faaa
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/ActionType.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.hub
+
+import com.fasterxml.jackson.annotation.{JsonCreator, JsonValue}
+
+/**
+ * Defines all possible user action types for tracking.
+ * Supports JSON ↔ enum conversion and lowercase string representation.
+ */
+sealed trait ActionType {
+ @JsonValue
+ def value: String
+ override def toString: String = value
+}
+
+object ActionType {
+ case object View extends ActionType { val value = "view" }
+ case object Like extends ActionType { val value = "like" }
+ case object Clone extends ActionType { val value = "clone" }
+ case object Unlike extends ActionType { val value = "unlike" }
+
+ private val values = Seq(View, Like, Clone, Unlike)
+
+ @JsonCreator
+ def fromString(s: String): ActionType =
+ values
+ .find(_.value.equalsIgnoreCase(s))
+ .getOrElse(
+ throw new IllegalArgumentException(s"Unsupported actionType '$s'")
+ )
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/EntityTables.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/EntityTables.scala
new file mode 100644
index 00000000000..b152086d4f7
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/EntityTables.scala
@@ -0,0 +1,145 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.hub
+
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.tables.records._
+import org.jooq._
+
+object EntityTables {
+ // ==================== BASE TABLE ====================
+ sealed trait BaseEntityTable {
+ type R <: Record
+ val table: Table[R]
+ val isPublicColumn: TableField[R, java.lang.Boolean]
+ val idColumn: TableField[R, Integer]
+ }
+
+ object BaseEntityTable {
+ case object WorkflowTable extends BaseEntityTable {
+ override type R = WorkflowRecord
+ override val table: Table[WorkflowRecord] = WORKFLOW
+ override val isPublicColumn: TableField[WorkflowRecord, java.lang.Boolean] =
+ WORKFLOW.IS_PUBLIC
+ override val idColumn: TableField[WorkflowRecord, Integer] = WORKFLOW.WID
+ }
+
+ case object DatasetTable extends BaseEntityTable {
+ override type R = DatasetRecord
+ override val table: Table[DatasetRecord] = DATASET
+ override val isPublicColumn: TableField[DatasetRecord, java.lang.Boolean] = DATASET.IS_PUBLIC
+ override val idColumn: TableField[DatasetRecord, Integer] = DATASET.DID
+ }
+
+ def apply(entityType: EntityType): BaseEntityTable =
+ entityType match {
+ case EntityType.Workflow => WorkflowTable
+ case EntityType.Dataset => DatasetTable
+ }
+ }
+
+ // ==================== BASE LC (like & clone) TABLE ====================
+ sealed trait BaseLCTable {
+ type R <: Record
+ val table: Table[R]
+ val uidColumn: TableField[R, Integer]
+ val idColumn: TableField[R, Integer]
+ }
+
+ // ==================== LIKE TABLE ====================
+ sealed trait LikeTable extends BaseLCTable
+
+ object LikeTable {
+ case object WorkflowLikeTable extends LikeTable {
+ override type R = WorkflowUserLikesRecord
+ override val table: Table[WorkflowUserLikesRecord] = WORKFLOW_USER_LIKES
+ override val uidColumn: TableField[WorkflowUserLikesRecord, Integer] =
+ WORKFLOW_USER_LIKES.UID
+ override val idColumn: TableField[WorkflowUserLikesRecord, Integer] = WORKFLOW_USER_LIKES.WID
+ }
+
+ case object DatasetLikeTable extends LikeTable {
+ override type R = DatasetUserLikesRecord
+ override val table: Table[DatasetUserLikesRecord] = DATASET_USER_LIKES
+ override val uidColumn: TableField[DatasetUserLikesRecord, Integer] =
+ DATASET_USER_LIKES.UID
+ override val idColumn: TableField[DatasetUserLikesRecord, Integer] = DATASET_USER_LIKES.DID
+ }
+
+ def apply(entityType: EntityType): LikeTable =
+ entityType match {
+ case EntityType.Workflow => WorkflowLikeTable
+ case EntityType.Dataset => DatasetLikeTable
+ }
+ }
+
+ // ==================== CLONE TABLE ====================
+ sealed trait CloneTable extends BaseLCTable
+
+ object CloneTable {
+ case object WorkflowCloneTable extends CloneTable {
+ override type R = WorkflowUserClonesRecord
+ override val table: Table[WorkflowUserClonesRecord] = WORKFLOW_USER_CLONES
+ override val uidColumn: TableField[WorkflowUserClonesRecord, Integer] =
+ WORKFLOW_USER_CLONES.UID
+ override val idColumn: TableField[WorkflowUserClonesRecord, Integer] =
+ WORKFLOW_USER_CLONES.WID
+ }
+
+ def apply(entityType: EntityType): CloneTable =
+ entityType match {
+ case EntityType.Workflow => WorkflowCloneTable
+ case _ =>
+ throw new IllegalArgumentException(s"Unsupported entity type: $entityType for clone")
+ }
+ }
+
+ // ==================== VIEW COUNT TABLE ====================
+ sealed trait ViewCountTable {
+ type R <: Record
+ val table: Table[R]
+ val idColumn: TableField[R, Integer]
+ val viewCountColumn: TableField[R, Integer]
+ }
+
+ object ViewCountTable {
+ case object WorkflowViewCountTable extends ViewCountTable {
+ override type R = WorkflowViewCountRecord
+ override val table: Table[WorkflowViewCountRecord] = WORKFLOW_VIEW_COUNT
+ override val idColumn: TableField[WorkflowViewCountRecord, Integer] = WORKFLOW_VIEW_COUNT.WID
+ override val viewCountColumn: TableField[WorkflowViewCountRecord, Integer] =
+ WORKFLOW_VIEW_COUNT.VIEW_COUNT
+ }
+
+ case object DatasetViewCountTable extends ViewCountTable {
+ override type R = DatasetViewCountRecord
+ override val table: Table[DatasetViewCountRecord] = DATASET_VIEW_COUNT
+ override val idColumn: TableField[DatasetViewCountRecord, Integer] = DATASET_VIEW_COUNT.DID
+ override val viewCountColumn: TableField[DatasetViewCountRecord, Integer] =
+ DATASET_VIEW_COUNT.VIEW_COUNT
+ }
+
+ def apply(entityType: EntityType): ViewCountTable =
+ entityType match {
+ case EntityType.Workflow => WorkflowViewCountTable
+ case EntityType.Dataset => DatasetViewCountTable
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/EntityType.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/EntityType.scala
new file mode 100644
index 00000000000..0cf1bd4cc8d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/EntityType.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.hub
+
+import com.fasterxml.jackson.annotation.{JsonCreator, JsonValue}
+
+/**
+ * Defines all supported entity types for Hub resources.
+ * Enables JSON ↔ enum conversion with lowercase string representation.
+ */
+sealed trait EntityType {
+ @JsonValue
+ def value: String
+
+ override def toString: String = value
+}
+
+object EntityType {
+ case object Workflow extends EntityType { val value = "workflow" }
+ case object Dataset extends EntityType { val value = "dataset" }
+
+ private val values = Seq(Workflow, Dataset)
+
+ @JsonCreator
+ def fromString(s: String): EntityType =
+ values
+ .find(_.value.equalsIgnoreCase(s))
+ .getOrElse(
+ throw new IllegalArgumentException(s"Unsupported entityType '$s'")
+ )
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/HubResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/HubResource.scala
new file mode 100644
index 00000000000..fe687aa1d7e
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/hub/HubResource.scala
@@ -0,0 +1,699 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.hub
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.amber.core.storage.util.LakeFSStorageClient
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.enums.ActionEnum
+import org.apache.texera.dao.jooq.generated.tables.Dataset.DATASET
+import org.apache.texera.dao.jooq.generated.tables.DatasetUserAccess.DATASET_USER_ACCESS
+import org.apache.texera.dao.jooq.generated.tables.User.USER
+import org.apache.texera.dao.jooq.generated.tables.pojos.{Dataset, DatasetUserAccess}
+import org.apache.texera.web.resource.dashboard.DashboardResource.DashboardClickableFileEntry
+import org.apache.texera.web.resource.dashboard.hub.ActionType.{Clone, Like, Unlike, View}
+import org.apache.texera.web.resource.dashboard.hub.EntityTables._
+import org.apache.texera.web.resource.dashboard.hub.HubResource._
+import org.apache.texera.web.resource.dashboard.user.dataset.DatasetResource.DashboardDataset
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource.{
+ DashboardWorkflow,
+ baseWorkflowSelect,
+ mapWorkflowEntries
+}
+import org.jooq.Table
+import org.jooq.impl.DSL
+
+import java.util.regex.Pattern
+import javax.servlet.http.HttpServletRequest
+import javax.ws.rs._
+import javax.ws.rs.core.{Context, MediaType}
+import scala.collection.mutable.ListBuffer
+import scala.jdk.CollectionConverters._
+import scala.language.existentials
+
+object HubResource {
+ // Represents an entity reference for general-purpose batch APIs.
+ // Used by: isLikedHelper, recordLikeAction, getCounts, userAccess
+ case class UserRequest(entityId: Integer, entityType: EntityType)
+
+ // Extends UserRequest by adding userId, used for view tracking.
+ // Used by: postView
+ case class ViewRequest(entityId: Integer, userId: Integer, entityType: EntityType)
+
+ // Response format indicating whether a given entity is liked by the user.
+ // Returned by: isLiked (which calls isLikedHelper), and by isLikedHelper directly.
+ case class LikedResponse(
+ entityId: Integer,
+ entityType: EntityType,
+ isLiked: Boolean
+ )
+
+ // Response containing all user IDs with access to a specific entity.
+ // Returned by: userAccess endpoint
+ case class AccessResponse(
+ entityType: EntityType,
+ entityId: Integer,
+ userIds: java.util.List[Integer]
+ )
+
+ // Contains aggregated counts (view/like/clone) for a given entity.
+ // Returned by: getCounts endpoint
+ case class CountResponse(
+ entityId: Integer,
+ entityType: EntityType,
+ counts: java.util.Map[ActionType, Int]
+ )
+
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+
+ final private val ipv4Pattern: Pattern = Pattern.compile(
+ "^([0-9]{1,3}\\.){3}[0-9]{1,3}$"
+ )
+
+ /**
+ * Checks if a given user has liked a specific entity.
+ *
+ * @param userId The ID of the user.
+ * @param entityId The ID of the entity.
+ * @param entityType The type of entity being checked (must be validated).
+ * @return `true` if the user has liked the entity, otherwise `false`.
+ */
+ def isLikedHelper(
+ userId: Integer,
+ entityIds: java.util.List[Integer],
+ entityTypes: java.util.List[EntityType]
+ ): java.util.List[LikedResponse] = {
+ val reqs: List[UserRequest] =
+ entityTypes.asScala
+ .zip(entityIds.asScala)
+ .map { case (etype, id) => UserRequest(id, etype) }
+ .toList
+
+ val buffer = ListBuffer[LikedResponse]()
+ reqs
+ .groupBy(_.entityType)
+ .foreach {
+ case (etype, groupReqs) =>
+ val tbl = LikeTable(etype)
+ val ids = groupReqs.map(_.entityId)
+
+ val likedSet: Set[Int] = context
+ .select(tbl.idColumn)
+ .from(tbl.table)
+ .where(tbl.uidColumn.eq(userId))
+ .and(tbl.idColumn.in(ids: _*))
+ .fetch()
+ .asScala
+ .map(r => r.get(tbl.idColumn).intValue())
+ .toSet
+
+ groupReqs.foreach { req =>
+ val flag = likedSet.contains(req.entityId.intValue())
+ buffer += LikedResponse(req.entityId, etype, flag)
+ }
+ }
+
+ buffer.toList.asJava
+ }
+
+ /**
+ * Records a user's action in the system.
+ *
+ * @param request The HTTP request object to extract the user's IP address.
+ * @param userId The ID of the user performing the action (default is 0 for anonymous users).
+ * @param entityId The ID of the entity associated with the action.
+ * @param entityType The type of entity being acted upon (validated before processing).
+ * @param action The action performed by the user ("like", "unlike", "view", "clone").
+ */
+ def recordUserAction(
+ request: HttpServletRequest,
+ userId: Integer = Integer.valueOf(0),
+ entityId: Integer,
+ entityType: EntityType,
+ action: ActionType
+ ): Unit = {
+ val userIp = request.getRemoteAddr
+ val actionEnum = ActionEnum.values().find(_.getLiteral.equalsIgnoreCase(action.value)).get
+
+ val query = context
+ .insertInto(USER_ACTION)
+ .set(USER_ACTION.UID, userId)
+ .set(USER_ACTION.RESOURCE_ID, entityId)
+ .set(USER_ACTION.RESOURCE_TYPE, entityType.value)
+ .set(USER_ACTION.ACTION, actionEnum)
+
+ if (ipv4Pattern.matcher(userIp).matches()) {
+ query.set(USER_ACTION.IP, userIp)
+ }
+
+ query.execute()
+ }
+
+ /**
+ * Records a user's like or unlike action for a given entity.
+ *
+ * @param request The HTTP request object to extract the user's IP address.
+ * @param userRequest An object containing entityId, userId, and entityType.
+ * @param isLike A boolean flag indicating whether the action is a like (`true`) or unlike (`false`).
+ * @return `true` if the like/unlike action was recorded successfully, otherwise `false`.
+ */
+ def recordLikeAction(
+ request: HttpServletRequest,
+ userId: Integer,
+ userRequest: UserRequest,
+ isLike: Boolean
+ ): Boolean = {
+ val (entityId, entityType) =
+ (userRequest.entityId, userRequest.entityType)
+ val entityTables = LikeTable(entityType)
+ val (table, uidColumn, idColumn) =
+ (entityTables.table, entityTables.uidColumn, entityTables.idColumn)
+
+ val likedResponses = isLikedHelper(
+ userId,
+ List(entityId).asJava,
+ List(entityType).asJava
+ ).asScala
+ val alreadyLiked = likedResponses.headOption.exists(_.isLiked)
+
+ if (isLike && !alreadyLiked) {
+ context
+ .insertInto(table)
+ .set(uidColumn, userId)
+ .set(idColumn, entityId)
+ .execute()
+
+ recordUserAction(request, userId, entityId, entityType, Like)
+ true
+ } else if (!isLike && alreadyLiked) {
+ context
+ .deleteFrom(table)
+ .where(uidColumn.eq(userId).and(idColumn.eq(entityId)))
+ .execute()
+
+ recordUserAction(request, userId, entityId, entityType, Unlike)
+ true
+ } else {
+ false
+ }
+ }
+
+ /**
+ * Records a user's clone action for a given entity.
+ *
+ * @param request The HTTP request object to extract the user's IP address.
+ * @param userId The ID of the user performing the clone action.
+ * @param entityId The ID of the entity being cloned.
+ * @param entityType The type of entity being cloned (must be validated).
+ */
+ def recordCloneAction(
+ request: HttpServletRequest,
+ userId: Integer,
+ entityId: Integer,
+ entityType: EntityType
+ ): Unit = {
+
+ val entityTables = CloneTable(entityType)
+ val (table, uidColumn, idColumn) =
+ (entityTables.table, entityTables.uidColumn, entityTables.idColumn)
+
+ recordUserAction(request, userId, entityId, entityType, Clone)
+
+ val existingCloneRecord = context
+ .selectFrom(table)
+ .where(uidColumn.eq(userId))
+ .and(idColumn.eq(entityId))
+ .fetchOne()
+
+ if (existingCloneRecord == null) {
+ context
+ .insertInto(table)
+ .set(uidColumn, userId)
+ .set(idColumn, entityId)
+ .execute()
+ }
+ }
+
+ def fetchDashboardWorkflowsByWids(wids: Seq[Integer], uid: Integer): List[DashboardWorkflow] = {
+ if (wids.isEmpty) {
+ return List.empty[DashboardWorkflow]
+ }
+
+ val records = baseWorkflowSelect()
+ .where(WORKFLOW.WID.in(wids: _*))
+ .groupBy(
+ WORKFLOW.WID,
+ WORKFLOW.NAME,
+ WORKFLOW.DESCRIPTION,
+ WORKFLOW.CREATION_TIME,
+ WORKFLOW.LAST_MODIFIED_TIME,
+ WORKFLOW_USER_ACCESS.PRIVILEGE,
+ WORKFLOW_OF_USER.UID,
+ USER.NAME
+ )
+ .fetch()
+
+ mapWorkflowEntries(records, uid)
+ }
+
+ def fetchDashboardDatasetsByDids(dids: Seq[Integer], uid: Integer): List[DashboardDataset] = {
+ if (dids.isEmpty) {
+ return List.empty[DashboardDataset]
+ }
+
+ val records = context
+ .select()
+ .from(
+ DATASET
+ .leftJoin(DATASET_USER_ACCESS)
+ .on(DATASET_USER_ACCESS.DID.eq(DATASET.DID))
+ .leftJoin(USER)
+ .on(USER.UID.eq(DATASET.OWNER_UID))
+ )
+ .where(DATASET.DID.in(dids: _*))
+ .groupBy(
+ DATASET.DID,
+ DATASET.NAME,
+ DATASET.DESCRIPTION,
+ DATASET.OWNER_UID,
+ USER.NAME,
+ DATASET_USER_ACCESS.DID,
+ DATASET_USER_ACCESS.UID,
+ USER.UID
+ )
+ .fetch()
+
+ records.asScala
+ .map { record =>
+ val dataset = record.into(DATASET).into(classOf[Dataset])
+ val datasetAccess = record.into(DATASET_USER_ACCESS).into(classOf[DatasetUserAccess])
+ val ownerEmail = record.into(USER).getEmail
+ DashboardDataset(
+ isOwner = if (uid == null) false else dataset.getOwnerUid == uid,
+ dataset = dataset,
+ accessPrivilege = datasetAccess.getPrivilege,
+ ownerEmail = ownerEmail,
+ size = LakeFSStorageClient.retrieveRepositorySize(dataset.getRepositoryName)
+ )
+ }
+ .toList
+ .distinctBy(_.dataset.getDid)
+ }
+}
+
+@Produces(Array(MediaType.APPLICATION_JSON))
+@Path("/hub")
+class HubResource {
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+
+ @GET
+ @Path("/count")
+ def getCount(@QueryParam("entityType") entityType: EntityType): Integer = {
+ val entityTables = BaseEntityTable(entityType)
+ val (table, isPublicColumn) = (entityTables.table, entityTables.isPublicColumn)
+
+ context
+ .selectCount()
+ .from(table)
+ .where(isPublicColumn.eq(true))
+ .fetchOne(0, classOf[Integer])
+ }
+
+ @GET
+ @Path("/isLiked")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def isLiked(
+ @Auth user: SessionUser,
+ @QueryParam("entityId") entityIds: java.util.List[Integer],
+ @QueryParam("entityType") entityTypes: java.util.List[EntityType]
+ ): java.util.List[LikedResponse] = {
+ isLikedHelper(user.getUid, entityIds, entityTypes)
+ }
+
+ @POST
+ @Path("/like")
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ def postLike(
+ @Auth user: SessionUser,
+ @Context request: HttpServletRequest,
+ likeRequest: UserRequest
+ ): Boolean = {
+ recordLikeAction(request, user.getUid, likeRequest, isLike = true)
+ }
+
+ @POST
+ @Path("/unlike")
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ def postUnlike(
+ @Auth user: SessionUser,
+ @Context request: HttpServletRequest,
+ unlikeRequest: UserRequest
+ ): Boolean = {
+ recordLikeAction(request, user.getUid, unlikeRequest, isLike = false)
+ }
+
+ @POST
+ @Path("/view")
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ def postView(
+ @Context request: HttpServletRequest,
+ viewRequest: ViewRequest
+ ): Int = {
+
+ val (entityID, userId, entityType) =
+ (viewRequest.entityId, viewRequest.userId, viewRequest.entityType)
+
+ val entityTables = ViewCountTable(entityType)
+ val (table, idColumn, viewCountColumn) =
+ (entityTables.table, entityTables.idColumn, entityTables.viewCountColumn)
+
+ val record = context
+ .insertInto(table)
+ .set(idColumn, entityID)
+ .set(viewCountColumn, Integer.valueOf(1))
+ .onDuplicateKeyUpdate()
+ .set(viewCountColumn, viewCountColumn.add(1))
+ .returning(viewCountColumn)
+ .fetchOne()
+
+ recordUserAction(request, userId, entityID, entityType, View)
+
+ record.get(viewCountColumn)
+ }
+
+ /**
+ * Unified endpoint to fetch the top N (here N = 8) public entities for a given entity type,
+ * grouped by specified action types, with optional user context.
+ *
+ * @param entityType The EntityType enum value (Workflow or Dataset) to query.
+ * @param actionTypes Optional list of ActionType enums to include (Like, Clone).
+ * If omitted or empty, defaults to [Like, Clone].
+ * @param uid Optional user ID (Integer) for user-specific context.
+ * If null or -1, no per-user flags are applied.
+ * @param limit Optional maximum number of items to return per action type.
+ * Must be > 0; defaults to 8 if not provided or invalid.
+ * @return A Map from each actionType.value (e.g. "like", "clone")
+ * to a List of DashboardClickableFileEntry containing the top 8
+ * public entities of that type.
+ */
+ @GET
+ @Path("/getTops")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getTops(
+ @QueryParam("entityType") entityType: EntityType,
+ @QueryParam("actionTypes") actionTypes: java.util.List[ActionType],
+ @QueryParam("uid") uid: Integer,
+ @QueryParam("limit") limit: Integer
+ ): java.util.Map[String, java.util.List[DashboardClickableFileEntry]] = {
+ val baseTable = BaseEntityTable(entityType)
+ val isPublicColumn = baseTable.isPublicColumn
+ val baseIdColumn = baseTable.idColumn
+ val topN: Int = Option(limit).filter(_ > 0).map(_.intValue).getOrElse(8)
+
+ val currentUid: Integer =
+ if (uid == null || uid == -1) null
+ else Integer.valueOf(uid)
+
+ val types: Seq[ActionType] =
+ if (actionTypes != null && !actionTypes.isEmpty)
+ actionTypes.asScala.toList.distinct
+ else
+ Seq(ActionType.Like, ActionType.Clone)
+
+ val result: Map[String, java.util.List[DashboardClickableFileEntry]] =
+ types.map { act =>
+ val (table, idColumn) = act match {
+ case ActionType.Like =>
+ val lt = LikeTable(entityType)
+ (lt.table, lt.idColumn)
+ case ActionType.Clone =>
+ val ct = CloneTable(entityType)
+ (ct.table, ct.idColumn)
+ case other =>
+ throw new BadRequestException(
+ s"Unsupported actionType: '$other'. Supported: [like, clone]"
+ )
+ }
+
+ val topIds: Seq[Integer] = context
+ .select(idColumn)
+ .from(table)
+ .join(baseTable.table)
+ .on(idColumn.eq(baseIdColumn))
+ .where(isPublicColumn.eq(true))
+ .groupBy(idColumn)
+ .orderBy(DSL.count(idColumn).desc())
+ .limit(topN)
+ .fetchInto(classOf[Integer])
+ .asScala
+ .toSeq
+
+ val entries: Seq[DashboardClickableFileEntry] =
+ if (entityType == EntityType.Workflow) {
+ fetchDashboardWorkflowsByWids(topIds, currentUid).map { w =>
+ DashboardClickableFileEntry(
+ resourceType = entityType.value,
+ workflow = Some(w),
+ project = None,
+ dataset = None
+ )
+ }
+ } else if (entityType == EntityType.Dataset) {
+ fetchDashboardDatasetsByDids(topIds, currentUid).map { d =>
+ DashboardClickableFileEntry(
+ resourceType = entityType.value,
+ workflow = None,
+ project = None,
+ dataset = Some(d)
+ )
+ }
+ } else {
+ Seq.empty
+ }
+
+ act.value -> entries.toList.asJava
+ }.toMap
+
+ result.asJava
+ }
+
+ /**
+ * Batch endpoint to fetch counts for one or more entities, optionally filtered by action types.
+ *
+ * Example requests:
+ * // All counts for two entities:
+ * // GET /hub/counts?
+ * // entityType=workflow&entityId=123&
+ * // entityType=dataset&entityId=456
+ *
+ * // Only "view" and "like" counts for the same pair:
+ * // GET /hub/counts?
+ * // entityType=workflow&entityId=123&
+ * // entityType=dataset&entityId=456&
+ * // actionType=view&actionType=like
+ *
+ * @param entityTypes List of entity types to query (enum EntityType), e.g. [Workflow, Dataset].
+ * @param entityIds Parallel list of entity IDs, must be the same length as entityTypes.
+ * @param actionTypes (Optional) List of action types to include (enum ActionType).
+ * Supported values: View, Like, Clone, Unlike. If empty or null, all actions are returned.
+ * @return A list of CountResponse objects, one per (entityType, entityId) pair,
+ * each containing the counts for the requested actions.
+ * @throws BadRequestException if entityTypes or entityIds are missing, empty, mismatched in length,
+ * or if actionTypes contains an unsupported value.
+ */
+ @GET
+ @Path("/counts")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getCounts(
+ @QueryParam("entityType") entityTypes: java.util.List[EntityType],
+ @QueryParam("entityId") entityIds: java.util.List[Integer],
+ @QueryParam("actionType") actionTypes: java.util.List[ActionType]
+ ): java.util.List[CountResponse] = {
+ if (
+ entityTypes == null || entityIds == null || entityTypes.isEmpty || entityTypes
+ .size() != entityIds.size()
+ )
+ throw new BadRequestException(
+ "Both 'entityType' and 'entityId' query parameters must be provided, and lists must have equal length."
+ )
+
+ val reqs: List[UserRequest] = entityTypes.asScala
+ .zip(entityIds.asScala)
+ .map {
+ case (etype, id) => UserRequest(id, etype)
+ }
+ .toList
+
+ val requestedActions: Seq[ActionType] =
+ if (actionTypes != null && !actionTypes.isEmpty)
+ actionTypes.asScala.toList.distinct
+ else
+ Seq(ActionType.View, ActionType.Like, ActionType.Clone)
+
+ val grouped: Map[EntityType, Seq[Integer]] =
+ reqs.groupBy(_.entityType).view.mapValues(_.map(_.entityId)).toMap
+
+ val buffer = ListBuffer[CountResponse]()
+
+ grouped.foreach {
+ case (etype, ids) =>
+ val viewTbl = ViewCountTable(etype)
+ val viewMap: Map[Int, Int] =
+ if (requestedActions.contains(ActionType.View)) {
+ val raw = context
+ .select(viewTbl.idColumn, viewTbl.viewCountColumn)
+ .from(viewTbl.table)
+ .where(viewTbl.idColumn.in(ids: _*))
+ .fetchMap(viewTbl.idColumn, viewTbl.viewCountColumn)
+ .asScala
+ .map { case (k, v) => k.intValue() -> v.intValue() }
+ .toMap
+
+ val missing = ids.filterNot(id => raw.contains(id.intValue()))
+
+ missing.foreach { id =>
+ context
+ .insertInto(viewTbl.table)
+ .set(viewTbl.idColumn, id)
+ .set(viewTbl.viewCountColumn, Integer.valueOf(0))
+ .onDuplicateKeyIgnore()
+ .execute()
+ }
+
+ raw ++ missing.map(id => id.intValue() -> 0).toMap
+ } else Map.empty
+
+ val likeTbl = LikeTable(etype)
+ val likeMap: Map[Int, Int] =
+ if (requestedActions.contains(ActionType.Like)) {
+ context
+ .select(likeTbl.idColumn, DSL.count().`as`("cnt"))
+ .from(likeTbl.table)
+ .where(likeTbl.idColumn.in(ids: _*))
+ .groupBy(likeTbl.idColumn)
+ .fetch()
+ .asScala
+ .map { r =>
+ r.get(likeTbl.idColumn).intValue() ->
+ r.get("cnt", classOf[Integer]).intValue()
+ }
+ .toMap
+ } else Map.empty
+
+ val cloneMap: Map[Int, Int] =
+ if (requestedActions.contains(ActionType.Clone) && etype != EntityType.Dataset) {
+ val cloneTbl = CloneTable(etype)
+ context
+ .select(cloneTbl.idColumn, DSL.count().`as`("cnt"))
+ .from(cloneTbl.table)
+ .where(cloneTbl.idColumn.in(ids: _*))
+ .groupBy(cloneTbl.idColumn)
+ .fetch()
+ .asScala
+ .map { r =>
+ r.get(cloneTbl.idColumn).intValue() ->
+ r.get("cnt", classOf[Integer]).intValue()
+ }
+ .toMap
+ } else Map.empty
+
+ reqs.filter(_.entityType == etype).foreach { req =>
+ val key = req.entityId.intValue()
+ val counts = scala.collection.mutable.Map[ActionType, Int]()
+ if (requestedActions.contains(ActionType.View))
+ counts(ActionType.View) = viewMap.getOrElse(key, 0)
+ if (requestedActions.contains(ActionType.Like))
+ counts(ActionType.Like) = likeMap.getOrElse(key, 0)
+ if (requestedActions.contains(ActionType.Clone))
+ counts(ActionType.Clone) = cloneMap.getOrElse(key, 0)
+
+ buffer += CountResponse(req.entityId, etype, counts.asJava)
+ }
+ }
+
+ buffer.toList.asJava
+ }
+
+ /**
+ * Batch-fetches the list of user IDs who have access rights for one or more entities.
+ * Supports multiple entityType/entityId pairs in a single request.
+ *
+ * @param entityTypes List of entity types (e.g. Workflow, Dataset) matching the entityIds.
+ * @param entityIds List of entity IDs matching the entityTypes.
+ * @return A list of AccessResponse objects, each containing:
+ * - entityType: the resource type
+ * - entityId: the resource ID
+ * - userIds: the list of user IDs with access to that resource
+ */
+ @GET
+ @Path("/user-access")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def userAccess(
+ @QueryParam("entityType") entityTypes: java.util.List[EntityType],
+ @QueryParam("entityId") entityIds: java.util.List[Integer]
+ ): java.util.List[AccessResponse] = {
+ val reqs =
+ entityIds.asScala
+ .zip(entityTypes.asScala)
+ .map { case (et, id) => UserRequest(et, id) }
+ .toList
+
+ val responses = ListBuffer[AccessResponse]()
+ reqs.groupBy(_.entityType).foreach {
+ case (etype, groupReqs) =>
+ val (tbl, idCol, uidCol) = etype match {
+ case EntityType.Workflow =>
+ (WORKFLOW_USER_ACCESS: Table[_], WORKFLOW_USER_ACCESS.WID, WORKFLOW_USER_ACCESS.UID)
+ case EntityType.Dataset =>
+ (DATASET_USER_ACCESS: Table[_], DATASET_USER_ACCESS.DID, DATASET_USER_ACCESS.UID)
+ }
+
+ val records = context
+ .select(idCol, uidCol)
+ .from(tbl)
+ .where(idCol.in(groupReqs.map(_.entityId).asJava))
+ .fetch()
+ .asScala
+
+ val accessMap =
+ records
+ .groupBy(r => r.get(idCol))
+ .map {
+ case (id, rs) =>
+ id -> rs.map(r => r.get(uidCol)).toList
+ }
+
+ groupReqs.map(_.entityId).distinct.foreach { eid =>
+ val uids = accessMap.getOrElse(eid, Nil).asJava
+ responses += AccessResponse(etype, eid, uids)
+ }
+ }
+
+ responses.toList.asJava
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/UserResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/UserResource.scala
new file mode 100644
index 00000000000..fd73f0ff4e5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/UserResource.scala
@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user
+
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.tables.daos.UserDao
+import javax.ws.rs._
+import javax.ws.rs.core.{MediaType, Response}
+
+case class RegistrationUpdateRequest(uid: Int, affiliation: String, joiningReason: String)
+
+object UserResource {
+ private def context = SqlServer.getInstance().createDSLContext()
+ private def userDao = new UserDao(context.configuration)
+}
+
+@Path("/user")
+class UserResource {
+
+ /**
+ * Checks whether the user needs to submit joining reason.
+ * null: never prompted, need to prompt -> return true
+ * not null: already prompted, no need to prompt -> return false
+ * @param uid: user id
+ * @return boolean value to whether prompt user to enter joining reason or not
+ */
+ @GET
+ @Path("/joining-reason/required")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def isJoiningReasonRequired(@QueryParam("uid") uid: Int): java.lang.Boolean = {
+ val user = UserResource.userDao.fetchOneByUid(uid)
+ if (user == null) {
+ throw new WebApplicationException("User not found", Response.Status.NOT_FOUND)
+ }
+ java.lang.Boolean.valueOf(user.getJoiningReason == null)
+ }
+
+ /**
+ * Updates the user's affiliation and joining reason.
+ * This is required and cannot be blank.
+ * @param request: provides uid, affiliation and joining reason
+ */
+ @PUT
+ @Path("/joining-reason")
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ def updateJoiningReason(request: RegistrationUpdateRequest): Unit = {
+ val affiliation = Option(request.affiliation).getOrElse("").trim
+ val reason = Option(request.joiningReason).getOrElse("").trim
+
+ if (reason.isEmpty) {
+ throw new WebApplicationException(
+ "Field 'Reason of joining Texera' cannot be empty",
+ Response.Status.BAD_REQUEST
+ )
+ }
+
+ val user = UserResource.userDao.fetchOneByUid(request.uid)
+ user.setAffiliation(affiliation)
+ user.setJoiningReason(reason)
+ UserResource.userDao.update(user)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/dataset/DatasetResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/dataset/DatasetResource.scala
new file mode 100644
index 00000000000..80ec0b9001d
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/dataset/DatasetResource.scala
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.dataset
+
+import org.apache.texera.dao.jooq.generated.tables.pojos.Dataset
+import org.jooq.EnumType
+
+object DatasetResource {
+ // TODO: move these community resource definitions to a centralized package, similar to workflow-core
+ case class DashboardDataset(
+ dataset: Dataset,
+ ownerEmail: String,
+ accessPrivilege: EnumType,
+ isOwner: Boolean,
+ size: Long
+ )
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/dataset/utils/DatasetStatisticsUtils.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/dataset/utils/DatasetStatisticsUtils.scala
new file mode 100644
index 00000000000..01a949e94c5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/dataset/utils/DatasetStatisticsUtils.scala
@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.dataset.utils
+
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.tables.Dataset.DATASET
+import org.apache.texera.web.resource.dashboard.user.quota.UserQuotaResource.DatasetQuota
+
+import scala.jdk.CollectionConverters._
+
+object DatasetStatisticsUtils {
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ // this function retrieves the total counts of dataset that belongs to the user
+ def getUserCreatedDatasetCount(uid: Integer): Int = {
+ val count = context
+ .selectCount()
+ .from(DATASET)
+ .where(DATASET.OWNER_UID.eq(uid))
+ .fetchOne(0, classOf[Int])
+
+ count
+ }
+
+ // this function would return a list of dataset ids that belongs to the user
+ private def getUserCreatedDatasetList(uid: Integer): List[DatasetQuota] = {
+ val result = context
+ .select(
+ DATASET.DID,
+ DATASET.NAME,
+ DATASET.CREATION_TIME
+ )
+ .from(DATASET)
+ .where(DATASET.OWNER_UID.eq(uid))
+ .fetch()
+
+ result.asScala
+ .map(record =>
+ DatasetQuota(
+ did = record.getValue(DATASET.DID),
+ name = record.getValue(DATASET.NAME),
+ creationTime = record.getValue(DATASET.CREATION_TIME).getTime,
+ size = 0
+ )
+ )
+ .toList
+ }
+
+ def getUserCreatedDatasets(uid: Integer): List[DatasetQuota] = {
+ val datasetList = getUserCreatedDatasetList(uid)
+ datasetList.map { dataset =>
+ val size = 0 // we disabled the size calculation due to the switch of dataset implementation
+ dataset.copy(size = size)
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/ProjectAccessResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/ProjectAccessResource.scala
new file mode 100644
index 00000000000..1e3340973da
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/ProjectAccessResource.scala
@@ -0,0 +1,176 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.project
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables.{
+ DATASET_USER_ACCESS,
+ PROJECT_USER_ACCESS,
+ USER,
+ WORKFLOW_USER_ACCESS
+}
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.{ProjectDao, ProjectUserAccessDao, UserDao}
+import org.apache.texera.dao.jooq.generated.tables.pojos.ProjectUserAccess
+import org.apache.texera.web.model.common.AccessEntry
+import org.apache.texera.web.resource.dashboard.user.project.ProjectAccessResource.userHasWriteAccess
+import org.jooq.DSLContext
+
+import java.util
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+
+object ProjectAccessResource {
+ private def context: DSLContext =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+
+ def userHasWriteAccess(pid: Integer, uid: Integer): Boolean = {
+ getProjectAccessPrivilege(pid, uid) == PrivilegeEnum.WRITE
+ }
+
+ def getProjectAccessPrivilege(pid: Integer, uid: Integer): PrivilegeEnum = {
+ Option(
+ context
+ .select(PROJECT_USER_ACCESS.PRIVILEGE)
+ .from(WORKFLOW_USER_ACCESS)
+ .where(
+ PROJECT_USER_ACCESS.PID
+ .eq(pid)
+ .and(DATASET_USER_ACCESS.UID.eq(uid))
+ )
+ .fetchOneInto(classOf[PrivilegeEnum])
+ ).getOrElse(PrivilegeEnum.NONE)
+ }
+}
+
+@Produces(Array(MediaType.APPLICATION_JSON))
+@RolesAllowed(Array("REGULAR", "ADMIN"))
+@Path("/access/project")
+class ProjectAccessResource() {
+ private def context: DSLContext =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def userDao = new UserDao(context.configuration())
+ private def projectDao = new ProjectDao(context.configuration)
+ private def projectUserAccessDao = new ProjectUserAccessDao(context.configuration)
+
+ /**
+ * This method returns the owner of a project
+ *
+ * @param pid , project id
+ * @return ownerEmail, the owner's email
+ */
+ @GET
+ @Path("/owner/{pid}")
+ def getOwner(@PathParam("pid") pid: Integer): String = {
+ userDao.fetchOneByUid(projectDao.fetchOneByPid(pid).getOwnerId).getEmail
+ }
+
+ /**
+ * Returns information about all current shared access of the given project
+ *
+ * @param pid project id
+ * @return a List of email/permission pair
+ */
+ @GET
+ @Path("/list/{pid}")
+ def getAccessList(
+ @PathParam("pid") pid: Integer
+ ): util.List[AccessEntry] = {
+ context
+ .select(
+ USER.EMAIL,
+ USER.NAME,
+ PROJECT_USER_ACCESS.PRIVILEGE
+ )
+ .from(PROJECT_USER_ACCESS)
+ .join(USER)
+ .on(USER.UID.eq(PROJECT_USER_ACCESS.UID))
+ .where(
+ PROJECT_USER_ACCESS.PID
+ .eq(pid)
+ .and(PROJECT_USER_ACCESS.UID.notEqual(projectDao.fetchOneByPid(pid).getOwnerId))
+ )
+ .fetchInto(classOf[AccessEntry])
+ }
+
+ /**
+ * This method shares a project to a user with a specific access type
+ *
+ * @param pid the given project
+ * @param email the email which the access is given to
+ * @param privilege the type of Access given to the target user
+ * @return rejection if user not permitted to share the project or Success Message
+ */
+ @PUT
+ @Path("/grant/{pid}/{email}/{privilege}")
+ def grantAccess(
+ @PathParam("pid") pid: Integer,
+ @PathParam("email") email: String,
+ @PathParam("privilege") privilege: String,
+ @Auth user: SessionUser
+ ): Unit = {
+ if (!userHasWriteAccess(pid, user.getUid)) {
+ throw new ForbiddenException(s"You do not have permission to modify project $pid")
+ }
+
+ projectUserAccessDao.merge(
+ new ProjectUserAccess(
+ userDao.fetchOneByEmail(email).getUid,
+ pid,
+ PrivilegeEnum.valueOf(privilege)
+ )
+ )
+ }
+
+ /**
+ * Revoke a user's access to a file
+ *
+ * @param pid the id of the file
+ * @param email the email of target user whose access is about to be revoked
+ * @return A successful resp if granted, failed resp otherwise
+ */
+ @DELETE
+ @Path("/revoke/{pid}/{email}")
+ def revokeAccess(
+ @PathParam("pid") pid: Integer,
+ @PathParam("email") email: String,
+ @Auth user: SessionUser
+ ): Unit = {
+ if (!userHasWriteAccess(pid, user.getUid)) {
+ throw new ForbiddenException(s"You do not have permission to modify project $pid")
+ }
+
+ context
+ .delete(PROJECT_USER_ACCESS)
+ .where(
+ PROJECT_USER_ACCESS.UID
+ .eq(userDao.fetchOneByEmail(email).getUid)
+ .and(PROJECT_USER_ACCESS.PID.eq(pid))
+ )
+ .execute()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/ProjectResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/ProjectResource.scala
new file mode 100644
index 00000000000..be72fd21d02
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/ProjectResource.scala
@@ -0,0 +1,337 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.project
+
+import io.dropwizard.auth.Auth
+import org.apache.commons.lang3.StringUtils
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.{
+ ProjectDao,
+ ProjectUserAccessDao,
+ WorkflowOfProjectDao
+}
+import org.apache.texera.dao.jooq.generated.tables.pojos._
+import org.apache.texera.web.resource.dashboard.DashboardResource
+import org.apache.texera.web.resource.dashboard.DashboardResource.SearchQueryParams
+import org.apache.texera.web.resource.dashboard.user.project.ProjectResource._
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowAccessResource.hasReadAccess
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource.DashboardWorkflow
+
+import java.sql.Timestamp
+import java.util
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+import scala.jdk.CollectionConverters.IterableHasAsScala
+
+/**
+ * This file handles various request related to projects.
+ * It sends mysql queries to the MysqlDB regarding the 'user_project',
+ * 'workflow_of_project', and 'file_of_project' Tables
+ * The details of these tables can be found in /sql/texera_ddl.sql
+ */
+
+object ProjectResource {
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def userProjectDao = new ProjectDao(context.configuration)
+ private def workflowOfProjectDao = new WorkflowOfProjectDao(context.configuration)
+ private def projectUserAccessDao = new ProjectUserAccessDao(context.configuration)
+
+ /**
+ * This method is used to insert any CSV files created from ResultExportService
+ * handleCSVRequest function into all project(s) that the workflow belongs to.
+ *
+ * No insertion occurs if the workflow does not belong to any projects.
+ *
+ * @param uid user ID
+ * @param wid workflow ID
+ * @param fileName name of exported file
+ * @return String containing status of adding exported file to project(s)
+ */
+ def addExportedFileToProject(uid: Integer, wid: Integer, fileName: String): String = {
+ // get map of PIDs and project names
+ val pidMap = context
+ .select(WORKFLOW_OF_PROJECT.PID, PROJECT.NAME)
+ .from(WORKFLOW_OF_PROJECT)
+ .leftJoin(PROJECT)
+ .on(WORKFLOW_OF_PROJECT.PID.eq(PROJECT.PID))
+ .where(WORKFLOW_OF_PROJECT.WID.eq(wid))
+ .fetch()
+ .intoMap(WORKFLOW_OF_PROJECT.PID, PROJECT.NAME)
+
+ if (pidMap.size() > 0) { // workflow belongs to project(s)
+ // generate string for ResultExportResponse
+ if (pidMap.size() == 1) {
+ s"and added to project: ${pidMap.values().toArray()(0)}"
+ } else {
+ s"and added to projects: ${pidMap.values().asScala.mkString(", ")}"
+ }
+ } else { // workflow does not belong to a project
+ ""
+ }
+ }
+
+ private def workflowOfProjectExists(wid: Integer, pid: Integer): Boolean = {
+ workflowOfProjectDao.existsById(
+ context
+ .newRecord(WORKFLOW_OF_PROJECT.WID, WORKFLOW_OF_PROJECT.PID)
+ .values(wid, pid)
+ )
+ }
+
+ case class DashboardProject(
+ pid: Integer,
+ name: String,
+ description: String,
+ ownerID: Integer,
+ creationTime: Timestamp,
+ color: String,
+ accessLevel: String
+ )
+}
+
+@Path("/project")
+@RolesAllowed(Array("REGULAR", "ADMIN"))
+@Produces(Array(MediaType.APPLICATION_JSON))
+class ProjectResource {
+
+ /**
+ * This method returns the specified project
+ *
+ * @param pid project id
+ * @return project specified by the project id
+ */
+ @GET
+ @Path("/{pid}")
+ def getProject(@PathParam("pid") pid: Integer): Project = {
+ userProjectDao.fetchOneByPid(pid)
+ }
+
+ /**
+ * This method returns the list of projects owned by the session user.
+ *
+ * @param user the session user
+ * @return a list of projects belonging to owner
+ */
+ @GET
+ @Path("/list")
+ def getProjectList(@Auth user: SessionUser): util.List[DashboardProject] = {
+ context
+ .selectDistinct(
+ PROJECT.PID,
+ PROJECT.NAME,
+ PROJECT.DESCRIPTION,
+ PROJECT.OWNER_ID,
+ PROJECT.CREATION_TIME,
+ PROJECT.COLOR,
+ PROJECT_USER_ACCESS.PRIVILEGE
+ )
+ .from(PROJECT_USER_ACCESS)
+ .join(PROJECT)
+ .on(PROJECT_USER_ACCESS.PID.eq(PROJECT.PID))
+ .where(PROJECT.OWNER_ID.eq(user.getUid).or(PROJECT_USER_ACCESS.UID.eq(user.getUid)))
+ .fetchInto(classOf[DashboardProject])
+ }
+
+ /**
+ * This method returns a list of DashboardWorkflow objects, which represents
+ * all the workflows that are part of the specified project.
+ *
+ * @param pid project ID
+ * @param user the session user
+ * @return list of DashboardWorkflow objects
+ */
+ @GET
+ @Path("/{pid}/workflows")
+ def listProjectWorkflows(
+ @PathParam("pid") pid: Integer,
+ @Auth user: SessionUser
+ ): List[DashboardWorkflow] = {
+ val result = DashboardResource.searchAllResources(
+ user,
+ SearchQueryParams(resourceType = "workflow", projectIds = util.Arrays.asList(pid))
+ )
+ result.results.map(_.workflow.get)
+ }
+
+ /**
+ * This method inserts a new project into the database belonging to the session user
+ * and with the specified name.
+ *
+ * @param user the session user
+ * @param name project name
+ */
+ @POST
+ @Path("/create/{name}")
+ def createProject(
+ @Auth user: SessionUser,
+ @PathParam("name") name: String
+ ): Project = {
+ val project = new Project(null, name, null, user.getUid, null, null)
+ try {
+ userProjectDao.insert(project)
+ projectUserAccessDao.merge(
+ new ProjectUserAccess(user.getUid, project.getPid, PrivilegeEnum.WRITE)
+ )
+ } catch {
+ case _: Throwable =>
+ throw new BadRequestException("Cannot create a new project with provided name.");
+ }
+ userProjectDao.fetchOneByPid(project.getPid)
+ }
+
+ /**
+ * This method adds a mapping between the specified workflow to the specified project into the database.
+ *
+ * @param pid project ID
+ * @param wid workflow ID
+ */
+ @POST
+ @Path("/{pid}/workflow/{wid}/add")
+ def addWorkflowToProject(
+ @PathParam("pid") pid: Integer,
+ @PathParam("wid") wid: Integer,
+ @Auth user: SessionUser
+ ): Unit = {
+ if (!hasReadAccess(wid, user.getUid)) {
+ throw new ForbiddenException("No sufficient access privilege to workflow.")
+ }
+
+ if (!workflowOfProjectExists(wid, pid)) {
+ workflowOfProjectDao.insert(new WorkflowOfProject(wid, pid))
+ }
+ }
+
+ /**
+ * This method updates the project name of the specified, existing project
+ *
+ * @param pid project ID
+ * @param name new name
+ */
+ @POST
+ @Path("/{pid}/rename/{name}")
+ def updateProjectName(
+ @PathParam("pid") pid: Integer,
+ @PathParam("name") name: String
+ ): Unit = {
+ val userProject: Project = userProjectDao.fetchOneByPid(pid)
+ if (StringUtils.isBlank(name)) {
+ throw new BadRequestException("Cannot rename project to empty or blank name.")
+ }
+
+ try {
+ userProject.setName(name)
+ userProjectDao.update(userProject)
+ } catch {
+ case _: Throwable => throw new BadRequestException("Cannot rename project to provided name.");
+ }
+ }
+
+ /**
+ * This method updates the description of a specified, existing project
+ *
+ * @param pid project ID
+ */
+ @POST
+ @Path("/{pid}/update/description")
+ @Consumes(Array(MediaType.TEXT_PLAIN))
+ def updateProjectDescription(
+ @PathParam("pid") pid: Integer,
+ description: String
+ ): Unit = {
+ val userProject: Project = userProjectDao.fetchOneByPid(pid)
+ try {
+ userProject.setDescription(description)
+ userProjectDao.update(userProject)
+ } catch {
+ case _: Throwable =>
+ throw new BadRequestException("Cannot update project description to provided text.");
+ }
+ }
+
+ /**
+ * This method updates a project's color.
+ *
+ * @param pid id of project to be updated
+ * @param colorHex new HEX formatted color to be set
+ */
+ @POST
+ @Path("/{pid}/color/{colorHex}/add")
+ def updateProjectColor(
+ @PathParam("pid") pid: Integer,
+ @PathParam("colorHex") colorHex: String,
+ @Auth sessionUser: SessionUser
+ ): Unit = {
+ val userProject: Project = userProjectDao.fetchOneByPid(pid)
+ if (
+ colorHex == null || colorHex.length != 6 && colorHex.length != 3 || !colorHex.matches(
+ "^[A-Fa-f0-9]{6}|[A-Fa-f0-9]{3}$"
+ )
+ ) {
+ throw new BadRequestException("Cannot assign invalid HEX format color to project.")
+ }
+
+ userProject.setColor(colorHex)
+ userProjectDao.update(userProject)
+ }
+
+ @POST
+ @Path("/{pid}/color/delete")
+ def deleteProjectColor(@PathParam("pid") pid: Integer): Unit = {
+ val userProject: Project = userProjectDao.fetchOneByPid(pid)
+ userProject.setColor(null)
+ userProjectDao.update(userProject)
+ }
+
+ /**
+ * This method deletes an existing project from the database
+ *
+ * @param pid projectID
+ */
+ @DELETE
+ @Path("/delete/{pid}")
+ def deleteProject(@PathParam("pid") pid: Integer): Unit = {
+ userProjectDao.deleteById(pid)
+ }
+
+ /**
+ * This method deletes an existing mapping between a workflow and project from
+ * the database
+ *
+ * @param pid project ID
+ * @param wid workflow ID
+ */
+ @DELETE
+ @Path("/{pid}/workflow/{wid}/delete")
+ def deleteWorkflowFromProject(
+ @PathParam("pid") pid: Integer,
+ @PathParam("wid") wid: Integer
+ ): Unit = {
+ workflowOfProjectDao.deleteById(
+ context.newRecord(WORKFLOW_OF_PROJECT.WID, WORKFLOW_OF_PROJECT.PID).values(wid, pid)
+ )
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/PublicProjectResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/PublicProjectResource.scala
new file mode 100644
index 00000000000..6983bf13ee0
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/project/PublicProjectResource.scala
@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.project
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables.{PROJECT, PUBLIC_PROJECT, USER}
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.{ProjectUserAccessDao, PublicProjectDao}
+import org.apache.texera.dao.jooq.generated.tables.pojos.{ProjectUserAccess, PublicProject}
+import org.jooq.DSLContext
+
+import java.sql.Timestamp
+import java.util
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+
+case class DashboardPublicProject(
+ pid: Integer,
+ name: String,
+ owner: String,
+ creationTime: Timestamp
+) {}
+
+@Path("/public/project")
+class PublicProjectResource {
+
+ private def context: DSLContext =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def publicProjectDao = new PublicProjectDao(context.configuration)
+ private def projectUserAccessDao = new ProjectUserAccessDao(context.configuration)
+
+ @GET
+ @RolesAllowed(Array("ADMIN"))
+ @Path("/type/{pid}")
+ def getType(@PathParam("pid") pid: Integer): String = {
+ if (publicProjectDao.fetchOneByPid(pid) == null)
+ "Private"
+ else
+ "Public"
+ }
+
+ @PUT
+ @RolesAllowed(Array("ADMIN"))
+ @Path("/public/{pid}")
+ def makePublic(@PathParam("pid") pid: Integer, @Auth user: SessionUser): Unit = {
+ publicProjectDao.insert(new PublicProject(pid, user.getUid))
+ }
+
+ @PUT
+ @RolesAllowed(Array("ADMIN"))
+ @Path("/private/{pid}")
+ def makePrivate(@PathParam("pid") pid: Integer): Unit = {
+ publicProjectDao.deleteById(pid)
+ }
+
+ @PUT
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/add")
+ def addPublicProjects(checkedList: util.List[Integer], @Auth user: SessionUser): Unit = {
+ checkedList.forEach(pid => {
+ projectUserAccessDao.merge(
+ new ProjectUserAccess(
+ user.getUid,
+ pid,
+ PrivilegeEnum.READ
+ )
+ )
+ })
+ }
+
+ @GET
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/list")
+ def listPublicProjects(): util.List[DashboardPublicProject] = {
+ context
+ .select(PUBLIC_PROJECT.PID, PROJECT.NAME, USER.NAME, PROJECT.CREATION_TIME)
+ .from(PUBLIC_PROJECT)
+ .leftJoin(PROJECT)
+ .on(PUBLIC_PROJECT.PID.eq(PROJECT.PID))
+ .leftJoin(USER)
+ .on(USER.UID.eq(PUBLIC_PROJECT.UID))
+ .fetchInto(classOf[DashboardPublicProject])
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/quota/UserQuotaResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/quota/UserQuotaResource.scala
new file mode 100644
index 00000000000..6f1bb79b150
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/quota/UserQuotaResource.scala
@@ -0,0 +1,224 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.quota
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.web.resource.dashboard.user.dataset.utils.DatasetStatisticsUtils.getUserCreatedDatasets
+import org.apache.texera.web.resource.dashboard.user.quota.UserQuotaResource._
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+
+import java.util
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+import scala.jdk.CollectionConverters.IterableHasAsScala
+
+object UserQuotaResource {
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+
+ case class Workflow(
+ userId: Integer,
+ workflowId: Integer,
+ workflowName: String,
+ creationTime: Long,
+ lastModifiedTime: Long
+ )
+
+ case class DatasetQuota(
+ did: Integer,
+ name: String,
+ creationTime: Long,
+ size: Long
+ )
+
+ case class QuotaStorage(
+ eid: Integer,
+ workflowId: Integer,
+ workflowName: String,
+ resultBytes: Long,
+ runTimeStatsBytes: Long,
+ logBytes: Long
+ )
+
+ def getUserCreatedWorkflow(uid: Integer): List[Workflow] = {
+ val userWorkflowEntries = context
+ .select(
+ WORKFLOW_OF_USER.UID,
+ WORKFLOW_OF_USER.WID,
+ WORKFLOW.NAME,
+ WORKFLOW.CREATION_TIME,
+ WORKFLOW.LAST_MODIFIED_TIME
+ )
+ .from(
+ WORKFLOW_OF_USER
+ )
+ .leftJoin(
+ WORKFLOW
+ )
+ .on(
+ WORKFLOW.WID.eq(WORKFLOW_OF_USER.WID)
+ )
+ .where(
+ WORKFLOW_OF_USER.UID.eq(uid)
+ )
+ .fetch()
+
+ userWorkflowEntries
+ .map(workflowRecord => {
+ Workflow(
+ workflowRecord.get(WORKFLOW_OF_USER.UID),
+ workflowRecord.get(WORKFLOW_OF_USER.WID),
+ workflowRecord.get(WORKFLOW.NAME),
+ workflowRecord.get(WORKFLOW.CREATION_TIME).getTime,
+ workflowRecord.get(WORKFLOW.LAST_MODIFIED_TIME).getTime
+ )
+ })
+ .asScala
+ .toList
+ }
+
+ def getUserAccessedWorkflow(uid: Integer): util.List[Integer] = {
+ val availableWorkflowIds = context
+ .select(
+ WORKFLOW_USER_ACCESS.WID
+ )
+ .from(
+ WORKFLOW_USER_ACCESS
+ )
+ .where(
+ WORKFLOW_USER_ACCESS.UID.eq(uid)
+ )
+ .fetchInto(classOf[Integer])
+
+ availableWorkflowIds
+ }
+
+ def getUserQuotaSize(uid: Integer): Array[QuotaStorage] = {
+ val executions = context
+ .select(
+ WORKFLOW_EXECUTIONS.EID,
+ WORKFLOW_EXECUTIONS.RUNTIME_STATS_SIZE,
+ WORKFLOW.WID,
+ WORKFLOW.NAME
+ )
+ .from(WORKFLOW_EXECUTIONS)
+ .leftJoin(WORKFLOW_VERSION)
+ .on(WORKFLOW_EXECUTIONS.VID.eq(WORKFLOW_VERSION.VID))
+ .leftJoin(WORKFLOW)
+ .on(WORKFLOW_VERSION.WID.eq(WORKFLOW.WID))
+ .where(WORKFLOW_EXECUTIONS.UID.eq(uid))
+ .orderBy(WORKFLOW_EXECUTIONS.EID.desc)
+ .fetch()
+
+ if (executions == null || executions.isEmpty) {
+ return Array.empty
+ }
+
+ executions.asScala.map { record =>
+ val eid = record.get(WORKFLOW_EXECUTIONS.EID)
+ val wid = record.get(WORKFLOW.WID)
+ val workflowName = record.get(WORKFLOW.NAME)
+ val runTimeStatsSize =
+ Option(record.get(WORKFLOW_EXECUTIONS.RUNTIME_STATS_SIZE)).map(_.toLong).getOrElse(0L)
+
+ val resultSize = context
+ .select(OPERATOR_PORT_EXECUTIONS.RESULT_SIZE)
+ .from(OPERATOR_PORT_EXECUTIONS)
+ .where(OPERATOR_PORT_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid))
+ .fetch()
+ .asScala
+ .map(r =>
+ Option(r.get(OPERATOR_PORT_EXECUTIONS.RESULT_SIZE)).getOrElse(0).asInstanceOf[Integer]
+ )
+ .map(_.toLong)
+ .sum
+
+ val logSize = context
+ .select(OPERATOR_EXECUTIONS.CONSOLE_MESSAGES_SIZE)
+ .from(OPERATOR_EXECUTIONS)
+ .where(OPERATOR_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid))
+ .fetch()
+ .asScala
+ .map(r =>
+ Option(r.get(OPERATOR_EXECUTIONS.CONSOLE_MESSAGES_SIZE))
+ .getOrElse(0)
+ .asInstanceOf[Integer]
+ )
+ .map(_.toLong)
+ .sum
+
+ QuotaStorage(
+ eid,
+ wid,
+ workflowName,
+ resultSize,
+ runTimeStatsSize,
+ logSize
+ )
+ }.toArray
+ }
+
+ def deleteExecutionCollection(eid: Integer): Unit = {
+ WorkflowExecutionsResource.removeAllExecutionFiles(Array(eid))
+ }
+}
+
+@Path("/quota")
+class UserQuotaResource {
+
+ @GET
+ @Path("/created_datasets")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getCreatedDatasets(@Auth current_user: SessionUser): List[DatasetQuota] = {
+ getUserCreatedDatasets(current_user.getUid)
+ }
+
+ @GET
+ @Path("/created_workflows")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getCreatedWorkflow(@Auth current_user: SessionUser): List[Workflow] = {
+ getUserCreatedWorkflow(current_user.getUid)
+ }
+
+ @GET
+ @Path("/access_workflows")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getAccessedWorkflow(@Auth current_user: SessionUser): util.List[Integer] = {
+ getUserAccessedWorkflow(current_user.getUid)
+ }
+
+ @GET
+ @Path("/user_quota_size")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getUserQuota(@Auth current_user: SessionUser): Array[QuotaStorage] = {
+ getUserQuotaSize(current_user.getUid)
+ }
+
+ @DELETE
+ @Path("/deleteCollection/{eid}")
+ def deleteCollection(@PathParam("eid") eid: Integer): Unit = {
+ deleteExecutionCollection(eid)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowAccessResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowAccessResource.scala
new file mode 100644
index 00000000000..a439238aae0
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowAccessResource.scala
@@ -0,0 +1,259 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.workflow
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.{
+ UserDao,
+ WorkflowOfUserDao,
+ WorkflowUserAccessDao
+}
+import org.apache.texera.dao.jooq.generated.tables.pojos.WorkflowUserAccess
+import org.apache.texera.web.model.common.AccessEntry
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowAccessResource.{
+ context,
+ getPrivilege,
+ hasWriteAccess
+}
+import org.jooq.DSLContext
+
+import java.util
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+
+object WorkflowAccessResource {
+ private def context: DSLContext =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+
+ /**
+ * Identifies whether the given user has read-only access over the given workflow
+ *
+ * @param wid workflow id
+ * @param uid user id, works with workflow id as primary keys in database
+ * @return boolean value indicating yes/no
+ */
+ def hasReadAccess(wid: Integer, uid: Integer): Boolean = {
+ isPublic(wid) || getPrivilege(wid, uid).eq(PrivilegeEnum.READ) || hasWriteAccess(
+ wid,
+ uid
+ )
+ }
+
+ /**
+ * Identifies whether the given user has write access over the given workflow
+ *
+ * @param wid workflow id
+ * @param uid user id, works with workflow id as primary keys in database
+ * @return boolean value indicating yes/no
+ */
+ def hasWriteAccess(wid: Integer, uid: Integer): Boolean = {
+ getPrivilege(wid, uid).eq(PrivilegeEnum.WRITE)
+ }
+
+ /**
+ * @param wid workflow id
+ * @param uid user id, works with workflow id as primary keys in database
+ * @return PrivilegeEnum value indicating NONE/READ/WRITE
+ */
+ def getPrivilege(wid: Integer, uid: Integer): PrivilegeEnum = {
+ val access = context
+ .select()
+ .from(WORKFLOW_USER_ACCESS)
+ .where(WORKFLOW_USER_ACCESS.WID.eq(wid).and(WORKFLOW_USER_ACCESS.UID.eq(uid)))
+ .fetchOneInto(classOf[WorkflowUserAccess])
+ if (access == null) {
+ val projectAccess = context
+ .select()
+ .from(PROJECT_USER_ACCESS)
+ .join(WORKFLOW_OF_PROJECT)
+ .on(WORKFLOW_OF_PROJECT.PID.eq(PROJECT_USER_ACCESS.PID))
+ .where(WORKFLOW_OF_PROJECT.WID.eq(wid).and(PROJECT_USER_ACCESS.UID.eq(uid)))
+ .fetchOneInto(classOf[WorkflowUserAccess])
+ if (projectAccess == null) {
+ PrivilegeEnum.NONE
+ } else {
+ projectAccess.getPrivilege
+ }
+ } else {
+ access.getPrivilege
+ }
+ }
+
+ def isPublic(wid: Integer): Boolean = {
+ context
+ .select(WORKFLOW.IS_PUBLIC)
+ .from(WORKFLOW)
+ .where(WORKFLOW.WID.eq(wid))
+ .fetchOneInto(classOf[Boolean])
+ }
+}
+
+@Produces(Array(MediaType.APPLICATION_JSON))
+@RolesAllowed(Array("REGULAR", "ADMIN"))
+@Path("/access/workflow")
+class WorkflowAccessResource() {
+ final private val userDao = new UserDao(context.configuration())
+ final private val workflowOfUserDao = new WorkflowOfUserDao(context.configuration)
+ final private val workflowUserAccessDao = new WorkflowUserAccessDao(context.configuration)
+
+ /**
+ * This method returns the owner of a workflow
+ *
+ * @param wid , workflow id
+ * @return ownerEmail, the owner's email
+ */
+ @GET
+ @Path("/owner/{wid}")
+ def getOwner(@PathParam("wid") wid: Integer): String = {
+ userDao.fetchOneByUid(workflowOfUserDao.fetchByWid(wid).get(0).getUid).getEmail
+ }
+
+ /**
+ * Returns information about all current shared access of the given workflow
+ *
+ * @param wid workflow id
+ * @return a List of email/name/permission
+ */
+ @GET
+ @Path("/list/{wid}")
+ def getAccessList(
+ @PathParam("wid") wid: Integer
+ ): util.List[AccessEntry] = {
+ context
+ .select(
+ USER.EMAIL,
+ USER.NAME,
+ WORKFLOW_USER_ACCESS.PRIVILEGE
+ )
+ .from(WORKFLOW_USER_ACCESS)
+ .join(USER)
+ .on(USER.UID.eq(WORKFLOW_USER_ACCESS.UID))
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(wid)
+ .and(WORKFLOW_USER_ACCESS.UID.notEqual(workflowOfUserDao.fetchByWid(wid).get(0).getUid))
+ )
+ .fetchInto(classOf[AccessEntry])
+ }
+
+ /**
+ * This method shares a workflow to a user with a specific access type
+ *
+ * @param wid the given workflow
+ * @param email the email which the access is given to
+ * @param privilege the type of Access given to the target user
+ * @return rejection if user not permitted to share the workflow or Success Message
+ */
+ @PUT
+ @Path("/grant/{wid}/{email}/{privilege}")
+ def grantAccess(
+ @PathParam("wid") wid: Integer,
+ @PathParam("email") email: String,
+ @PathParam("privilege") privilege: String,
+ @Auth user: SessionUser
+ ): Unit = {
+ val isModifyingOwnAccess = email.equals(user.getEmail)
+ val currentPrivilege = getPrivilege(wid, user.getUid)
+ val hasExistingAccess = !currentPrivilege.eq(PrivilegeEnum.NONE)
+
+ // Users can only modify their own access if they already have access
+ if (isModifyingOwnAccess && !hasExistingAccess) {
+ throw new BadRequestException("You cannot grant access to yourself!")
+ }
+
+ // Must have write access to modify access levels (including your own)
+ if (!hasWriteAccess(wid, user.getUid)) {
+ throw new ForbiddenException(s"You do not have permission to modify workflow $wid")
+ }
+
+ val userUid = userDao.fetchOneByEmail(email).getUid
+ val workflowOwnerUid = context
+ .select(WORKFLOW_OF_USER.UID)
+ .from(WORKFLOW_OF_USER)
+ .where(WORKFLOW_OF_USER.WID.eq(wid))
+ .fetchOneInto(classOf[Integer])
+ if (userUid == workflowOwnerUid) {
+ throw new ForbiddenException("You cannot modify the owner's permissions!")
+ }
+
+ try {
+ workflowUserAccessDao.merge(
+ new WorkflowUserAccess(
+ userUid,
+ wid,
+ PrivilegeEnum.valueOf(privilege)
+ )
+ )
+ } catch {
+ case _: NullPointerException =>
+ throw new BadRequestException(s"User $email Not Found!")
+ }
+ }
+
+ /**
+ * This method identifies the user access level of the given workflow
+ *
+ * @param wid the given workflow
+ * @param email the email of the use whose access is about to be removed
+ * @return message indicating a success message
+ */
+ @DELETE
+ @Path("/revoke/{wid}/{email}")
+ def revokeAccess(
+ @PathParam("wid") wid: Integer,
+ @PathParam("email") email: String,
+ @Auth user: SessionUser
+ ): Unit = {
+ try {
+ val targetUserUid = userDao.fetchOneByEmail(email).getUid
+ val workflowOwnerUid = workflowOfUserDao.fetchByWid(wid).get(0).getUid
+
+ // Prevent owner from revoking their own access
+ if (targetUserUid == workflowOwnerUid) {
+ throw new ForbiddenException("The owner cannot revoke their own access")
+ }
+
+ // Allow if: (1) user has WRITE access, OR (2) user is revoking their own access
+ val isRevokingOwnAccess = targetUserUid == user.getUid
+ if (!hasWriteAccess(wid, user.getUid) && !isRevokingOwnAccess) {
+ throw new ForbiddenException(s"You do not have permission to modify workflow $wid")
+ }
+
+ context
+ .delete(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.UID
+ .eq(targetUserUid)
+ .and(WORKFLOW_USER_ACCESS.WID.eq(wid))
+ )
+ .execute()
+ } catch {
+ case _: NullPointerException =>
+ throw new BadRequestException(s"User $email Not Found!")
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowExecutionsResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowExecutionsResource.scala
new file mode 100644
index 00000000000..72fb1c364e5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowExecutionsResource.scala
@@ -0,0 +1,882 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.workflow
+
+import io.dropwizard.auth.Auth
+import org.apache.texera.amber.core.storage.{
+ DocumentFactory,
+ FileResolver,
+ VFSResourceType,
+ VFSURIFactory
+}
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity._
+import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PortIdentity}
+import org.apache.texera.amber.engine.architecture.logreplay.{ReplayDestination, ReplayLogRecord}
+import org.apache.texera.amber.engine.common.Utils.{maptoStatusCode, stringToAggregatedState}
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.util.JSONUtils.objectMapper
+import org.apache.texera.amber.util.serde.GlobalPortIdentitySerde.SerdeOps
+import org.apache.texera.auth.{JwtParser, SessionUser}
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.SqlServer.withTransaction
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.WorkflowExecutionsDao
+import org.apache.texera.dao.jooq.generated.tables.pojos.{WorkflowExecutions, User => UserPojo}
+import org.apache.texera.web.model.http.request.result.ResultExportRequest
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource._
+import org.apache.texera.web.service.{ExecutionsMetadataPersistService, ResultExportService}
+import org.jooq.DSLContext
+import play.api.libs.json.Json
+
+import java.net.URI
+import java.sql.Timestamp
+import java.util.concurrent.TimeUnit
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.{MediaType, Response}
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+object WorkflowExecutionsResource {
+ private def context: DSLContext =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def executionsDao = new WorkflowExecutionsDao(context.configuration)
+
+ private def getExecutionById(eId: Integer): WorkflowExecutions = {
+ executionsDao.fetchOneByEid(eId)
+ }
+
+ def getExpiredExecutionsWithResultOrLog(timeToLive: Int): List[WorkflowExecutions] = {
+ val deadline = new Timestamp(
+ System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(timeToLive)
+ )
+ context
+ .selectFrom(WORKFLOW_EXECUTIONS)
+ .where(
+ WORKFLOW_EXECUTIONS.LAST_UPDATE_TIME.isNull
+ .and(WORKFLOW_EXECUTIONS.STARTING_TIME.lt(deadline))
+ .or(WORKFLOW_EXECUTIONS.LAST_UPDATE_TIME.lt(deadline))
+ )
+ .and(
+ WORKFLOW_EXECUTIONS.RESULT.ne("").or(WORKFLOW_EXECUTIONS.LOG_LOCATION.ne(""))
+ )
+ .fetchInto(classOf[WorkflowExecutions])
+ .asScala
+ .toList
+ }
+
+ /**
+ * This function retrieves the latest execution id of a workflow
+ *
+ * @param wid workflow id
+ * @return Integer
+ */
+ def getLatestExecutionID(wid: Integer, cuid: Integer): Option[Integer] = {
+ val executions = context
+ .select(WORKFLOW_EXECUTIONS.EID)
+ .from(WORKFLOW_EXECUTIONS)
+ .join(WORKFLOW_VERSION)
+ .on(WORKFLOW_EXECUTIONS.VID.eq(WORKFLOW_VERSION.VID))
+ .where(WORKFLOW_VERSION.WID.eq(wid).and(WORKFLOW_EXECUTIONS.CUID.eq(cuid)))
+ .fetchInto(classOf[Integer])
+ .asScala
+ .toList
+ if (executions.isEmpty) {
+ None
+ } else {
+ Some(executions.max)
+ }
+ }
+
+ /**
+ * Computes which operators in a workflow are restricted due to dataset access controls.
+ *
+ * This function:
+ * 1. Parses the workflow JSON to find all operators and their dataset dependencies
+ * 2. Identifies operators using non-downloadable datasets that the user doesn't own
+ * 3. Uses BFS to propagate restrictions through the workflow graph
+ * 4. Returns a map of operator IDs to the restricted datasets they depend on
+ *
+ * @param wid The workflow ID
+ * @param currentUser The current user making the export request
+ * @return Map of operator ID -> Set of (ownerEmail, datasetName) tuples that block its export
+ */
+ private def getNonDownloadableOperatorMap(
+ wid: Int,
+ currentUser: UserPojo
+ ): Map[String, Set[(String, String)]] = {
+ // Load workflow
+ val workflowRecord = context
+ .select(WORKFLOW.CONTENT)
+ .from(WORKFLOW)
+ .where(WORKFLOW.WID.eq(wid).and(WORKFLOW.CONTENT.isNotNull).and(WORKFLOW.CONTENT.ne("")))
+ .fetchOne()
+
+ if (workflowRecord == null) {
+ return Map.empty
+ }
+
+ val content = workflowRecord.value1()
+
+ val rootNode =
+ try {
+ objectMapper.readTree(content)
+ } catch {
+ case _: Exception => return Map.empty
+ }
+
+ val operatorsNode = rootNode.path("operators")
+ val linksNode = rootNode.path("links")
+
+ // Collect all datasets used by operators (that user doesn't own)
+ val operatorDatasets = mutable.Map.empty[String, (String, String)]
+
+ operatorsNode.elements().asScala.foreach { operatorNode =>
+ val operatorId = operatorNode.path("operatorID").asText("")
+ if (operatorId.nonEmpty) {
+ val fileNameNode = operatorNode.path("operatorProperties").path("fileName")
+ if (fileNameNode.isTextual) {
+ FileResolver.parseDatasetOwnerAndName(fileNameNode.asText()).foreach {
+ case (ownerEmail, datasetName) =>
+ val isOwner =
+ Option(currentUser.getEmail)
+ .exists(_.equalsIgnoreCase(ownerEmail))
+ if (!isOwner) {
+ operatorDatasets.update(operatorId, (ownerEmail, datasetName))
+ }
+ }
+ }
+ }
+ }
+
+ if (operatorDatasets.isEmpty) {
+ return Map.empty
+ }
+
+ // Query all datasets
+ val uniqueDatasets = operatorDatasets.values.toSet
+ val conditions = uniqueDatasets.map {
+ case (ownerEmail, datasetName) =>
+ USER.EMAIL.equalIgnoreCase(ownerEmail).and(DATASET.NAME.equalIgnoreCase(datasetName))
+ }
+
+ val nonDownloadableDatasets = context
+ .select(USER.EMAIL, DATASET.NAME)
+ .from(DATASET)
+ .join(USER)
+ .on(DATASET.OWNER_UID.eq(USER.UID))
+ .where(conditions.reduce((a, b) => a.or(b)))
+ .and(DATASET.IS_DOWNLOADABLE.eq(false))
+ .fetch()
+ .asScala
+ .map(record => (record.value1(), record.value2()))
+ .toSet
+
+ // Filter to only operators with non-downloadable datasets
+ val restrictedSourceMap = operatorDatasets.filter {
+ case (_, dataset) =>
+ nonDownloadableDatasets.contains(dataset)
+ }
+
+ // Build dependency graph
+ val adjacency = mutable.Map.empty[String, mutable.ListBuffer[String]]
+
+ linksNode.elements().asScala.foreach { linkNode =>
+ val sourceId = linkNode.path("source").path("operatorID").asText("")
+ val targetId = linkNode.path("target").path("operatorID").asText("")
+ if (sourceId.nonEmpty && targetId.nonEmpty) {
+ adjacency.getOrElseUpdate(sourceId, mutable.ListBuffer.empty[String]) += targetId
+ }
+ }
+
+ // BFS to propagate restrictions
+ val restrictionMap = mutable.Map.empty[String, Set[(String, String)]]
+ val queue = mutable.Queue.empty[(String, Set[(String, String)])]
+
+ restrictedSourceMap.foreach {
+ case (operatorId, dataset) =>
+ queue.enqueue(operatorId -> Set(dataset))
+ }
+
+ while (queue.nonEmpty) {
+ val (currentOperatorId, datasetSet) = queue.dequeue()
+ val existing = restrictionMap.getOrElse(currentOperatorId, Set.empty)
+ val merged = existing ++ datasetSet
+ if (merged != existing) {
+ restrictionMap.update(currentOperatorId, merged)
+ adjacency
+ .get(currentOperatorId)
+ .foreach(_.foreach(nextOperator => queue.enqueue(nextOperator -> merged)))
+ }
+ }
+
+ restrictionMap.toMap
+ }
+
+ def insertOperatorPortResultUri(
+ eid: ExecutionIdentity,
+ globalPortId: GlobalPortIdentity,
+ uri: URI
+ ): Unit = {
+ context
+ .insertInto(OPERATOR_PORT_EXECUTIONS)
+ .columns(
+ OPERATOR_PORT_EXECUTIONS.WORKFLOW_EXECUTION_ID,
+ OPERATOR_PORT_EXECUTIONS.GLOBAL_PORT_ID,
+ OPERATOR_PORT_EXECUTIONS.RESULT_URI
+ )
+ .values(eid.id.toInt, globalPortId.serializeAsString, uri.toString)
+ .execute()
+ }
+
+ def insertOperatorExecutions(
+ eid: Long,
+ opId: String,
+ uri: URI
+ ): Unit = {
+ context
+ .insertInto(OPERATOR_EXECUTIONS)
+ .columns(
+ OPERATOR_EXECUTIONS.WORKFLOW_EXECUTION_ID,
+ OPERATOR_EXECUTIONS.OPERATOR_ID,
+ OPERATOR_EXECUTIONS.CONSOLE_MESSAGES_URI
+ )
+ .values(eid.toInt, opId, uri.toString)
+ .execute()
+ }
+
+ def updateRuntimeStatsUri(wid: Long, eid: Long, uri: URI): Unit = {
+ context
+ .update(WORKFLOW_EXECUTIONS)
+ .set(WORKFLOW_EXECUTIONS.RUNTIME_STATS_URI, uri.toString)
+ .where(
+ WORKFLOW_EXECUTIONS.EID
+ .eq(eid.toInt)
+ .and(
+ WORKFLOW_EXECUTIONS.VID.in(
+ context
+ .select(WORKFLOW_VERSION.VID)
+ .from(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.WID.eq(wid.toInt))
+ )
+ )
+ )
+ .execute()
+ }
+
+ def getResultUrisByExecutionId(eid: ExecutionIdentity): List[URI] = {
+ context
+ .select(OPERATOR_PORT_EXECUTIONS.RESULT_URI)
+ .from(OPERATOR_PORT_EXECUTIONS)
+ .where(OPERATOR_PORT_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .fetchInto(classOf[String])
+ .asScala
+ .toList
+ .filter(uri => uri != null && uri.nonEmpty)
+ .map(URI.create)
+ }
+
+ def getConsoleMessagesUriByExecutionId(eid: ExecutionIdentity): List[URI] =
+ context
+ .select(OPERATOR_EXECUTIONS.CONSOLE_MESSAGES_URI)
+ .from(OPERATOR_EXECUTIONS)
+ .where(OPERATOR_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .fetchInto(classOf[String])
+ .asScala
+ .toList
+ .filter(uri => uri != null && uri.nonEmpty)
+ .map(URI.create)
+
+ def getRuntimeStatsUriByExecutionId(eid: ExecutionIdentity): Option[URI] =
+ Option(
+ context
+ .select(WORKFLOW_EXECUTIONS.RUNTIME_STATS_URI)
+ .from(WORKFLOW_EXECUTIONS)
+ .where(WORKFLOW_EXECUTIONS.EID.eq(eid.id.toInt))
+ .fetchOneInto(classOf[String])
+ ).filter(_.nonEmpty)
+ .map(URI.create)
+
+ def getWorkflowExecutions(
+ wid: Integer,
+ context: DSLContext,
+ statusCodes: Set[Byte] = Set.empty
+ ): List[WorkflowExecutionEntry] = {
+ var condition = WORKFLOW_VERSION.WID.eq(wid)
+
+ if (statusCodes.nonEmpty) {
+ condition = condition.and(
+ WORKFLOW_EXECUTIONS.STATUS.in(statusCodes.map(Byte.box).asJava)
+ )
+ }
+
+ context
+ .select(
+ WORKFLOW_EXECUTIONS.EID,
+ WORKFLOW_EXECUTIONS.VID,
+ WORKFLOW_EXECUTIONS.CUID,
+ USER.NAME,
+ USER.GOOGLE_AVATAR,
+ WORKFLOW_EXECUTIONS.STATUS,
+ WORKFLOW_EXECUTIONS.RESULT,
+ WORKFLOW_EXECUTIONS.STARTING_TIME,
+ WORKFLOW_EXECUTIONS.LAST_UPDATE_TIME,
+ WORKFLOW_EXECUTIONS.BOOKMARKED,
+ WORKFLOW_EXECUTIONS.NAME,
+ WORKFLOW_EXECUTIONS.LOG_LOCATION
+ )
+ .from(WORKFLOW_EXECUTIONS)
+ .join(WORKFLOW_VERSION)
+ .on(WORKFLOW_VERSION.VID.eq(WORKFLOW_EXECUTIONS.VID))
+ .join(USER)
+ .on(WORKFLOW_EXECUTIONS.UID.eq(USER.UID))
+ .where(condition)
+ .orderBy(WORKFLOW_EXECUTIONS.EID.desc())
+ .fetchInto(classOf[WorkflowExecutionEntry])
+ .asScala
+ .toList
+ }
+
+ def deleteConsoleMessageAndExecutionResultUris(eid: ExecutionIdentity): Unit = {
+ context
+ .delete(OPERATOR_PORT_EXECUTIONS)
+ .where(OPERATOR_PORT_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .execute()
+ context
+ .delete(OPERATOR_EXECUTIONS)
+ .where(OPERATOR_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .execute()
+ }
+
+ /**
+ * Removes all resources related to the specified execution IDs,
+ * including runtime statistics, console messages, result documents, and database records.
+ *
+ * @param eids Array of execution IDs to be cleaned up.
+ */
+ def removeAllExecutionFiles(eids: Array[Integer]): Unit = {
+ val eIdsLong = eids.map(_.toLong)
+ val eIdsList = eIdsLong.toSeq.asJava
+
+ // Collect all related document URIs (runtime stats, console logs, results)
+ val uris: Seq[URI] = eIdsLong.toIndexedSeq.flatMap { eid =>
+ val execId = ExecutionIdentity(eid)
+ WorkflowExecutionsResource
+ .getRuntimeStatsUriByExecutionId(execId)
+ .toList ++
+ WorkflowExecutionsResource.getConsoleMessagesUriByExecutionId(execId) ++
+ WorkflowExecutionsResource.getResultUrisByExecutionId(execId)
+ }
+
+ // Delete execution-related URIs from database tables
+ context
+ .deleteFrom(WORKFLOW_EXECUTIONS)
+ .where(WORKFLOW_EXECUTIONS.EID.in(eIdsList))
+ .execute()
+
+ // Clear corresponding Iceberg documents
+ uris.foreach { uri =>
+ try {
+ DocumentFactory.openDocument(uri)._1.clear()
+ } catch {
+ case _: Throwable =>
+ // Document already deleted – safe to ignore
+ }
+ }
+ }
+
+ /**
+ * Updates the result size of the corresponding Iceberg document in the database.
+ *
+ * @param eid Execution ID associated with the result.
+ * @param globalPortId Global port identifier for the operator output.
+ * @param size Size of the result in bytes.
+ */
+ def updateResultSize(
+ eid: ExecutionIdentity,
+ globalPortId: GlobalPortIdentity,
+ size: Long
+ ): Unit = {
+ context
+ .update(OPERATOR_PORT_EXECUTIONS)
+ .set(OPERATOR_PORT_EXECUTIONS.RESULT_SIZE, Integer.valueOf(size.toInt))
+ .where(OPERATOR_PORT_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .and(OPERATOR_PORT_EXECUTIONS.GLOBAL_PORT_ID.eq(globalPortId.serializeAsString))
+ .execute()
+ }
+
+ /**
+ * Updates the size of the runtime statistics stored via Iceberg document.
+ *
+ * @param eid Execution ID associated with the runtime statistics document.
+ */
+ def updateRuntimeStatsSize(eid: ExecutionIdentity): Unit = {
+ val statsUriOpt = context
+ .select(WORKFLOW_EXECUTIONS.RUNTIME_STATS_URI)
+ .from(WORKFLOW_EXECUTIONS)
+ .where(WORKFLOW_EXECUTIONS.EID.eq(eid.id.toInt))
+ .fetchOptionalInto(classOf[String])
+ .map(URI.create)
+
+ if (statsUriOpt.isPresent) {
+ val size = DocumentFactory.openDocument(statsUriOpt.get)._1.getTotalFileSize
+ context
+ .update(WORKFLOW_EXECUTIONS)
+ .set(WORKFLOW_EXECUTIONS.RUNTIME_STATS_SIZE, Integer.valueOf(size.toInt))
+ .where(WORKFLOW_EXECUTIONS.EID.eq(eid.id.toInt))
+ .execute()
+ }
+ }
+
+ /**
+ * Updates the size of the console message stored via Iceberg document.
+ *
+ * @param eid Execution ID associated with the console message.
+ * @param opId Operator ID of the corresponding operator.
+ */
+ def updateConsoleMessageSize(eid: ExecutionIdentity, opId: OperatorIdentity): Unit = {
+ val uriOpt = context
+ .select(OPERATOR_EXECUTIONS.CONSOLE_MESSAGES_URI)
+ .from(OPERATOR_EXECUTIONS)
+ .where(OPERATOR_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .and(OPERATOR_EXECUTIONS.OPERATOR_ID.eq(opId.id))
+ .fetchOptionalInto(classOf[String])
+ .map(URI.create)
+
+ if (uriOpt.isPresent) {
+ val size = DocumentFactory.openDocument(uriOpt.get)._1.getTotalFileSize
+ context
+ .update(OPERATOR_EXECUTIONS)
+ .set(OPERATOR_EXECUTIONS.CONSOLE_MESSAGES_SIZE, Integer.valueOf(size.toInt))
+ .where(OPERATOR_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .and(OPERATOR_EXECUTIONS.OPERATOR_ID.eq(opId.id))
+ .execute()
+ }
+ }
+
+ /**
+ * This method is mainly used for frontend requests. Given a logicalOpId and an outputPortId of an execution,
+ * this method finds the URI for a globalPortId that both: 1. matches the logicalOpId and outputPortId, and
+ * 2. is an external port. Currently the lookup is O(n), where n is the number of globalPortIds for this execution.
+ * TODO: Optimize the lookup once the frontend also has information about physical operators.
+ */
+ def getResultUriByLogicalPortId(
+ eid: ExecutionIdentity,
+ opId: OperatorIdentity,
+ portId: PortIdentity
+ ): Option[URI] = {
+ def isMatchingExternalPortURI(uri: URI): Boolean = {
+ val (_, _, globalPortIdOption, resourceType) = VFSURIFactory.decodeURI(uri)
+ globalPortIdOption.exists { globalPortId =>
+ !globalPortId.portId.internal &&
+ globalPortId.opId.logicalOpId == opId &&
+ globalPortId.portId == portId &&
+ resourceType == VFSResourceType.RESULT
+ }
+ }
+
+ val urisOfEid: List[URI] =
+ context
+ .select(OPERATOR_PORT_EXECUTIONS.RESULT_URI)
+ .from(OPERATOR_PORT_EXECUTIONS)
+ .where(OPERATOR_PORT_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(eid.id.toInt))
+ .fetchInto(classOf[String])
+ .asScala
+ .toList
+ .map(URI.create)
+
+ urisOfEid.find(isMatchingExternalPortURI)
+ }
+
+ case class WorkflowExecutionEntry(
+ eId: Integer,
+ vId: Integer,
+ cuId: Integer,
+ userName: String,
+ googleAvatar: String,
+ status: Byte,
+ result: String,
+ startingTime: Timestamp,
+ completionTime: Timestamp,
+ bookmarked: Boolean,
+ name: String,
+ logLocation: String
+ )
+
+ case class WorkflowRuntimeStatistics(
+ operatorId: String,
+ timestamp: Timestamp,
+ inputTupleCount: Long,
+ inputTupleSize: Long,
+ outputTupleCount: Long,
+ outputTupleSize: Long,
+ dataProcessingTime: Long,
+ controlProcessingTime: Long,
+ idleTime: Long,
+ numWorkers: Int,
+ status: Int
+ )
+}
+
+case class ExecutionGroupBookmarkRequest(
+ wid: Integer,
+ eIds: Array[Integer],
+ isBookmarked: Boolean
+)
+
+case class ExecutionGroupDeleteRequest(wid: Integer, eIds: Array[Integer])
+
+case class ExecutionRenameRequest(wid: Integer, eId: Integer, executionName: String)
+
+@Produces(Array(MediaType.APPLICATION_JSON, MediaType.APPLICATION_OCTET_STREAM, "application/zip"))
+@Path("/executions")
+class WorkflowExecutionsResource {
+
+ @GET
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @Path("/{wid}/latest")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def retrieveLatestExecutionEntry(
+ @PathParam("wid") wid: Integer,
+ @Auth sessionUser: SessionUser
+ ): WorkflowExecutionEntry = {
+
+ validateUserCanAccessWorkflow(sessionUser.getUser.getUid, wid)
+
+ withTransaction(context) { ctx =>
+ val latestEntryOpt =
+ ctx
+ .select(
+ WORKFLOW_EXECUTIONS.EID,
+ WORKFLOW_EXECUTIONS.VID,
+ WORKFLOW_EXECUTIONS.CUID,
+ USER.NAME,
+ USER.GOOGLE_AVATAR,
+ WORKFLOW_EXECUTIONS.STATUS,
+ WORKFLOW_EXECUTIONS.RESULT,
+ WORKFLOW_EXECUTIONS.STARTING_TIME,
+ WORKFLOW_EXECUTIONS.LAST_UPDATE_TIME,
+ WORKFLOW_EXECUTIONS.BOOKMARKED,
+ WORKFLOW_EXECUTIONS.NAME,
+ WORKFLOW_EXECUTIONS.LOG_LOCATION
+ )
+ .from(WORKFLOW_EXECUTIONS)
+ .join(WORKFLOW_VERSION)
+ .on(WORKFLOW_VERSION.VID.eq(WORKFLOW_EXECUTIONS.VID))
+ .join(USER)
+ .on(WORKFLOW_EXECUTIONS.UID.eq(USER.UID))
+ .where(WORKFLOW_VERSION.WID.eq(wid))
+ // sort by latest VID first, then latest start-time
+ .orderBy(
+ WORKFLOW_EXECUTIONS.VID.desc(),
+ WORKFLOW_EXECUTIONS.EID.desc()
+ )
+ .limit(1)
+ .fetchInto(classOf[WorkflowExecutionEntry])
+ .asScala
+ .headOption
+
+ latestEntryOpt.getOrElse {
+ throw new ForbiddenException("Executions doesn't exist")
+ }
+ }
+ }
+
+ @GET
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @Path("/{wid}/interactions/{eid}")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def retrieveInteractionHistory(
+ @PathParam("wid") wid: Integer,
+ @PathParam("eid") eid: Integer,
+ @Auth sessionUser: SessionUser
+ ): List[String] = {
+ val user = sessionUser.getUser
+ if (!WorkflowAccessResource.hasReadAccess(wid, user.getUid)) {
+ List()
+ } else {
+ ExecutionsMetadataPersistService.tryGetExistingExecution(
+ ExecutionIdentity(eid.longValue())
+ ) match {
+ case Some(value) =>
+ val logLocation = value.getLogLocation
+ if (logLocation != null && logLocation.nonEmpty) {
+ val storage =
+ SequentialRecordStorage.getStorage[ReplayLogRecord](Some(new URI(logLocation)))
+ val result = new mutable.ArrayBuffer[EmbeddedControlMessageIdentity]()
+ storage.getReader("CONTROLLER").mkRecordIterator().foreach {
+ case destination: ReplayDestination =>
+ result.append(destination.id)
+ case _ =>
+ }
+ result.map(_.id).toList
+ } else {
+ List()
+ }
+ case None => List()
+ }
+ }
+ }
+
+ /**
+ * This method returns the executions of a workflow given by its ID
+ *
+ * @return executions[]
+ */
+ @GET
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @Path("/{wid}")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def retrieveExecutionsOfWorkflow(
+ @PathParam("wid") wid: Integer,
+ @Auth sessionUser: SessionUser,
+ @QueryParam("status") status: String
+ ): List[WorkflowExecutionEntry] = {
+ val user = sessionUser.getUser
+ if (!WorkflowAccessResource.hasReadAccess(wid, user.getUid)) {
+ List()
+ } else {
+ val statusCodes: Set[Byte] =
+ Option(status)
+ .map(_.trim)
+ .filter(_.nonEmpty)
+ .map { raw =>
+ val tokens = raw.split(',').map(_.trim.toLowerCase).filter(_.nonEmpty)
+ try {
+ tokens.map(stringToAggregatedState).map(maptoStatusCode).toSet
+ } catch {
+ case e: IllegalArgumentException =>
+ throw new BadRequestException(e.getMessage)
+ }
+ }
+ .getOrElse(Set.empty[Byte])
+ getWorkflowExecutions(wid, context, statusCodes)
+ }
+ }
+
+ @GET
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @Path("/{wid}/stats/{eid}")
+ def retrieveWorkflowRuntimeStatistics(
+ @PathParam("wid") wid: Integer,
+ @PathParam("eid") eid: Integer
+ ): List[WorkflowRuntimeStatistics] = {
+ // Create URI for runtime statistics
+ val uriString: String = context
+ .select(WORKFLOW_EXECUTIONS.RUNTIME_STATS_URI)
+ .from(WORKFLOW_EXECUTIONS)
+ .where(
+ WORKFLOW_EXECUTIONS.EID
+ .eq(eid)
+ .and(
+ WORKFLOW_EXECUTIONS.VID.in(
+ context
+ .select(WORKFLOW_VERSION.VID)
+ .from(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.WID.eq(wid))
+ )
+ )
+ )
+ .fetchOneInto(classOf[String])
+
+ if (uriString == null || uriString.isEmpty) {
+ throw new NoSuchElementException(
+ "No runtime statistics URI found for the given execution ID."
+ )
+ }
+
+ val uri: URI = new URI(uriString)
+ val document = DocumentFactory.openDocument(uri)._1
+
+ // Read all records from Iceberg and convert to WorkflowRuntimeStatistics
+ document
+ .get()
+ .map(tuple => {
+ val record = tuple.asInstanceOf[Tuple]
+ WorkflowRuntimeStatistics(
+ operatorId = record.getField(0).asInstanceOf[String],
+ timestamp = record.getField(1).asInstanceOf[Timestamp],
+ inputTupleCount = record.getField(2).asInstanceOf[Long],
+ inputTupleSize = record.getField(3).asInstanceOf[Long],
+ outputTupleCount = record.getField(4).asInstanceOf[Long],
+ outputTupleSize = record.getField(5).asInstanceOf[Long],
+ dataProcessingTime = record.getField(6).asInstanceOf[Long],
+ controlProcessingTime = record.getField(7).asInstanceOf[Long],
+ idleTime = record.getField(8).asInstanceOf[Long],
+ numWorkers = record.getField(9).asInstanceOf[Int],
+ status = record.getField(10).asInstanceOf[Int]
+ )
+ })
+ .toList
+ }
+
+ /** Sets a group of executions' bookmarks to the payload passed in the body. */
+ @PUT
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @Path("/set_execution_bookmarks")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def setExecutionAreBookmarked(
+ request: ExecutionGroupBookmarkRequest,
+ @Auth sessionUser: SessionUser
+ ): Unit = {
+ validateUserCanAccessWorkflow(sessionUser.getUser.getUid, request.wid)
+ val eIdsList = request.eIds.toSeq.asJava
+ if (request.isBookmarked) {
+ // If currently bookmarked, un-bookmark (set bookmarked = false)
+ context
+ .update(WORKFLOW_EXECUTIONS)
+ .set(WORKFLOW_EXECUTIONS.BOOKMARKED, java.lang.Boolean.valueOf(false))
+ .where(WORKFLOW_EXECUTIONS.EID.in(eIdsList))
+ .execute()
+ } else {
+ // If currently not bookmarked, bookmark (set bookmarked = true)
+ context
+ .update(WORKFLOW_EXECUTIONS)
+ .set(WORKFLOW_EXECUTIONS.BOOKMARKED, java.lang.Boolean.valueOf(true))
+ .where(WORKFLOW_EXECUTIONS.EID.in(eIdsList))
+ .execute()
+ }
+
+ }
+
+ /** Determine if the user is authorized to access the workflow, if not raise 401 */
+ private def validateUserCanAccessWorkflow(uid: Integer, wid: Integer): Unit = {
+ if (!WorkflowAccessResource.hasReadAccess(wid, uid))
+ throw new WebApplicationException(Response.Status.UNAUTHORIZED)
+ }
+
+ /** Delete a group of executions */
+ @PUT
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @Path("/delete_executions")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def groupDeleteExecutionsOfWorkflow(
+ request: ExecutionGroupDeleteRequest,
+ @Auth sessionUser: SessionUser
+ ): Unit = {
+ validateUserCanAccessWorkflow(sessionUser.getUser.getUid, request.wid)
+ removeAllExecutionFiles(request.eIds)
+ }
+
+ /** Name a single execution * */
+ @POST
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @Path("/update_execution_name")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def updateWorkflowExecutionsName(
+ request: ExecutionRenameRequest,
+ @Auth sessionUser: SessionUser
+ ): Unit = {
+ validateUserCanAccessWorkflow(sessionUser.getUser.getUid, request.wid)
+ val execution = getExecutionById(request.eId)
+ execution.setName(request.executionName)
+ executionsDao.update(execution)
+ }
+
+ /**
+ * Returns which operators are restricted from export due to dataset access controls.
+ * This endpoint allows the frontend to check restrictions before attempting export.
+ *
+ * @param wid The workflow ID to check
+ * @param user The authenticated user
+ * @return JSON map of operator ID -> array of {ownerEmail, datasetName} that block its export
+ */
+ @GET
+ @Path("/{wid}/result/downloadability")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def getWorkflowResultDownloadability(
+ @PathParam("wid") wid: Integer,
+ @Auth user: SessionUser
+ ): Response = {
+ validateUserCanAccessWorkflow(user.getUser.getUid, wid)
+
+ val datasetRestrictions = getNonDownloadableOperatorMap(wid, user.user)
+
+ // Convert to frontend-friendly format: Map[operatorId -> Array[datasetLabel]]
+ val restrictionMap = datasetRestrictions.map {
+ case (operatorId, datasets) =>
+ operatorId -> datasets.map {
+ case (ownerEmail, datasetName) => s"$datasetName ($ownerEmail)"
+ }.toArray
+ }.asJava
+
+ Response.ok(restrictionMap).build()
+ }
+
+ @POST
+ @Path("/result/export/dataset")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def exportResultToDataset(request: ResultExportRequest, @Auth user: SessionUser): Response = {
+ try {
+ val resultExportService =
+ new ResultExportService(WorkflowIdentity(request.workflowId), request.computingUnitId)
+ resultExportService.exportToDataset(user.user, request)
+
+ } catch {
+ case ex: Exception =>
+ Response
+ .status(Response.Status.INTERNAL_SERVER_ERROR)
+ .`type`(MediaType.APPLICATION_JSON)
+ .entity(Map("error" -> ex.getMessage).asJava)
+ .build()
+ }
+ }
+
+ @POST
+ @Path("/result/export/local")
+ @Consumes(Array(MediaType.APPLICATION_FORM_URLENCODED))
+ def exportResultToLocal(
+ @FormParam("request") requestJson: String,
+ @FormParam("token") token: String
+ ): Response = {
+
+ try {
+ val userOpt = JwtParser.parseToken(token)
+ if (userOpt.isPresent) {
+ val user = userOpt.get()
+ val role = user.getUser.getRole
+ val RolesAllowed = Set(UserRoleEnum.REGULAR, UserRoleEnum.ADMIN)
+ if (!RolesAllowed.contains(role)) {
+ throw new RuntimeException("User role is not allowed to perform this download")
+ }
+ } else {
+ throw new RuntimeException("Invalid or expired token")
+ }
+
+ val request = Json.parse(requestJson).as[ResultExportRequest]
+ val resultExportService =
+ new ResultExportService(WorkflowIdentity(request.workflowId), request.computingUnitId)
+ resultExportService.exportToLocal(request)
+
+ } catch {
+ case ex: Exception =>
+ Response
+ .status(Response.Status.INTERNAL_SERVER_ERROR)
+ .`type`(MediaType.APPLICATION_JSON)
+ .entity(Map("error" -> ex.getMessage).asJava)
+ .build()
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowResource.scala
new file mode 100644
index 00000000000..cb910d11c3c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowResource.scala
@@ -0,0 +1,801 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.workflow
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import com.typesafe.scalalogging.LazyLogging
+import io.dropwizard.auth.Auth
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.virtualidentity.ExecutionIdentity
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.{
+ WorkflowDao,
+ WorkflowOfProjectDao,
+ WorkflowOfUserDao,
+ WorkflowUserAccessDao
+}
+import org.apache.texera.dao.jooq.generated.tables.pojos._
+import org.apache.texera.service.util.LargeBinaryManager
+import org.apache.texera.web.resource.dashboard.hub.EntityType
+import org.apache.texera.web.resource.dashboard.hub.HubResource.recordCloneAction
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowAccessResource.hasReadAccess
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource._
+import org.jooq.impl.DSL.{groupConcatDistinct, noCondition}
+import org.jooq.{Condition, DSLContext, Record9, Result, SelectOnConditionStep}
+
+import java.sql.Timestamp
+import java.util
+import java.util.UUID
+import javax.annotation.security.RolesAllowed
+import javax.servlet.http.HttpServletRequest
+import javax.ws.rs._
+import javax.ws.rs.core.{Context, MediaType}
+import scala.collection.mutable.ListBuffer
+import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
+
+/**
+ * This file handles various request related to saved-workflows.
+ * It sends mysql queries to the MysqlDB regarding the UserWorkflow Table
+ * The details of UserWorkflowTable can be found in /sql/texera_ddl.sql
+ */
+
+object WorkflowResource {
+ private def context: DSLContext =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def workflowDao = new WorkflowDao(context.configuration)
+ private def workflowOfUserDao =
+ new WorkflowOfUserDao(
+ context.configuration
+ )
+ private def workflowUserAccessDao =
+ new WorkflowUserAccessDao(
+ context.configuration()
+ )
+ private def workflowOfProjectDao = new WorkflowOfProjectDao(context.configuration)
+
+ def getWorkflowName(wid: Integer): String = {
+ val workflow = workflowDao.fetchOneByWid(wid)
+ if (workflow == null) {
+ throw new NotFoundException(s"Workflow with id $wid not found")
+ }
+ workflow.getName
+ }
+
+ private def insertWorkflow(workflow: Workflow, user: User): Unit = {
+ workflowDao.insert(workflow)
+ workflowOfUserDao.insert(new WorkflowOfUser(user.getUid, workflow.getWid))
+ workflowUserAccessDao.insert(
+ new WorkflowUserAccess(
+ user.getUid,
+ workflow.getWid,
+ PrivilegeEnum.WRITE
+ )
+ )
+ }
+
+ private def workflowOfUserExists(wid: Integer, uid: Integer): Boolean = {
+ workflowOfUserDao.existsById(
+ context
+ .newRecord(WORKFLOW_OF_USER.UID, WORKFLOW_OF_USER.WID)
+ .values(uid, wid)
+ )
+ }
+
+ private def workflowOfProjectExists(wid: Integer, pid: Integer): Boolean = {
+ workflowOfProjectDao.existsById(
+ context
+ .newRecord(WORKFLOW_OF_PROJECT.WID, WORKFLOW_OF_PROJECT.PID)
+ .values(wid, pid)
+ )
+ }
+
+ case class DashboardWorkflow(
+ isOwner: Boolean,
+ accessLevel: String,
+ ownerName: String,
+ workflow: Workflow,
+ projectIDs: List[Integer],
+ ownerId: Integer
+ )
+
+ case class WorkflowWithPrivilege(
+ name: String,
+ description: String,
+ wid: Integer,
+ content: String,
+ creationTime: Timestamp,
+ lastModifiedTime: Timestamp,
+ isPublished: Boolean,
+ readonly: Boolean
+ )
+
+ case class WorkflowIDs(wids: List[Integer], pid: Option[Integer])
+
+ private def updateWorkflowField(
+ workflow: Workflow,
+ sessionUser: SessionUser,
+ updateFunction: Workflow => Unit
+ ): Unit = {
+ val wid = workflow.getWid
+ val user = sessionUser.getUser
+
+ if (
+ workflowOfUserExists(wid, user.getUid) || WorkflowAccessResource.hasWriteAccess(
+ wid,
+ user.getUid
+ )
+ ) {
+ val userWorkflow = workflowDao.fetchOneByWid(wid)
+ updateFunction(userWorkflow)
+ workflowDao.update(userWorkflow)
+ } else {
+ throw new ForbiddenException("No sufficient access privilege.")
+ }
+ }
+
+ /**
+ * Updates operator IDs in the given workflow content by assigning new unique IDs.
+ * Each operator ID in the "operators" section is replaced with a new ID of the form:
+ * "-operator-"
+ *
+ * @param workflowContent JSON string representing the workflow, containing operator details.
+ * @return The updated workflow content with new operator IDs.
+ */
+ def assignNewOperatorIds(workflowContent: String): String = {
+ val objectMapper = new ObjectMapper().registerModule(DefaultScalaModule)
+ val operatorIdMap = objectMapper
+ .readValue(workflowContent, classOf[Map[String, List[Map[String, String]]]])("operators")
+ .map(operator => {
+ val oldOperatorId = operator("operatorID")
+ val operatorType = operator("operatorType")
+ // operator id in frontend: operatorSchema.operatorType + "-operator-" + uuid(); // v4 = UUID.randomUUID().toString
+ val newOperatorId = s"$operatorType-operator-${UUID.randomUUID()}"
+ oldOperatorId -> newOperatorId
+ })
+ .toMap
+
+ // replace all old operator ids with new operator ids
+ operatorIdMap.foldLeft(workflowContent) {
+ case (updatedContent, (oldId, newId)) =>
+ updatedContent.replace(oldId, newId)
+ }
+ }
+
+ def baseWorkflowSelect(): SelectOnConditionStep[Record9[
+ Integer,
+ String,
+ String,
+ Timestamp,
+ Timestamp,
+ PrivilegeEnum,
+ Integer,
+ String,
+ String
+ ]] = {
+ context
+ .select(
+ WORKFLOW.WID,
+ WORKFLOW.NAME,
+ WORKFLOW.DESCRIPTION,
+ WORKFLOW.CREATION_TIME,
+ WORKFLOW.LAST_MODIFIED_TIME,
+ WORKFLOW_USER_ACCESS.PRIVILEGE,
+ WORKFLOW_OF_USER.UID,
+ USER.NAME,
+ groupConcatDistinct(WORKFLOW_OF_PROJECT.PID).as("projects")
+ )
+ .from(WORKFLOW)
+ .leftJoin(WORKFLOW_USER_ACCESS)
+ .on(WORKFLOW_USER_ACCESS.WID.eq(WORKFLOW.WID))
+ .leftJoin(WORKFLOW_OF_USER)
+ .on(WORKFLOW_OF_USER.WID.eq(WORKFLOW.WID))
+ .leftJoin(USER)
+ .on(USER.UID.eq(WORKFLOW_OF_USER.UID))
+ .leftJoin(WORKFLOW_OF_PROJECT)
+ .on(WORKFLOW.WID.eq(WORKFLOW_OF_PROJECT.WID))
+ }
+
+ def mapWorkflowEntries(
+ workflowEntries: Result[Record9[
+ Integer,
+ String,
+ String,
+ Timestamp,
+ Timestamp,
+ PrivilegeEnum,
+ Integer,
+ String,
+ String
+ ]],
+ uid: Integer
+ ): List[DashboardWorkflow] = {
+ workflowEntries
+ .map(workflowRecord =>
+ DashboardWorkflow(
+ if (uid != null)
+ workflowRecord.into(WORKFLOW_OF_USER).getUid.eq(uid)
+ else false,
+ workflowRecord
+ .into(WORKFLOW_USER_ACCESS)
+ .into(classOf[WorkflowUserAccess])
+ .getPrivilege
+ .toString,
+ workflowRecord.into(USER).getName,
+ workflowRecord.into(WORKFLOW).into(classOf[Workflow]),
+ if (workflowRecord.component9() == null) List[Integer]()
+ else
+ workflowRecord.component9().split(',').map(str => Integer.valueOf(str)).toList,
+ workflowRecord.into(WORKFLOW_OF_USER).getUid
+ )
+ )
+ .asScala
+ .toList
+ }
+}
+
+@Produces(Array(MediaType.APPLICATION_JSON))
+@Path("/workflow")
+class WorkflowResource extends LazyLogging {
+
+ /**
+ * This method returns all workflow IDs that the user has access to
+ *
+ * @return WorkflowID[]
+ */
+ @GET
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/user-workflow-ids")
+ def retrieveIDs(@Auth user: SessionUser): util.List[String] = {
+ context
+ .select(WORKFLOW_USER_ACCESS.WID)
+ .from(WORKFLOW_USER_ACCESS)
+ .where(WORKFLOW_USER_ACCESS.UID.eq(user.getUid))
+ .fetchInto(classOf[String])
+ }
+
+ /**
+ * This method returns all owner user names of the workflows that the user has access to
+ *
+ * @return OwnerName[]
+ */
+ @GET
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/user-workflow-owners")
+ def retrieveOwners(@Auth user: SessionUser): util.List[String] = {
+ context
+ .selectDistinct(USER.EMAIL)
+ .from(WORKFLOW_USER_ACCESS)
+ .join(WORKFLOW_OF_USER)
+ .on(WORKFLOW_USER_ACCESS.WID.eq(WORKFLOW_OF_USER.WID))
+ .join(USER)
+ .on(WORKFLOW_OF_USER.UID.eq(USER.UID))
+ .where(WORKFLOW_USER_ACCESS.UID.eq(user.getUid))
+ .fetchInto(classOf[String])
+ }
+
+ /**
+ * This method returns workflow IDs, that contain the selected operators, as strings
+ *
+ * @return WorkflowID[]
+ */
+ @GET
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/search-by-operators")
+ def searchWorkflowByOperator(
+ @QueryParam("operator") operator: String,
+ @Auth sessionUser: SessionUser
+ ): List[String] = {
+ // Example GET url: localhost:8080/workflow/searchOperators?operator=Regex,CSVFileScan
+ val user = sessionUser.getUser
+ val quotes = "\""
+ val operatorArray =
+ operator.replace(" ", "").stripPrefix("[").stripSuffix("]").split(',')
+ var orCondition: Condition = noCondition()
+ for (i <- operatorArray.indices) {
+ val operatorName = operatorArray(i)
+ orCondition = orCondition.or(
+ WORKFLOW.CONTENT
+ .likeIgnoreCase(
+ "%" + quotes + "operatorType" + quotes + ":" + quotes + s"$operatorName" + quotes + "%"
+ //gives error when I try to combine escape character with formatted string
+ //may be due to old scala version bug
+ )
+ )
+
+ }
+
+ val workflowEntries =
+ context
+ .select(
+ WORKFLOW.WID
+ )
+ .from(WORKFLOW)
+ .join(WORKFLOW_USER_ACCESS)
+ .on(WORKFLOW_USER_ACCESS.WID.eq(WORKFLOW.WID))
+ .where(
+ orCondition
+ .and(WORKFLOW_USER_ACCESS.UID.eq(user.getUid))
+ )
+ .fetch()
+
+ workflowEntries
+ .map(workflowRecord => {
+ workflowRecord.into(WORKFLOW).getWid.intValue().toString
+ })
+ .asScala
+ .toList
+ }
+
+ /**
+ * This method returns the current in-session user's workflow list based on all workflows he/she has access to
+ *
+ * @return Workflow[]
+ */
+ @GET
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/list")
+ def retrieveWorkflowsBySessionUser(
+ @Auth sessionUser: SessionUser
+ ): List[DashboardWorkflow] = {
+ val user = sessionUser.getUser
+ val workflowEntries = baseWorkflowSelect()
+ .where(WORKFLOW_USER_ACCESS.UID.eq(user.getUid))
+ .groupBy(
+ WORKFLOW.WID,
+ WORKFLOW.NAME,
+ WORKFLOW.DESCRIPTION,
+ WORKFLOW.CREATION_TIME,
+ WORKFLOW.LAST_MODIFIED_TIME,
+ WORKFLOW_USER_ACCESS.PRIVILEGE,
+ WORKFLOW_OF_USER.UID,
+ USER.NAME
+ )
+ .fetch()
+ mapWorkflowEntries(workflowEntries, user.getUid)
+ }
+
+ /**
+ * This method handles the client request to get a specific workflow to be displayed in canvas
+ * at current design, it only takes the workflowID and searches within the database for the matching workflow
+ * for future design, it should also take userID as an parameter.
+ *
+ * @param wid workflow id, which serves as the primary key in the UserWorkflow database
+ * @return a json string representing an savedWorkflow
+ */
+ @GET
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/{wid}")
+ def retrieveWorkflow(
+ @PathParam("wid") wid: Integer,
+ @Auth user: SessionUser
+ ): WorkflowWithPrivilege = {
+ if (WorkflowAccessResource.hasReadAccess(wid, user.getUid)) {
+ val workflow = workflowDao.fetchOneByWid(wid)
+ WorkflowWithPrivilege(
+ workflow.getName,
+ workflow.getDescription,
+ workflow.getWid,
+ workflow.getContent,
+ workflow.getCreationTime,
+ workflow.getLastModifiedTime,
+ workflow.getIsPublic,
+ !WorkflowAccessResource.hasWriteAccess(wid, user.getUid)
+ )
+ } else {
+ throw new ForbiddenException("No sufficient access privilege.")
+ }
+ }
+
+ /**
+ * This method persists the workflow into database
+ *
+ * @param workflow , a workflow
+ * @return Workflow, which contains the generated wid if not provided//
+ * TODO: divide into two endpoints -> one for new-workflow and one for updating existing workflow
+ * TODO: if the persist is triggered in parallel, the none atomic actions currently might cause an issue.
+ * Should consider making the operations atomic
+ */
+ @POST
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/persist")
+ def persistWorkflow(workflow: Workflow, @Auth sessionUser: SessionUser): Workflow = {
+ val user = sessionUser.getUser
+ if (user == org.apache.texera.web.auth.GuestAuthFilter.GUEST) {
+ throw new ForbiddenException("Guest user does not have access to db.")
+ }
+
+ if (workflowOfUserExists(workflow.getWid, user.getUid)) {
+ WorkflowVersionResource.insertVersion(workflow, insertingNewWorkflow = false)
+ workflowDao.update(workflow)
+ } else {
+ if (!WorkflowAccessResource.hasReadAccess(workflow.getWid, user.getUid)) {
+ // Check if this workflow exists in the database
+ val workflowExistsInDb =
+ workflow.getWid != null && workflowDao.existsById(workflow.getWid)
+ if (workflowExistsInDb) {
+ // User trying to persist an existing workflow without access - reject
+ throw new ForbiddenException("No sufficient access privilege.")
+ }
+ // This is a new workflow being created (wid is null or doesn't exist in DB)
+ workflow.setWid(null)
+ insertWorkflow(workflow, user)
+ WorkflowVersionResource.insertVersion(workflow, insertingNewWorkflow = true)
+ } else if (WorkflowAccessResource.hasWriteAccess(workflow.getWid, user.getUid)) {
+ WorkflowVersionResource.insertVersion(workflow, insertingNewWorkflow = false)
+ // not owner but has write access
+ workflowDao.update(workflow)
+ } else {
+ // not owner and no write access -> rejected
+ throw new ForbiddenException("No sufficient access privilege.")
+ }
+ }
+
+ val wid = workflow.getWid
+ workflowDao.fetchOneByWid(wid)
+ }
+
+ /**
+ * This method duplicates the target workflow, the new workflow name is appended with `_copy`
+ *
+ * @param workflow , a workflow to be duplicated
+ * @return Workflow, which contains the generated wid if not provided
+ */
+ @POST
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/duplicate")
+ def duplicateWorkflow(
+ workflowIDs: WorkflowIDs,
+ @Auth sessionUser: SessionUser
+ ): List[DashboardWorkflow] = {
+
+ val user = sessionUser.getUser
+ // do the permission check first
+ for (wid <- workflowIDs.wids) {
+ if (!WorkflowAccessResource.hasReadAccess(wid, user.getUid)) {
+ throw new ForbiddenException("No sufficient access privilege.")
+ }
+ }
+
+ val resultWorkflows: ListBuffer[DashboardWorkflow] = ListBuffer()
+ val addToProject = workflowIDs.pid.nonEmpty
+ // then start a transaction and do the duplication
+ try {
+ context.transaction { txConfig =>
+ for (wid <- workflowIDs.wids) {
+ val oldWorkflow: Workflow = workflowDao.fetchOneByWid(wid)
+ val newWorkflow = createWorkflow(
+ new Workflow(
+ null,
+ oldWorkflow.getName + "_copy",
+ oldWorkflow.getDescription,
+ assignNewOperatorIds(oldWorkflow.getContent),
+ null,
+ null,
+ false
+ ),
+ sessionUser
+ )
+ // if workflows also need to be added to the project
+ if (addToProject) {
+ val newWid = newWorkflow.workflow.getWid
+ if (!hasReadAccess(newWid, user.getUid)) {
+ throw new ForbiddenException("No sufficient access privilege to workflow.")
+ }
+ val pid = workflowIDs.pid.get
+ if (!workflowOfProjectExists(newWid, pid)) {
+ workflowOfProjectDao.insert(new WorkflowOfProject(newWid, pid))
+ } else {
+ throw new BadRequestException("Workflow already exists in the project")
+ }
+ }
+ resultWorkflows += newWorkflow
+ }
+ }
+ } catch {
+ case _: BadRequestException | _: ForbiddenException =>
+ case NonFatal(exception) =>
+ throw new WebApplicationException(exception)
+ }
+ resultWorkflows.toList
+ }
+
+ @POST
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/clone/{wid}")
+ def cloneWorkflow(
+ @PathParam("wid") wid: Integer,
+ @Auth sessionUser: SessionUser,
+ @Context request: HttpServletRequest
+ ): Integer = {
+ val oldWorkflow: Workflow = workflowDao.fetchOneByWid(wid)
+ val newWorkflow: DashboardWorkflow = createWorkflow(
+ new Workflow(
+ null,
+ oldWorkflow.getName + "_clone",
+ oldWorkflow.getDescription,
+ assignNewOperatorIds(oldWorkflow.getContent),
+ null,
+ null,
+ false
+ ),
+ sessionUser
+ )
+
+ recordCloneAction(request, sessionUser.getUid, wid, EntityType.Workflow)
+
+ newWorkflow.workflow.getWid
+ }
+
+ /**
+ * This method creates and insert a new workflow to database
+ *
+ * @param workflow , a workflow to be created
+ * @return Workflow, which contains the generated wid if not provided
+ */
+ @POST
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/create")
+ def createWorkflow(workflow: Workflow, @Auth sessionUser: SessionUser): DashboardWorkflow = {
+ val user = sessionUser.getUser
+ if (workflow.getWid != null) {
+ throw new BadRequestException("Cannot create a new workflow with a provided id.")
+ } else {
+ insertWorkflow(workflow, user)
+ WorkflowVersionResource.insertVersion(workflow, insertingNewWorkflow = true)
+ DashboardWorkflow(
+ isOwner = true,
+ PrivilegeEnum.WRITE.toString,
+ user.getName,
+ workflowDao.fetchOneByWid(workflow.getWid),
+ List[Integer](),
+ user.getUid
+ )
+ }
+
+ }
+
+ /**
+ * Deletes workflows from the database and cleans up associated resources.
+ *
+ * @param workflowIDs The IDs of workflows to delete
+ * @param sessionUser Current authenticated user
+ * @return Unit, with appropriate HTTP status: 200 if deleted, 400 if not exists
+ */
+ @POST
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/delete")
+ def deleteWorkflow(workflowIDs: WorkflowIDs, @Auth sessionUser: SessionUser): Unit = {
+ val user = sessionUser.getUser
+
+ try {
+ // Find all execution IDs related to these workflows
+ val eids = context
+ .select(WORKFLOW_EXECUTIONS.EID)
+ .from(WORKFLOW_EXECUTIONS)
+ .join(WORKFLOW_VERSION)
+ .on(WORKFLOW_EXECUTIONS.VID.eq(WORKFLOW_VERSION.VID))
+ .join(WORKFLOW)
+ .on(WORKFLOW_VERSION.WID.eq(WORKFLOW.WID))
+ .where(WORKFLOW.WID.in(workflowIDs.wids.asJava))
+ .fetchInto(classOf[Integer])
+ .asScala
+ .toList
+
+ LargeBinaryManager.deleteAllObjects()
+
+ // Collect all URIs related to executions for cleanup
+ val uris = eids.flatMap { eid =>
+ val executionId = ExecutionIdentity(eid.longValue())
+
+ // Gather URIs from all execution resources
+ val resultUris = WorkflowExecutionsResource.getResultUrisByExecutionId(executionId)
+ val consoleMessagesUris =
+ WorkflowExecutionsResource.getConsoleMessagesUriByExecutionId(executionId)
+ val runtimeStatsUris =
+ WorkflowExecutionsResource.getRuntimeStatsUriByExecutionId(executionId).toList
+
+ resultUris ++ consoleMessagesUris ++ runtimeStatsUris
+ }
+
+ // Delete workflows in a transaction
+ context.transaction { _ =>
+ for (wid <- workflowIDs.wids) {
+ if (workflowOfUserExists(wid, user.getUid)) {
+ workflowDao.deleteById(wid)
+ } else {
+ throw new BadRequestException("The workflow does not exist.")
+ }
+ }
+ }
+
+ // Clean up document storage
+ try {
+ uris.foreach { uri =>
+ try {
+ val (document, _) = DocumentFactory.openDocument(uri)
+ document.clear()
+ } catch {
+ case e: IllegalArgumentException if e.getMessage.contains("No storage is found") =>
+ logger.warn(s"Storage for URI $uri not found, ignoring: ${e.getMessage}")
+ case NonFatal(e) =>
+ logger.error(s"Failed to clear document for URI $uri", e)
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.error("Failed to clean up execution results", e)
+ }
+ } catch {
+ case _: BadRequestException =>
+ case NonFatal(exception) => throw new WebApplicationException(exception)
+ }
+ }
+
+ @POST
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/update/name")
+ def updateWorkflowName(
+ workflow: Workflow,
+ @Auth sessionUser: SessionUser
+ ): Unit = {
+ updateWorkflowField(workflow, sessionUser, _.setName(workflow.getName))
+ }
+
+ @POST
+ @Consumes(Array(MediaType.APPLICATION_JSON))
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/update/description")
+ def updateWorkflowDescription(
+ workflow: Workflow,
+ @Auth sessionUser: SessionUser
+ ): Unit = {
+ updateWorkflowField(workflow, sessionUser, _.setDescription(workflow.getDescription))
+ }
+
+ @PUT
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/public/{wid}")
+ def makePublic(@PathParam("wid") wid: Integer, @Auth user: SessionUser): Unit = {
+ if (!WorkflowAccessResource.hasWriteAccess(wid, user.getUid)) {
+ throw new ForbiddenException(s"You do not have permission to modify workflow $wid")
+ }
+ val workflow: Workflow = workflowDao.fetchOneByWid(wid)
+ workflow.setIsPublic(true)
+ workflowDao.update(workflow)
+ }
+
+ @PUT
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/private/{wid}")
+ def makePrivate(@PathParam("wid") wid: Integer, @Auth user: SessionUser): Unit = {
+ if (!WorkflowAccessResource.hasWriteAccess(wid, user.getUid)) {
+ throw new ForbiddenException(s"You do not have permission to modify workflow $wid")
+ }
+ val workflow: Workflow = workflowDao.fetchOneByWid(wid)
+ workflow.setIsPublic(false)
+ workflowDao.update(workflow)
+ }
+
+ @GET
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ @Path("/type/{wid}")
+ def getWorkflowType(@PathParam("wid") wid: Integer): String = {
+ val workflow: Workflow = workflowDao.fetchOneByWid(wid)
+ if (workflow.getIsPublic) {
+ "Public"
+ } else {
+ "Private"
+ }
+ }
+
+ @GET
+ @Produces(Array(MediaType.TEXT_PLAIN))
+ @Path("/owner_name")
+ def getOwnerName(@QueryParam("wid") wid: Integer): String = {
+ context
+ .select(USER.NAME)
+ .from(USER)
+ .join(WORKFLOW_OF_USER)
+ .on(USER.UID.eq(WORKFLOW_OF_USER.UID))
+ .where(WORKFLOW_OF_USER.WID.eq(wid))
+ .fetchOneInto(classOf[String])
+ }
+
+ @GET
+ @Path("/workflow_name")
+ def getWorkflowName(@QueryParam("wid") wid: Integer): String = {
+ context
+ .select(
+ WORKFLOW.NAME
+ )
+ .from(WORKFLOW)
+ .where(WORKFLOW.WID.eq(wid))
+ .fetchOneInto(classOf[String])
+ }
+
+ @GET
+ @Path("/publicised/{wid}")
+ def retrievePublicWorkflow(
+ @PathParam("wid") wid: Integer
+ ): WorkflowWithPrivilege = {
+ val workflow = workflowDao.ctx
+ .selectFrom(WORKFLOW)
+ .where(WORKFLOW.WID.eq(wid))
+ .and(WORKFLOW.IS_PUBLIC.isTrue)
+ .fetchOne()
+ WorkflowWithPrivilege(
+ workflow.getName,
+ workflow.getDescription,
+ workflow.getWid,
+ workflow.getContent,
+ workflow.getCreationTime,
+ workflow.getLastModifiedTime,
+ workflow.getIsPublic,
+ readonly = true
+ )
+ }
+
+ @GET
+ @Path("/workflow_description")
+ def getWorkflowDescription(@QueryParam("wid") wid: Integer): String = {
+ context
+ .select(
+ WORKFLOW.DESCRIPTION
+ )
+ .from(WORKFLOW)
+ .where(WORKFLOW.WID.eq(wid))
+ .fetchOneInto(classOf[String])
+ }
+
+ //TODO Get size from database
+ @GET
+ @Path("/size")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getSize(@QueryParam("wid") wids: java.util.List[Integer]): java.util.Map[Integer, Int] = {
+ val result = new java.util.HashMap[Integer, Int]()
+ if (wids != null && !wids.isEmpty) {
+ workflowDao.ctx
+ .selectFrom(WORKFLOW)
+ .where(WORKFLOW.WID.in(wids))
+ .fetch()
+ .asScala
+ .foreach { wf =>
+ result.put(wf.getWid, wf.getContent.length)
+ }
+ }
+ result
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowVersionResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowVersionResource.scala
new file mode 100644
index 00000000000..7be74ae5b00
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowVersionResource.scala
@@ -0,0 +1,450 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.workflow
+
+import com.flipkart.zjsonpatch.{JsonDiff, JsonPatch}
+import io.dropwizard.auth.Auth
+import org.apache.texera.amber.util.JSONUtils.objectMapper
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.config.UserSystemConfig
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.Tables.WORKFLOW_VERSION
+import org.apache.texera.dao.jooq.generated.tables.daos.{WorkflowDao, WorkflowVersionDao}
+import org.apache.texera.dao.jooq.generated.tables.pojos.{Workflow, WorkflowVersion}
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource.{
+ DashboardWorkflow,
+ assignNewOperatorIds
+}
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowVersionResource._
+import org.jooq.DSLContext
+
+import java.sql.Timestamp
+import javax.annotation.security.RolesAllowed
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+import scala.jdk.CollectionConverters.IterableHasAsScala
+
+/**
+ * This file handles various request related to workflows versions.
+ * The details of the mysql tables can be found in /sql/texera_ddl.sql
+ */
+
+object WorkflowVersionResource {
+ private def context: DSLContext =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def workflowVersionDao = new WorkflowVersionDao(context.configuration)
+ private def workflowDao = new WorkflowDao(context.configuration)
+ // constant to indicate versions should be aggregated if they are within the specified time limit
+ private final val AGGREGATE_TIME_LIMIT_MILLSEC =
+ UserSystemConfig.workflowVersionCollapseIntervalInMinutes * 60000
+ // list of Json keys in the diff patch that are considered UNimportant
+ private final val VERSION_UNIMPORTANCE_RULES = List("/operatorPositions/")
+ private final val SNAPSHOT_UNIMPORTANCE_RULES = List("replace")
+
+ /**
+ * This function does the check of the difference between the current workflow and its previous version if it exists and inserts a new version
+ *
+ * @param workflow
+ * @param insertingNewWorkflow indicates if the workflow didn't exist before
+ */
+ def insertVersion(workflow: Workflow, insertingNewWorkflow: Boolean): Unit = {
+ val wid = workflow.getWid
+ // retrieve current workflow from DB
+ val currentWorkflow = workflowDao.fetchOneByWid(wid)
+ // if the workflow is new then previous workflow is empty
+ val existingWorkflowContent = if (insertingNewWorkflow) "{}" else currentWorkflow.getContent
+ // compute diff
+ val patch = JsonDiff.asJson(
+ objectMapper.readTree(workflow.getContent),
+ objectMapper.readTree(existingWorkflowContent)
+ )
+
+ // if we are creating a new workflow, even if it is empty, create a new version
+ // otherwise, only when there is a diff we would create a new version
+ if (insertingNewWorkflow || !patch.isEmpty) {
+ insertNewVersion(wid, patch.toString)
+ }
+ }
+
+ /**
+ * This function updates the content of the latest version and inserts a new empty version for the current workflow
+ *
+ * @param patch to update latest version
+ * @param wid
+ */
+ private def updateLatestVersion(patch: String, wid: Integer): Unit = {
+ // get the latest version to update its content
+ val vid = getLatestVersion(wid)
+ val workflowVersion = workflowVersionDao.fetchOneByVid(vid)
+ workflowVersion.setContent(patch)
+ workflowVersionDao.update(workflowVersion)
+ }
+
+ /**
+ * This function retrieves the latest version of a workflow
+ *
+ * @param wid
+ * @return vid
+ */
+ def getLatestVersion(wid: Integer): Integer = {
+ val versions = context
+ .select(WORKFLOW_VERSION.VID)
+ .from(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.WID.eq(wid))
+ .fetchInto(classOf[Integer])
+ .asScala
+ .toList
+ // for backwards compatibility check, old constructed versions would follow the old design by not saving the current
+ // version as an empty delta, so should do the check and create one once
+ // TODO should remove the check when all versions in the DB follow latest design
+ if (versions.isEmpty) {
+ return insertNewVersion(wid).getVid
+ }
+ versions.max
+ }
+
+ /**
+ * This function inserts a new version for a workflow
+ *
+ * @param wid
+ */
+ def insertNewVersion(wid: Integer, content: String = "[]"): WorkflowVersion = {
+ val workflowVersion = new WorkflowVersion()
+ workflowVersion.setContent(content)
+ workflowVersion.setWid(wid)
+ workflowVersionDao.insert(workflowVersion)
+ workflowVersion
+ }
+
+ /**
+ * This function retrieves the content of versions from a specific workflow in a range
+ *
+ * @param lowerBound lower bound of the version search range
+ * @param UpperBound upper bound of the search range
+ * @param wid workflow id
+ * @return a list of contents as strings
+ */
+ def isSnapshotInRangeUnimportant(
+ lowerBound: Integer,
+ UpperBound: Integer,
+ wid: Integer
+ ): Boolean = {
+ if (lowerBound == UpperBound) {
+ return true
+ }
+ val contents = context
+ .select(WORKFLOW_VERSION.CONTENT)
+ .from(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.WID.eq(wid))
+ .and(WORKFLOW_VERSION.VID.between(lowerBound).and(UpperBound))
+ .fetchInto(classOf[String])
+ .asScala
+ .toList
+ contents.forall(content => !isSnapshotImportant(content))
+ }
+
+ /**
+ * This function parses the content of the delta to determine if it is positional only
+ *
+ * @param versionContent
+ * @return
+ */
+ private def isSnapshotImportant(versionContent: String): Boolean = {
+ val jsonTreeIterator = objectMapper.readTree(versionContent).iterator()
+ while (jsonTreeIterator.hasNext) {
+ // if the change(which is marked by the key `path` using the Json patch library
+ // doesn't contain any of the specified keywords then it shall be deemed important
+ if (
+ !SNAPSHOT_UNIMPORTANCE_RULES.exists(jsonTreeIterator.next().path("op").asText().contains)
+ ) {
+ return true
+ }
+ }
+ false
+ }
+
+ /**
+ * This function gives a label to each version whether it is significant or not based on a few rules
+ * The reason why it is computed AFTER retrieving the list of versions is due to multiple reasons:
+ * * 1. minimize the changes to the database.
+ * * 2. since the frontend sends persisting workflow request often, we don't want to slow down the
+ * insertion to DB because of computing the version's importance especially because the request is
+ * async, the versions can quickly become inconsistent if there is delay.
+ * * 3. The rules can be changed in the future so we want this logic to be changed flexibly.
+ *
+ * @param versions the version from DB sorted from latest to earliest
+ * @return
+ */
+ private def encodeVersionImportance(
+ currentVersions: List[WorkflowVersion]
+ ): List[VersionEntry] = {
+ var impEncodedVersions: List[VersionEntry] = List()
+
+ val lastVersion = currentVersions.head
+ var lastVersionTime = lastVersion.getCreationTime
+ impEncodedVersions = impEncodedVersions :+ VersionEntry(
+ lastVersion.getVid,
+ lastVersion.getCreationTime,
+ lastVersion.getContent,
+ true
+ ) // the first (latest)
+ // version is important even if it is positional
+ var versionImportance: Boolean = true
+ for (version <- currentVersions.tail) {
+ if (
+ isWithinTimeLimit(
+ lastVersionTime,
+ version.getCreationTime
+ )
+ ) {
+ versionImportance = false
+ } // try reducing unnecessary check of positional versions
+ // because parsing the Json string is expensive
+ else {
+ lastVersionTime = version.getCreationTime
+ versionImportance = isVersionImportant(version.getContent)
+ }
+ impEncodedVersions = impEncodedVersions :+ VersionEntry(
+ version.getVid,
+ version.getCreationTime,
+ version.getContent,
+ versionImportance
+ )
+ }
+ impEncodedVersions
+ }
+
+ /**
+ * This function determines whether this version is still within the time range of previous versions
+ *
+ * @param latestTime
+ * @param currentVersionTimestamp
+ * @return
+ */
+ private def isWithinTimeLimit(
+ latestTime: Timestamp,
+ currentVersionTimestamp: Timestamp
+ ): Boolean = {
+ (latestTime.getTime - currentVersionTimestamp.getTime) < AGGREGATE_TIME_LIMIT_MILLSEC
+ }
+
+ /**
+ * This function parses the content of the delta to determine if it is positional only
+ *
+ * @param versionContent
+ * @return
+ */
+ private def isVersionImportant(versionContent: String): Boolean = {
+ val jsonTreeIterator = objectMapper.readTree(versionContent).iterator()
+ while (jsonTreeIterator.hasNext) {
+ // if the change(which is marked by the key `path` using the Json patch library
+ // doesn't contain any of the specified keywords then it shall be deemed important
+ if (
+ !VERSION_UNIMPORTANCE_RULES.exists(jsonTreeIterator.next().path("path").asText().contains)
+ ) {
+ return true
+ }
+ }
+ false
+ }
+
+ /**
+ * Fetches all versions of a workflow from a specific version ID to the latest version
+ *
+ * @param wid workflow ID to query
+ * @param vid starting version ID (inclusive)
+ * @return List of workflow versions ordered from latest to earliest
+ */
+ def fetchSubsequentVersions(
+ wid: Integer,
+ vid: Integer,
+ context: DSLContext
+ ): List[WorkflowVersion] = {
+ context
+ .select(WORKFLOW_VERSION.VID, WORKFLOW_VERSION.CREATION_TIME, WORKFLOW_VERSION.CONTENT)
+ .from(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.WID.eq(wid).and(WORKFLOW_VERSION.VID.ge(vid)))
+ .orderBy(WORKFLOW_VERSION.VID.desc())
+ .fetchInto(classOf[WorkflowVersion])
+ .asScala
+ .toList
+ }
+
+ /**
+ * This function applies all the diff versions to a workflow
+ *
+ * @param versions list of computed delta in each version
+ * @param workflow beginning workflow ( more recent)
+ * @return the (old) workflow is computed after applying all the patches
+ */
+ def applyPatch(versions: List[WorkflowVersion], workflow: Workflow): Workflow = {
+ // loop all versions and apply the patch
+ for (patch <- versions) {
+ workflow.setContent(
+ JsonPatch
+ .apply(
+ objectMapper.readTree(patch.getContent),
+ objectMapper.readTree(workflow.getContent)
+ )
+ .toString
+ )
+ workflow.setCreationTime(patch.getCreationTime)
+ workflow.setLastModifiedTime(patch.getCreationTime)
+ }
+ // the checked out version is returned
+ workflow
+ }
+
+ /**
+ * This class is to add version importance encoding to the existing `VersionEntry` from DB
+ *
+ * @param vId
+ * @param creationTime
+ * @param content
+ * @param importance false is not an important version and true is an important version
+ */
+ case class VersionEntry(
+ vId: Integer,
+ creationTime: Timestamp,
+ content: String,
+ importance: Boolean
+ )
+
+}
+
+@Path("/version")
+@Produces(Array(MediaType.APPLICATION_JSON))
+class WorkflowVersionResource {
+
+ /**
+ * This method returns the versions of a workflow given by its ID
+ *
+ * @return versions[]
+ */
+ @GET
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @Path("/{wid}")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def retrieveVersionsOfWorkflow(
+ @PathParam("wid") wid: Integer,
+ @Auth sessionUser: SessionUser
+ ): List[VersionEntry] = {
+ val user = sessionUser.getUser
+ if (!WorkflowAccessResource.hasReadAccess(wid, user.getUid)) {
+ List()
+ } else {
+ encodeVersionImportance(
+ context
+ .select(WORKFLOW_VERSION.VID, WORKFLOW_VERSION.CREATION_TIME, WORKFLOW_VERSION.CONTENT)
+ .from(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.WID.eq(wid))
+ .orderBy(WORKFLOW_VERSION.CREATION_TIME.desc())
+ .fetchInto(classOf[WorkflowVersion])
+ .asScala
+ .toList
+ )
+ }
+ }
+
+ /**
+ * This method returns a particular version of a workflow given the vid and wid
+ * first, list the versions of the workflow; second, from the current version(last) apply the differences until the requested version
+ * third, return the requested workflow
+ *
+ * @param wid workflowID of the current workflow the user is working on
+ * @param vid versionID of the checked-out version to be computed and returned
+ * @return workflow of a particular version
+ */
+ @GET
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @Path("/{wid}/{vid}")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def retrieveWorkflowVersion(
+ @PathParam("wid") wid: Integer,
+ @PathParam("vid") vid: Integer,
+ @Auth sessionUser: SessionUser
+ ): Workflow = {
+ val user = sessionUser.getUser
+ if (!WorkflowAccessResource.hasReadAccess(wid, user.getUid)) {
+ throw new ForbiddenException("No sufficient access privilege.")
+ } else {
+ // fetch all versions equal to and subsequent to the specified version
+ val versionEntries = fetchSubsequentVersions(wid, vid, context)
+ // apply patch
+ val currentWorkflow = workflowDao.fetchOneByWid(wid)
+ // return particular version of the workflow
+ val res: Workflow = applyPatch(versionEntries, currentWorkflow)
+ res
+ }
+ }
+
+ @POST
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ @Path("/clone/{vid}")
+ @RolesAllowed(Array("REGULAR", "ADMIN"))
+ def cloneVersion(
+ @PathParam("vid") vid: Integer,
+ @Auth sessionUser: SessionUser,
+ requestBody: java.util.Map[String, Int]
+ ): Integer = {
+ val displayedVersionId = requestBody.get("displayedVersionId")
+
+ // Fetch the workflow ID (`wid`) associated with the specified version (`vid`)
+ val versionRecord = Option(
+ context
+ .select(WORKFLOW_VERSION.WID)
+ .from(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.VID.eq(vid))
+ .fetchOne()
+ ).getOrElse {
+ throw new NotFoundException(s"Version ID $vid not found.")
+ }
+ val wid = versionRecord.get(WORKFLOW_VERSION.WID)
+ // Use retrieveWorkflowVersion to get the specified version of the workflow
+ val workflowVersion = retrieveWorkflowVersion(wid, vid, sessionUser)
+ // Generate a new name for the cloned workflow
+ val newWorkflowName = s"${workflowVersion.getName}_v${displayedVersionId}_copy"
+ // Create a new workflow based on the retrieved version
+ val workflowResource = new WorkflowResource()
+ val newWorkflow: DashboardWorkflow =
+ try {
+ workflowResource.createWorkflow(
+ new Workflow(
+ null,
+ newWorkflowName,
+ workflowVersion.getDescription,
+ assignNewOperatorIds(workflowVersion.getContent),
+ null,
+ null,
+ false
+ ),
+ sessionUser
+ )
+ } catch {
+ case e: Exception =>
+ throw new InternalServerErrorException(
+ "An error occurred while creating the cloned workflow."
+ )
+ }
+ newWorkflow.workflow.getWid
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveManager.scala b/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveManager.scala
new file mode 100644
index 00000000000..0613496c39c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveManager.scala
@@ -0,0 +1,315 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.pythonvirtualenvironment
+
+import java.nio.file.{Files, Path, Paths}
+import java.util.concurrent.BlockingQueue
+import scala.collection.mutable.Map
+import scala.jdk.CollectionConverters._
+import scala.sys.process._
+import java.util.Comparator
+import org.apache.texera.amber.config.PythonUtils
+
+/**
+ * PveManager is responsible for managing Python Virtual Environments (PVEs)
+ * for each Computing Unit
+ *
+ * It supports:
+ * - Creating and initializing isolated Python environments
+ * - Streaming pip output logs back to the caller
+ *
+ * Each PVE is stored under:
+ * /tmp/texera-pve/venvs/{cuid}/{pveName}/
+ */
+
+object PveManager {
+
+ case class PvePackageResponse(
+ pveName: String,
+ userPackages: Seq[String]
+ )
+
+ private val VenvRoot: Path = Paths.get("/tmp/texera-pve/venvs")
+
+ private def cuidDir(cuid: Int, pveName: String): Path = {
+ VenvRoot.resolve(cuid.toString).resolve(pveName)
+ }
+
+ private def pveDir(cuid: Int, pveName: String): Path =
+ cuidDir(cuid, pveName).resolve("pve")
+
+ private def pythonBinPath(cuid: Int, pveName: String): Path =
+ pveDir(cuid, pveName).resolve("bin").resolve("python")
+
+ private def pipEnv: Map[String, String] =
+ Map(
+ "PYTHONUNBUFFERED" -> "1",
+ "PIP_PROGRESS_BAR" -> "off",
+ "PIP_DISABLE_PIP_VERSION_CHECK" -> "1",
+ "PIP_NO_INPUT" -> "1"
+ )
+
+ def getSystemPackages(): Seq[String] = {
+ val python = PythonUtils.getPythonExecutable
+ Process(Seq(python, "-m", "pip", "freeze")).!!.split("\n").map(_.trim).filter(_.nonEmpty).toSeq
+ }
+
+ /**
+ * Creates a new PVE for a CU.
+ *
+ * Behavior:
+ * Creates a fresh venv and installs dependencies
+ *
+ * Steps:
+ * 1. Install system dependencies
+ * 2. Logs progress to the provided queue.
+ */
+ def createNewPve(
+ cuid: Int,
+ queue: BlockingQueue[String],
+ pveName: String,
+ isLocal: Boolean
+ ): Unit = {
+ queue.put(s"[PVE] Creating new PVE for cuid: $cuid with name: $pveName")
+
+ // NOTE: These paths are derived from computing-unit-master.dockerfile.
+ // If requirements.txt or operator-requirements.txt locations change, update these paths.
+ val requirementsPath =
+ if (isLocal) Paths.get("amber", "requirements.txt")
+ else Paths.get("/tmp", "requirements.txt")
+
+ val operatorRequirementsPath =
+ if (isLocal) Paths.get("amber", "operator-requirements.txt")
+ else Paths.get("/tmp", "operator-requirements.txt")
+
+ if (!Files.exists(requirementsPath) || !Files.exists(operatorRequirementsPath)) {
+ queue.put(s"[PVE][ERR] System requirements not found")
+ return
+ }
+
+ val venvDirPath = pveDir(cuid, pveName).toAbsolutePath
+ val python = pythonBinPath(cuid, pveName).toAbsolutePath.toString
+ val envVars = pipEnv
+
+ val createVenvPython = PythonUtils.getPythonExecutable
+
+ Files.createDirectories(venvDirPath.getParent)
+
+ val createCode = Process(Seq(createVenvPython, "-m", "venv", venvDirPath.toString)).!(
+ ProcessLogger(
+ out => queue.put(s"[pve] $out"),
+ err => queue.put(s"[pve][ERR] $err")
+ )
+ )
+
+ queue.put(s"[pve] venv creation finished with exit code $createCode")
+
+ if (createCode != 0) {
+ queue.put(s"[PVE][ERR] Failed to create venv (exit=$createCode)")
+ return
+ }
+
+ queue.put(
+ s"[PVE] Installing requirements from ${requirementsPath.toAbsolutePath} and ${operatorRequirementsPath.toAbsolutePath}"
+ )
+
+ val installReqCode = Process(
+ Seq(
+ python,
+ "-u",
+ "-m",
+ "pip",
+ "install",
+ "--progress-bar",
+ "off",
+ "-r",
+ requirementsPath.toString,
+ "-r",
+ operatorRequirementsPath.toString
+ ),
+ None,
+ envVars.toSeq: _*
+ ).!(
+ ProcessLogger(
+ out => queue.put(s"[pip] $out"),
+ err => queue.put(s"[pip][ERR] $err")
+ )
+ )
+
+ queue.put(s"[PVE] requirements install finished with exit code $installReqCode")
+
+ if (installReqCode != 0) {
+ queue.put(s"[PVE][ERR] Failed to install requirements files (exit=$installReqCode)")
+ return
+ }
+
+ queue.put(s"[PVE] Created new environment for cuid = $cuid")
+ }
+
+ def getEnvironments(cuid: Int): List[PvePackageResponse] = {
+
+ val cuPath = VenvRoot.resolve(cuid.toString)
+
+ if (!Files.isDirectory(cuPath)) {
+ return List()
+ }
+
+ val stream = Files.list(cuPath)
+
+ try {
+ stream
+ .iterator()
+ .asScala
+ .filter(path => Files.isDirectory(path))
+ .map { path =>
+ val pveName = path.getFileName.toString
+ val metadataPath = path.resolve("user-packages.txt")
+
+ val userPackages =
+ if (Files.exists(metadataPath)) {
+ Files
+ .readAllLines(metadataPath)
+ .asScala
+ .map(_.trim)
+ .filter(_.nonEmpty)
+ .toSeq
+ } else {
+ Seq()
+ }
+
+ PvePackageResponse(
+ pveName = pveName,
+ userPackages = userPackages
+ )
+ }
+ .toList
+ } finally {
+ stream.close()
+ }
+ }
+
+ // Deletes all PVE environments for a given CU (when running locally)
+ def deleteEnvironments(cuid: Int): Unit = {
+ val cuPath = VenvRoot.resolve(cuid.toString)
+
+ if (!Files.isDirectory(cuPath)) {
+ return
+ }
+
+ val stream = Files.walk(cuPath)
+
+ try {
+ stream
+ .sorted(Comparator.reverseOrder())
+ .iterator()
+ .asScala
+ .foreach(path => Files.deleteIfExists(path))
+ } finally {
+ stream.close()
+ }
+ }
+
+ /**
+ * Installs user requested Python packages into the PVE.
+ *
+ * 1. Executes pip install for each package
+ * 2. Updates user metadata file
+ * 3. Streams logs back via queue
+ */
+ def installUserPackages(
+ packages: List[String],
+ cuid: Int,
+ queue: BlockingQueue[String],
+ pveName: String
+ ): Unit = {
+
+ val python = pythonBinPath(cuid, pveName).toAbsolutePath.toString
+ val envVars = pipEnv
+
+ if (!Files.exists(Paths.get(python))) {
+ queue.put(s"[PVE][ERR] Python executable not found for PVE: $python")
+ return
+ }
+
+ val metadataPath = cuidDir(cuid, pveName).resolve("user-packages.txt")
+ Files.createDirectories(metadataPath.getParent)
+
+ var installedPackages =
+ if (Files.exists(metadataPath)) {
+ Files
+ .readAllLines(metadataPath)
+ .asScala
+ .map(_.trim)
+ .filter(_.nonEmpty)
+ .toSet
+ } else {
+ Set[String]()
+ }
+
+ packages.foreach { pkg =>
+ val trimmedPkg = pkg.trim
+
+ if (trimmedPkg.nonEmpty) {
+ queue.put(s"[PVE] Installing package: $trimmedPkg")
+
+ val code = Process(
+ Seq(
+ python,
+ "-u",
+ "-m",
+ "pip",
+ "install",
+ "--progress-bar",
+ "off",
+ "--no-input",
+ trimmedPkg
+ ),
+ None,
+ envVars.toSeq: _*
+ ).!(
+ ProcessLogger(
+ out => queue.put(s"[pip] $out"),
+ err => queue.put(s"[pip][ERR] $err")
+ )
+ )
+
+ queue.put(s"[pip] install($trimmedPkg) finished with exit code $code")
+
+ if (code != 0) {
+ queue.put(s"[PVE][ERR] Failed to install package: $trimmedPkg")
+ return
+ }
+
+ installedPackages = installedPackages + trimmedPkg
+
+ Files.write(
+ metadataPath,
+ installedPackages.toSeq.sorted.asJava
+ )
+ }
+ }
+
+ queue.put("[PVE] Final user package list:")
+
+ installedPackages.toSeq.sorted.foreach { pkg =>
+ queue.put(s"[user-package] $pkg")
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveResource.scala
new file mode 100644
index 00000000000..0a058ed6f5c
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveResource.scala
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.pythonvirtualenvironment
+
+import javax.ws.rs._
+import javax.ws.rs.core.MediaType
+import scala.jdk.CollectionConverters._
+import java.util
+
+@Path("/pve")
+@Consumes(Array(MediaType.APPLICATION_JSON))
+class PveResource {
+ // --------------------------------------------------
+ // Get system packages
+ // --------------------------------------------------
+ @GET
+ @Path("/system")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def getSystemPackages: util.Map[String, util.List[String]] = {
+ try {
+ val systemPkgs = PveManager.getSystemPackages().toList.asJava
+ Map("system" -> systemPkgs).asJava
+ } catch {
+ case e: Exception =>
+ e.printStackTrace()
+ throw new InternalServerErrorException("Failed to get system packages.")
+ }
+ }
+
+ // --------------------------------------------------
+ // Fetch PVEs and Installed User Packages
+ // --------------------------------------------------
+ @GET
+ @Path("/pves")
+ @Produces(Array(MediaType.APPLICATION_JSON))
+ def fetchPVEs(@QueryParam("cuid") cuid: Int): util.List[util.Map[String, Object]] = {
+ try {
+ PveManager
+ .getEnvironments(cuid)
+ .map { pve =>
+ Map(
+ "pveName" -> pve.pveName.asInstanceOf[Object],
+ "userPackages" -> pve.userPackages.asJava.asInstanceOf[Object]
+ ).asJava
+ }
+ .asJava
+
+ } catch {
+ case e: Exception =>
+ e.printStackTrace()
+ throw new InternalServerErrorException(s"Failed to get PVEs: ${e.getMessage}")
+ }
+ }
+
+ // --------------------------------------------------
+ // Delete PVEs
+ // --------------------------------------------------
+ @DELETE
+ @Path("/pves/{cuId}")
+ def deleteEnvironments(@PathParam("cuId") cuid: Int): Unit = {
+ PveManager.deleteEnvironments(cuid)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveWebsocketResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveWebsocketResource.scala
new file mode 100644
index 00000000000..577a8566ad5
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveWebsocketResource.scala
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.pythonvirtualenvironment
+
+import javax.websocket._
+import javax.websocket.server.ServerEndpoint
+import java.util.concurrent.LinkedBlockingQueue
+import scala.concurrent.Future
+import scala.concurrent.ExecutionContext.Implicits.global
+
+/**
+ * WebSocket endpoint for PVE creation and user pacakge installation that streams
+ * pip installation logs to the frontend in real time. The environment setup runs
+ * asynchronously, and output is pushed to the client until completion.
+ */
+
+@ServerEndpoint("/wsapi/pve")
+class PveWebsocketResource {
+
+ @OnOpen
+ def onOpen(session: Session): Unit = {
+
+ val params = session.getRequestParameterMap
+
+ val cuid = params.get("cuid").get(0).toInt
+ val pveName = params.get("pveName").get(0)
+ val isLocal = params.get("isLocal").get(0).toBoolean
+ val action = params.getOrDefault("action", java.util.List.of("create")).get(0)
+
+ val queue = new LinkedBlockingQueue[String]()
+
+ Future {
+ try {
+ action match {
+ case "create" =>
+ PveManager.createNewPve(cuid, queue, pveName, isLocal)
+
+ case "install" =>
+ val packages =
+ params
+ .getOrDefault("packages", java.util.List.of("[]"))
+ .get(0)
+ .stripPrefix("[")
+ .stripSuffix("]")
+ .split(",")
+ .toList
+ .map(_.replace("\"", "").trim)
+ .filter(_.nonEmpty)
+
+ PveManager.installUserPackages(packages, cuid, queue, pveName)
+
+ case _ =>
+ queue.put(s"[ERR] Unknown action: $action")
+ }
+ } catch {
+ case e: Exception =>
+ queue.put(s"[ERR] ${e.getMessage}")
+ } finally {
+ queue.put("__DONE__")
+ }
+ }
+
+ Future {
+ var done = false
+
+ while (!done && session.isOpen) {
+ val msg = queue.take()
+ session.getBasicRemote.sendText(msg)
+
+ if (msg == "__DONE__") {
+ done = true
+ session.close()
+ }
+ }
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/EmailNotificationService.scala b/amber/src/main/scala/org/apache/texera/web/service/EmailNotificationService.scala
new file mode 100644
index 00000000000..8f1c49c9595
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/EmailNotificationService.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+
+import java.util.concurrent.{ExecutorService, Executors}
+import scala.concurrent.{ExecutionContext, Future}
+
+trait EmailNotifier {
+ def shouldSendEmail(workflowState: WorkflowAggregatedState): Boolean
+
+ def sendStatusEmail(state: WorkflowAggregatedState): Unit
+}
+
+class EmailNotificationService(emailNotifier: EmailNotifier) {
+ private val executorService: ExecutorService = Executors.newSingleThreadExecutor()
+ private implicit val ec: ExecutionContext = ExecutionContext.fromExecutor(executorService)
+
+ def processEmailNotificationIfNeeded(
+ workflowState: WorkflowAggregatedState
+ ): Future[Unit] = {
+ Future {
+ if (emailNotifier.shouldSendEmail(workflowState)) {
+ emailNotifier.sendStatusEmail(workflowState)
+ }
+ }.recover {
+ case e: Exception =>
+ println(s"Failed to send email notification: ${e.getMessage}")
+ }
+ }
+
+ def shutdown(): Unit = {
+ executorService.shutdown()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/ExecutionConsoleService.scala b/amber/src/main/scala/org/apache/texera/web/service/ExecutionConsoleService.scala
new file mode 100644
index 00000000000..1678494e937
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/ExecutionConsoleService.scala
@@ -0,0 +1,332 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.google.protobuf.timestamp.Timestamp
+import com.twitter.util.{Await, Duration}
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.storage.model.BufferedItemWriter
+import org.apache.texera.amber.core.storage.result.ResultSchema
+import org.apache.texera.amber.core.storage.{DocumentFactory, VFSURIFactory}
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, OperatorIdentity}
+import org.apache.texera.amber.core.workflow.WorkflowContext
+import org.apache.texera.amber.engine.architecture.controller.ExecutionStateUpdate
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.ConsoleMessageType.COMMAND
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ ConsoleMessage,
+ ConsoleMessageType,
+ EvaluatePythonExpressionRequest,
+ DebugCommandRequest => AmberDebugCommandRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{
+ COMPLETED,
+ FAILED,
+ KILLED
+}
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.engine.common.executionruntimestate.{
+ EvaluatedValueList,
+ ExecutionConsoleStore,
+ OperatorConsole
+}
+import org.apache.texera.amber.util.VirtualIdentityUtils
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+import org.apache.texera.web.model.websocket.event.python.ConsoleUpdateEvent
+import org.apache.texera.web.model.websocket.request.RetryRequest
+import org.apache.texera.web.model.websocket.request.python.{
+ DebugCommandRequest,
+ PythonExpressionEvaluateRequest
+}
+import org.apache.texera.web.model.websocket.response.python.PythonExpressionEvaluateResponse
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+import org.apache.texera.web.storage.ExecutionStateStore
+import org.apache.texera.web.{SubscriptionManager, WebsocketInput}
+
+import java.time.Instant
+import java.util.concurrent.{ExecutorService, Executors}
+import scala.collection.mutable
+
+/**
+ * Utility object for processing console messages
+ * This is extracted to allow for easier testing and reuse
+ */
+object ConsoleMessageProcessor {
+
+ /**
+ * Processes a console message for display, performing truncation if needed.
+ *
+ * @param consoleMessage The original console message to process
+ * @param displayLength The maximum display length for the message title
+ * @return The truncated console message
+ */
+ def processConsoleMessage(
+ consoleMessage: ConsoleMessage,
+ displayLength: Int
+ ): ConsoleMessage = {
+ // Truncate message title if it exceeds the display length
+ val title = consoleMessage.title
+ if (title.getBytes.length > displayLength) {
+ val truncateIndicator = "..."
+ val truncatedTitle = title
+ .take(displayLength - truncateIndicator.length) + truncateIndicator
+ consoleMessage.copy(title = truncatedTitle)
+ } else {
+ consoleMessage
+ }
+ }
+
+ /**
+ * Updates the console store by adding a console message to an operator's console.
+ *
+ * @param consoleStore The console store to update
+ * @param opId The operator ID
+ * @param processedMessage The processed console message
+ * @param bufferSize The maximum number of messages to keep in the buffer
+ * @return The updated console store
+ */
+ def addMessageToOperatorConsole(
+ consoleStore: ExecutionConsoleStore,
+ opId: String,
+ processedMessage: ConsoleMessage,
+ bufferSize: Int
+ ): ExecutionConsoleStore = {
+ val opInfo = consoleStore.operatorConsole.getOrElse(opId, OperatorConsole())
+
+ val updatedOpInfo = if (opInfo.consoleMessages.size < bufferSize) {
+ opInfo.addConsoleMessages(processedMessage)
+ } else {
+ opInfo.withConsoleMessages(opInfo.consoleMessages.tail :+ processedMessage)
+ }
+
+ consoleStore.addOperatorConsole(opId -> updatedOpInfo)
+ }
+}
+
+class ExecutionConsoleService(
+ client: AmberClient,
+ stateStore: ExecutionStateStore,
+ wsInput: WebsocketInput,
+ workflowContext: WorkflowContext
+) extends SubscriptionManager
+ with LazyLogging {
+
+ registerCallbackOnPythonConsoleMessage()
+
+ val bufferSize: Int = ApplicationConfig.operatorConsoleBufferSize
+ val consoleMessageDisplayLength: Int = ApplicationConfig.consoleMessageDisplayLength
+
+ private val consoleMessageOpIdToWriterMap: mutable.Map[String, BufferedItemWriter[Tuple]] =
+ mutable.Map()
+
+ private val consoleWriterThread: ExecutorService = Executors.newSingleThreadExecutor()
+
+ private def getOrCreateWriter(opId: OperatorIdentity): BufferedItemWriter[Tuple] = {
+ consoleMessageOpIdToWriterMap.getOrElseUpdate(
+ opId.id, {
+ val uri = VFSURIFactory
+ .createConsoleMessagesURI(workflowContext.workflowId, workflowContext.executionId, opId)
+ val writer = DocumentFactory
+ .createDocument(uri, ResultSchema.consoleMessagesSchema)
+ .writer("console_messages")
+ .asInstanceOf[BufferedItemWriter[Tuple]]
+ WorkflowExecutionsResource.insertOperatorExecutions(
+ workflowContext.executionId.id,
+ opId.id,
+ uri
+ )
+ writer.open()
+ writer
+ }
+ )
+ }
+
+ addSubscription(
+ stateStore.consoleStore.registerDiffHandler((oldState, newState) => {
+ val output = new mutable.ArrayBuffer[TexeraWebSocketEvent]()
+ // For each operator, check if it has new python console message or breakpoint events
+ newState.operatorConsole
+ .foreach {
+ case (opId, info) =>
+ val oldConsole = oldState.operatorConsole.getOrElse(opId, new OperatorConsole())
+ val diff = info.consoleMessages.diff(oldConsole.consoleMessages)
+ output.append(ConsoleUpdateEvent(opId, diff))
+
+ info.evaluateExprResults.keys
+ .filterNot(oldConsole.evaluateExprResults.contains)
+ .foreach { key =>
+ output.append(
+ PythonExpressionEvaluateResponse(key, info.evaluateExprResults(key).values)
+ )
+ }
+ }
+ output
+ })
+ )
+
+ protected def registerCallbackOnPythonConsoleMessage(): Unit = {
+ addSubscription(
+ client
+ .registerCallback[ConsoleMessage]((evt: ConsoleMessage) => {
+ stateStore.consoleStore.updateState { consoleStore =>
+ val opId =
+ VirtualIdentityUtils.getPhysicalOpId(
+ ActorVirtualIdentity(evt.workerId)
+ )
+ addConsoleMessage(consoleStore, opId.logicalOpId.id, evt)
+ }
+ })
+ )
+
+ }
+
+ addSubscription(
+ client.registerCallback[ExecutionStateUpdate] {
+ case ExecutionStateUpdate(state: WorkflowAggregatedState.Recognized)
+ if Set(COMPLETED, FAILED, KILLED).contains(state) =>
+ logger.info("Workflow execution terminated. Commit console messages.")
+ consoleMessageOpIdToWriterMap.values.foreach { writer =>
+ try {
+ writer.close()
+ } catch {
+ case e: Exception =>
+ logger.error("Failed to close console message writer", e)
+ }
+ }
+ case _ =>
+ }
+ )
+
+ /**
+ * Processes a console message for display, performing truncation if needed.
+ * This method uses the shared implementation in ConsoleMessageProcessor.
+ *
+ * @param consoleMessage The original console message to process
+ * @return The truncated console message
+ */
+ def processConsoleMessage(consoleMessage: ConsoleMessage): ConsoleMessage = {
+ // Do not truncate debugger messages
+ if (consoleMessage.msgType == ConsoleMessageType.DEBUGGER) {
+ return consoleMessage
+ }
+ ConsoleMessageProcessor.processConsoleMessage(consoleMessage, consoleMessageDisplayLength)
+ }
+
+ /**
+ * Updates the console store by adding a console message to an operator's console.
+ * This method uses the shared implementation in ConsoleMessageProcessor.
+ *
+ * @param consoleStore The console store to update
+ * @param opId The operator ID
+ * @param processedMessage The processed console message
+ * @return The updated console store
+ */
+ def addMessageToOperatorConsole(
+ consoleStore: ExecutionConsoleStore,
+ opId: String,
+ processedMessage: ConsoleMessage
+ ): ExecutionConsoleStore = {
+ ConsoleMessageProcessor.addMessageToOperatorConsole(
+ consoleStore,
+ opId,
+ processedMessage,
+ bufferSize
+ )
+ }
+
+ private[this] def addConsoleMessage(
+ consoleStore: ExecutionConsoleStore,
+ opId: String,
+ consoleMessage: ConsoleMessage
+ ): ExecutionConsoleStore = {
+ // Write the original full message to the database
+ consoleWriterThread.execute(() => {
+ val writer = getOrCreateWriter(OperatorIdentity(opId))
+ try {
+ val tuple = new Tuple(
+ ResultSchema.consoleMessagesSchema,
+ Array(consoleMessage.toProtoString)
+ )
+ writer.putOne(tuple)
+ } catch {
+ case e: Exception =>
+ logger.error(s"Error while writing console message for operator $opId", e)
+ }
+ })
+
+ // Process the message (truncate if needed) and update store
+ val truncatedMessage = processConsoleMessage(consoleMessage)
+ addMessageToOperatorConsole(consoleStore, opId, truncatedMessage)
+ }
+
+ //Receive retry request
+ addSubscription(wsInput.subscribe((req: RetryRequest, uidOpt) => {
+ // empty implementation
+ }))
+
+ //Receive evaluate python expression
+ addSubscription(wsInput.subscribe((req: PythonExpressionEvaluateRequest, uidOpt) => {
+ val result = Await.result(
+ client.controllerInterface.evaluatePythonExpression(
+ EvaluatePythonExpressionRequest(req.expression, req.operatorId),
+ ()
+ ),
+ Duration.fromSeconds(10)
+ )
+ stateStore.consoleStore.updateState(consoleStore => {
+ val opInfo = consoleStore.operatorConsole.getOrElse(req.operatorId, OperatorConsole())
+ consoleStore.addOperatorConsole(
+ (
+ req.operatorId,
+ opInfo.addEvaluateExprResults((req.expression, EvaluatedValueList(result.values)))
+ )
+ )
+ })
+
+ // TODO: remove the following hack after fixing the frontend
+ // currently frontend is not prepared for re-receiving the eval-expr messages
+ // so we add it to the state and remove it from the state immediately
+ stateStore.consoleStore.updateState(consoleStore => {
+ val opInfo = consoleStore.operatorConsole.getOrElse(req.operatorId, OperatorConsole())
+ consoleStore.addOperatorConsole((req.operatorId, opInfo.clearEvaluateExprResults))
+ })
+ }))
+
+ //Receive debug command
+ addSubscription(wsInput.subscribe((req: DebugCommandRequest, uidOpt) => {
+ stateStore.consoleStore.updateState { consoleStore =>
+ val newMessage = new ConsoleMessage(
+ req.workerId,
+ Timestamp(Instant.now),
+ COMMAND,
+ "USER-" + uidOpt.getOrElse("UNKNOWN"),
+ req.cmd,
+ ""
+ )
+ addConsoleMessage(consoleStore, req.operatorId, newMessage)
+ }
+
+ client.controllerInterface.debugCommand(AmberDebugCommandRequest(req.workerId, req.cmd), ())
+
+ }))
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/ExecutionReconfigurationService.scala b/amber/src/main/scala/org/apache/texera/web/service/ExecutionReconfigurationService.scala
new file mode 100644
index 00000000000..e7617fdfe16
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/ExecutionReconfigurationService.scala
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.controller.{UpdateExecutorCompleted, Workflow}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ UpdateExecutorRequest,
+ WorkflowReconfigureRequest
+}
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.web.SubscriptionManager
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+import org.apache.texera.web.model.websocket.request.ModifyLogicRequest
+import org.apache.texera.web.model.websocket.response.{
+ ModifyLogicCompletedEvent,
+ ModifyLogicResponse
+}
+import org.apache.texera.web.storage.{ExecutionReconfigurationStore, ExecutionStateStore}
+
+import java.util.UUID
+import scala.util.{Failure, Success}
+
+class ExecutionReconfigurationService(
+ client: AmberClient,
+ stateStore: ExecutionStateStore,
+ workflow: Workflow
+) extends SubscriptionManager {
+
+ // monitors notification from the engine that a reconfiguration on a worker is completed
+ registerWorkerCompletionCallback()
+
+ // monitors the reconfiguration state (completed workers) change,
+ // notifies the frontend when all workers of an operator complete reconfiguration
+ registerCompletionDiffHandler()
+
+ // handles reconfigure workflow logic from frontend
+ // validate the modify logic request and notifies the frontend
+ // reconfigurations can only come when the workflow is paused,
+ // they are not actually performed until the workflow is resumed
+ def modifyOperatorLogic(modifyLogicRequest: ModifyLogicRequest): TexeraWebSocketEvent = {
+ val newOp = modifyLogicRequest.operator
+ val opId = newOp.operatorIdentifier
+ val currentOp = workflow.logicalPlan.getOperator(opId)
+ val reconfiguredPhysicalOp =
+ currentOp.runtimeReconfiguration(
+ workflow.context.workflowId,
+ workflow.context.executionId,
+ currentOp,
+ newOp
+ )
+ reconfiguredPhysicalOp match {
+ case Failure(exception) => ModifyLogicResponse(opId.id, isValid = false, exception.getMessage)
+ case Success(op) => {
+ stateStore.reconfigurationStore.updateState(old =>
+ old.copy(unscheduledReconfigurations = old.unscheduledReconfigurations :+ op)
+ )
+ ModifyLogicResponse(opId.id, isValid = true, "")
+ }
+ }
+ }
+
+ // actually performs all reconfiguration requests the user made during pause
+ // sends ModifyLogic messages to operators and workers,
+ // see the Fries reconfiguration paper for the algorithm.
+ // Note: StateTransferFunc is currently not threaded through to the engine —
+ // the new UpdateExecutorRequest only carries (targetOpId, newOpExecInitInfo).
+ def performReconfigurationOnResume(): Unit = {
+ val reconfigurations = stateStore.reconfigurationStore.getState.unscheduledReconfigurations
+ if (reconfigurations.isEmpty) {
+ return
+ }
+
+ val reconfigurationId = UUID.randomUUID().toString
+ val updateExecutorRequests = reconfigurations.map {
+ case (op, _) => UpdateExecutorRequest(op.id, op.opExecInitInfo)
+ }
+ dispatch(
+ WorkflowReconfigureRequest(
+ reconfiguration = updateExecutorRequests,
+ reconfigurationId = reconfigurationId
+ )
+ )
+
+ // clear all un-scheduled reconfigurations, start a new reconfiguration ID
+ stateStore.reconfigurationStore.updateState(_ =>
+ ExecutionReconfigurationStore(currentReconfigId = Some(reconfigurationId))
+ )
+ }
+
+ // Seam for unit testing the dispatch path without spinning up an AmberClient.
+ protected def dispatch(request: WorkflowReconfigureRequest): Unit = {
+ client.controllerInterface.reconfigureWorkflow(request, ())
+ }
+
+ // Seam for unit testing — production wires the engine's UpdateExecutorCompleted
+ // events into the reconfiguration store so the diff handler above can fire
+ // ModifyLogicCompletedEvent for the frontend.
+ protected def registerWorkerCompletionCallback(): Unit = {
+ client.registerCallback[UpdateExecutorCompleted]((evt: UpdateExecutorCompleted) => {
+ onWorkerReconfigured(evt.id)
+ })
+ }
+
+ // Exposed (instead of inlined in the callback) so tests can drive the
+ // completion path directly.
+ private[service] def onWorkerReconfigured(worker: ActorVirtualIdentity): Unit = {
+ stateStore.reconfigurationStore.updateState(old =>
+ old.copy(completedReconfigurations = old.completedReconfigurations + worker)
+ )
+ }
+
+ // Seam for unit testing — the diff handler dereferences workflow.physicalPlan
+ // to map worker → logical op, which makes constructing a service in tests
+ // require a full Workflow. Tests override to no-op.
+ protected def registerCompletionDiffHandler(): Unit = {
+ addSubscription(
+ stateStore.reconfigurationStore.registerDiffHandler((oldState, newState) => {
+ if (
+ oldState.completedReconfigurations != newState.completedReconfigurations
+ && oldState.currentReconfigId == newState.currentReconfigId
+ ) {
+ val diff = newState.completedReconfigurations -- oldState.completedReconfigurations
+ val newlyCompletedOps = diff
+ .map(workerId => workflow.physicalPlan.getPhysicalOpByWorkerId(workerId).id)
+ .map(opId => opId.logicalOpId.id)
+ if (newlyCompletedOps.nonEmpty) {
+ List(ModifyLogicCompletedEvent(newlyCompletedOps.toList))
+ } else {
+ List()
+ }
+ } else {
+ List()
+ }
+ })
+ )
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/ExecutionResultService.scala b/amber/src/main/scala/org/apache/texera/web/service/ExecutionResultService.scala
new file mode 100644
index 00000000000..3f0362f8242
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/ExecutionResultService.scala
@@ -0,0 +1,500 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import org.apache.pekko.actor.Cancellable
+import com.fasterxml.jackson.annotation.{JsonTypeInfo, JsonTypeName}
+import com.fasterxml.jackson.databind.node.ObjectNode
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.storage.model.VirtualDocument
+import org.apache.texera.amber.core.storage.result._
+import org.apache.texera.amber.core.storage.{DocumentFactory, VFSURIFactory}
+import org.apache.texera.amber.core.tuple.{AttributeType, Tuple, TupleUtils}
+import org.apache.texera.amber.core.virtualidentity.{
+ ExecutionIdentity,
+ OperatorIdentity,
+ WorkflowIdentity
+}
+import org.apache.texera.amber.core.workflow.OutputPort.OutputMode
+import org.apache.texera.amber.core.workflow.{PhysicalOp, PhysicalPlan, PortIdentity}
+import org.apache.texera.amber.engine.architecture.controller.{ExecutionStateUpdate, FatalError}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{
+ COMPLETED,
+ FAILED,
+ KILLED,
+ RUNNING,
+ TERMINATED
+}
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.engine.common.executionruntimestate.ExecutionMetadataStore
+import org.apache.texera.web.SubscriptionManager
+import org.apache.texera.web.model.websocket.event.{
+ PaginatedResultEvent,
+ TexeraWebSocketEvent,
+ WebResultUpdateEvent
+}
+import org.apache.texera.web.model.websocket.request.ResultPaginationRequest
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+import org.apache.texera.web.service.ExecutionResultService.convertTuplesToJson
+import org.apache.texera.web.service.WorkflowExecutionService.getLatestExecutionId
+import org.apache.texera.web.storage.{ExecutionStateStore, WorkflowStateStore}
+
+import java.util.UUID
+import scala.collection.mutable
+import scala.concurrent.duration.DurationInt
+
+object ExecutionResultService {
+
+ private val defaultPageSize: Int = 5
+
+ /**
+ * Converts a collection of Tuples to a list of JSON ObjectNodes.
+ *
+ * This function takes a collection of Tuples and converts each tuple into a JSON ObjectNode.
+ * For binary data, it formats the bytes into a readable hex string representation with length info.
+ * For string values longer than maxStringLength (100), it truncates them.
+ * NULL values are converted to the string "NULL".
+ *
+ * @param tuples The collection of Tuples to convert
+ * @param isVisualization Whether this is for visualization rendering (affects string truncation)
+ * @return A List of ObjectNodes containing the JSON representation of the tuples
+ */
+ def convertTuplesToJson(
+ tuples: Iterable[Tuple],
+ isVisualization: Boolean = false
+ ): List[ObjectNode] = {
+ val maxStringLength = 100
+
+ tuples.map { tuple =>
+ val processedFields = tuple.schema.getAttributes.zipWithIndex
+ .map {
+ case (attr, idx) =>
+ val fieldValue = tuple.getField[AnyRef](idx)
+
+ Option(fieldValue) match {
+ case None => "NULL"
+ case Some(value) =>
+ attr.getType match {
+ case AttributeType.BINARY =>
+ value match {
+ case byteArray: Array[Byte] =>
+ val totalSize = byteArray.length
+ val hexString = byteArrayToHexString(byteArray)
+
+ // 39 = 30 (leading bytes) + 9 (trailing bytes)
+ // 30 bytes = space for 10 hex values (each hex value takes 2 chars + 1 space)
+ // 9 bytes = space for 3 hex values at the end (2 chars each + 1 space)
+ if (hexString.length < 39) {
+ s"bytes'$hexString' (length: $totalSize)"
+ } else {
+ val leadingBytes = hexString.take(30)
+ val trailingBytes = hexString.takeRight(9)
+ s"bytes'$leadingBytes...$trailingBytes' (length: $totalSize)"
+ }
+
+ case _ =>
+ throw new RuntimeException(
+ s"Expected byte array for binary type field, but got: ${value.getClass.getName}"
+ )
+ }
+ case AttributeType.STRING =>
+ val stringValue = value.asInstanceOf[String]
+ if (stringValue.length > maxStringLength && !isVisualization)
+ stringValue.take(maxStringLength) + "..."
+ else
+ stringValue
+ case _ => value
+ }
+ }
+ }
+ .toArray[Any]
+
+ TupleUtils.tuple2json(tuple.schema, processedFields)
+ }.toList
+ }
+
+ /**
+ * Converts a byte array to a hex string representation.
+ *
+ * This helper function takes a byte array and converts its contents to a space-separated
+ * string of hexadecimal values. Each byte is formatted as a two-digit uppercase hex number.
+ *
+ * @param byteArray The byte array to convert
+ * @return A string containing the hex representation of the byte array's contents
+ */
+ private def byteArrayToHexString(byteArray: Array[Byte]): String = {
+ byteArray.map(b => String.format("%02X", Byte.box(b))).mkString(" ")
+ }
+
+ /**
+ * convert Tuple from engine's format to JSON format
+ */
+ private def tuplesToWebData(
+ mode: WebOutputMode,
+ table: List[Tuple]
+ ): WebDataUpdate = {
+ val tableInJson = convertTuplesToJson(table, mode == SetSnapshotMode())
+ WebDataUpdate(mode, tableInJson)
+ }
+
+ /**
+ * For SET_SNAPSHOT output mode: result is the latest snapshot
+ * FOR SET_DELTA output mode:
+ * - for insert-only delta: effectively the same as latest snapshot
+ * - for insert-retract delta: the union of all delta outputs, not compacted to a snapshot
+ *
+ * Produces the WebResultUpdate to send to frontend from a result update from the engine.
+ */
+ private def convertWebResultUpdate(
+ workflowIdentity: WorkflowIdentity,
+ executionId: ExecutionIdentity,
+ physicalOps: List[PhysicalOp],
+ oldTupleCount: Int,
+ newTupleCount: Int
+ ): WebResultUpdate = {
+ val outputMode = physicalOps
+ .flatMap(op => op.outputPorts)
+ .filter({
+ case (portId, (port, links, schema)) => !portId.internal
+ })
+ .map({
+ case (portId, (port, links, schema)) => port.mode
+ })
+ .head
+
+ val webOutputMode: WebOutputMode = {
+ outputMode match {
+ // currently, only table outputs are using these modes
+ case OutputMode.SET_DELTA => SetDeltaMode()
+ case OutputMode.SET_SNAPSHOT => PaginationMode()
+
+ // currently, only visualizations are using single snapshot mode
+ case OutputMode.SINGLE_SNAPSHOT => SetSnapshotMode()
+ case OutputMode.Unrecognized(_) =>
+ throw new RuntimeException(
+ s"Unrecognized output mode: $outputMode for workflow ${workflowIdentity.id}"
+ )
+ }
+ }
+
+ // Cannot assume the storage is available at this point. The storage object is only available
+ // after a region is scheduled to execute.
+ val storageUriOption = WorkflowExecutionsResource.getResultUriByLogicalPortId(
+ executionId,
+ physicalOps.head.id.logicalOpId,
+ PortIdentity()
+ )
+ storageUriOption match {
+ case Some(storageUri) =>
+ val storage: VirtualDocument[Tuple] =
+ DocumentFactory.openDocument(storageUri)._1.asInstanceOf[VirtualDocument[Tuple]]
+ val webUpdate = webOutputMode match {
+ case PaginationMode() =>
+ val numTuples = storage.getCount
+ val maxPageIndex =
+ Math.ceil(numTuples / defaultPageSize.toDouble).toInt
+ // This can be extremly expensive when we have a lot of pages.
+ // It causes delays in some obseved cases.
+ // TODO: try to optimize this.
+ WebPaginationUpdate(
+ PaginationMode(),
+ newTupleCount,
+ (1 to maxPageIndex).toList
+ )
+ case SetSnapshotMode() =>
+ tuplesToWebData(webOutputMode, storage.get().toList)
+ case SetDeltaMode() =>
+ val deltaList = storage.getAfter(oldTupleCount).toList
+ tuplesToWebData(webOutputMode, deltaList)
+
+ case _ =>
+ throw new RuntimeException(
+ "update mode combination not supported: " + (webOutputMode, outputMode)
+ )
+ }
+ webUpdate
+ case None =>
+ WebPaginationUpdate(
+ PaginationMode(),
+ 0,
+ List.empty
+ )
+ }
+ }
+
+ /**
+ * Behavior for different web output modes:
+ * - PaginationMode (used by view result operator)
+ * - send new number of tuples and dirty page index
+ * - SetSnapshotMode (used by visualization in snapshot mode)
+ * - send entire snapshot result to frontend
+ * - SetDeltaMode (used by visualization in delta mode)
+ * - send incremental delta result to frontend
+ */
+ @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
+ sealed abstract class WebOutputMode extends Product with Serializable
+
+ /**
+ * The result update of one operator that will be sent to the frontend.
+ * Can be either WebPaginationUpdate (for PaginationMode)
+ * or WebDataUpdate (for SetSnapshotMode or SetDeltaMode)
+ */
+ sealed abstract class WebResultUpdate extends Product with Serializable
+
+ @JsonTypeName("PaginationMode")
+ final case class PaginationMode() extends WebOutputMode
+
+ @JsonTypeName("SetSnapshotMode")
+ final case class SetSnapshotMode() extends WebOutputMode
+
+ @JsonTypeName("SetDeltaMode")
+ final case class SetDeltaMode() extends WebOutputMode
+
+ case class WebPaginationUpdate(
+ mode: PaginationMode,
+ totalNumTuples: Long,
+ dirtyPageIndices: List[Int]
+ ) extends WebResultUpdate
+
+ case class WebDataUpdate(mode: WebOutputMode, table: List[ObjectNode]) extends WebResultUpdate
+
+}
+
+/**
+ * ExecutionResultService manages all operator output ports that have storage in one workflow execution.
+ *
+ * On each result update from the engine, WorkflowResultService
+ * - update the result data for each operator,
+ * - send result update event to the frontend
+ */
+class ExecutionResultService(
+ workflowIdentity: WorkflowIdentity,
+ computingUnitId: Int,
+ val workflowStateStore: WorkflowStateStore
+) extends SubscriptionManager
+ with LazyLogging {
+ private val resultPullingFrequency = ApplicationConfig.executionResultPollingInSecs
+ private var resultUpdateCancellable: Cancellable = _
+
+ def attachToExecution(
+ executionId: ExecutionIdentity,
+ stateStore: ExecutionStateStore,
+ physicalPlan: PhysicalPlan,
+ client: AmberClient
+ ): Unit = {
+ if (resultUpdateCancellable != null && !resultUpdateCancellable.isCancelled) {
+ resultUpdateCancellable.cancel()
+ }
+
+ unsubscribeAll()
+
+ addSubscription(stateStore.metadataStore.getStateObservable.subscribe {
+ newState: ExecutionMetadataStore =>
+ {
+ if (newState.state == RUNNING) {
+ if (resultUpdateCancellable == null || resultUpdateCancellable.isCancelled) {
+ resultUpdateCancellable = AmberRuntime
+ .scheduleRecurringCallThroughActorSystem(
+ 2.seconds,
+ resultPullingFrequency.seconds
+ ) {
+ onResultUpdate(executionId, physicalPlan)
+ }
+ }
+ } else {
+ if (resultUpdateCancellable != null) resultUpdateCancellable.cancel()
+ }
+ }
+ })
+
+ addSubscription(
+ client
+ .registerCallback[ExecutionStateUpdate](evt => {
+ if (
+ evt.state == COMPLETED || evt.state == FAILED || evt.state == KILLED || evt.state == TERMINATED
+ ) {
+ logger.info("Workflow execution terminated. Stop update results.")
+ if (resultUpdateCancellable.cancel() || resultUpdateCancellable.isCancelled) {
+ // immediately perform final update
+ onResultUpdate(executionId, physicalPlan)
+ }
+ }
+ })
+ )
+
+ addSubscription(
+ client.registerCallback[FatalError](_ =>
+ if (resultUpdateCancellable != null) {
+ resultUpdateCancellable.cancel()
+ }
+ )
+ )
+
+ addSubscription(
+ workflowStateStore.resultStore.registerDiffHandler((oldState, newState) => {
+ val buf = mutable.HashMap[String, ExecutionResultService.WebResultUpdate]()
+ val allTableStats = mutable.Map[String, Map[String, Map[String, Any]]]()
+ newState.resultInfo
+ .filter(info => {
+ // only update those operators with changing tuple count.
+ !oldState.resultInfo
+ .contains(info._1) || oldState.resultInfo(info._1).tupleCount != info._2.tupleCount
+ })
+ .foreach {
+ case (opId, info) =>
+ val oldInfo = oldState.resultInfo.getOrElse(opId, OperatorResultMetadata())
+ buf(opId.id) = ExecutionResultService.convertWebResultUpdate(
+ workflowIdentity,
+ executionId,
+ physicalPlan.getPhysicalOpsOfLogicalOp(opId),
+ oldInfo.tupleCount,
+ info.tupleCount
+ )
+ // using the first port for now. TODO: support multiple ports
+ val outputPortsMap = physicalPlan
+ .getPhysicalOpsOfLogicalOp(opId)
+ .headOption
+ .map(_.outputPorts)
+ .getOrElse(Map.empty)
+ val hasSingleSnapshot = outputPortsMap.values.exists {
+ case (outputPort, _, _) =>
+ // SINGLE_SNAPSHOT is used for HTML content
+ outputPort.mode == OutputMode.SINGLE_SNAPSHOT
+ }
+
+ if (!hasSingleSnapshot) {
+ val storageUri = WorkflowExecutionsResource
+ .getResultUriByLogicalPortId(
+ executionId,
+ opId,
+ PortIdentity()
+ )
+
+ if (storageUri.nonEmpty) {
+ val (_, _, globalPortIdOption, _) = VFSURIFactory.decodeURI(storageUri.get)
+ val opStorage = DocumentFactory.openDocument(storageUri.get)._1
+
+ allTableStats(opId.id) = opStorage.getTableStatistics
+ WorkflowExecutionsResource.updateResultSize(
+ executionId,
+ globalPortIdOption.get,
+ opStorage.getTotalFileSize
+ )
+ WorkflowExecutionsResource.updateRuntimeStatsSize(executionId)
+ WorkflowExecutionsResource.updateConsoleMessageSize(executionId, opId)
+ }
+ }
+ }
+ Iterable(
+ WebResultUpdateEvent(
+ buf.toMap,
+ allTableStats.toMap
+ )
+ )
+ })
+ )
+
+ // clear all the result metadata
+ workflowStateStore.resultStore.updateState { _ =>
+ WorkflowResultStore() // empty result store
+ }
+
+ }
+
+ def handleResultPagination(request: ResultPaginationRequest): TexeraWebSocketEvent = {
+ // calculate from index (pageIndex starts from 1 instead of 0)
+ val from = request.pageSize * (request.pageIndex - 1)
+ val latestExecutionId = getLatestExecutionId(workflowIdentity, computingUnitId).getOrElse(
+ throw new IllegalStateException("No execution is recorded")
+ )
+
+ val storageUriOption = WorkflowExecutionsResource.getResultUriByLogicalPortId(
+ latestExecutionId,
+ OperatorIdentity(request.operatorID),
+ PortIdentity()
+ )
+
+ storageUriOption match {
+ case Some(storageUri) =>
+ val (document, schemaOption) = DocumentFactory.openDocument(storageUri)
+ val virtualDocument = document.asInstanceOf[VirtualDocument[Tuple]]
+
+ val columns = {
+ val schema = schemaOption.get
+ val allColumns = schema.getAttributeNames
+ val filteredColumns = request.columnSearch match {
+ case Some(search) =>
+ allColumns.filter(col => col.toLowerCase.contains(search.toLowerCase))
+ case None => allColumns
+ }
+ Some(
+ filteredColumns.slice(request.columnOffset, request.columnOffset + request.columnLimit)
+ )
+ }
+
+ val paginationIterable = {
+ virtualDocument
+ .getRange(from, from + request.pageSize, columns)
+ .to(Iterable)
+ }
+ val mappedResults = convertTuplesToJson(paginationIterable)
+ val attributes = paginationIterable.headOption
+ .map(_.getSchema.getAttributes)
+ .getOrElse(List.empty)
+ PaginatedResultEvent.apply(request, mappedResults, attributes)
+
+ case None =>
+ // Handle the case when storageUri is empty
+ PaginatedResultEvent.apply(request, List.empty, List.empty)
+ }
+ }
+
+ private def onResultUpdate(executionId: ExecutionIdentity, physicalPlan: PhysicalPlan): Unit = {
+ workflowStateStore.resultStore.updateState { _ =>
+ val newInfo: Map[OperatorIdentity, OperatorResultMetadata] = {
+ WorkflowExecutionsResource
+ .getResultUrisByExecutionId(executionId)
+ .map(uri => {
+ val count = DocumentFactory.openDocument(uri)._1.getCount.toInt
+
+ val (_, _, globalPortIdOption, _) = VFSURIFactory.decodeURI(uri)
+
+ // Retrieve the mode of the specified output port
+ val mode = physicalPlan
+ .getPhysicalOpsOfLogicalOp(globalPortIdOption.get.opId.logicalOpId)
+ .flatMap(_.outputPorts.get(globalPortIdOption.get.portId))
+ .map(_._1.mode)
+ .head
+
+ val changeDetector =
+ if (mode == OutputMode.SET_SNAPSHOT) {
+ UUID.randomUUID.toString
+ } else ""
+ (globalPortIdOption.get.opId.logicalOpId, OperatorResultMetadata(count, changeDetector))
+ })
+ .toMap
+ }
+ WorkflowResultStore(newInfo)
+ }
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/ExecutionRuntimeService.scala b/amber/src/main/scala/org/apache/texera/web/service/ExecutionRuntimeService.scala
new file mode 100644
index 00000000000..70b2f07920a
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/ExecutionRuntimeService.scala
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.virtualidentity.EmbeddedControlMessageIdentity
+import org.apache.texera.amber.engine.architecture.controller.ExecutionStateUpdate
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ EmptyRequest,
+ TakeGlobalCheckpointRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState._
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.FaultToleranceConfig
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.web.model.websocket.request._
+import org.apache.texera.web.storage.ExecutionStateStore
+import org.apache.texera.web.storage.ExecutionStateStore.updateWorkflowState
+import org.apache.texera.web.{SubscriptionManager, WebsocketInput}
+
+import java.net.URI
+import java.util.UUID
+
+class ExecutionRuntimeService(
+ client: AmberClient,
+ stateStore: ExecutionStateStore,
+ wsInput: WebsocketInput,
+ reconfigurationService: ExecutionReconfigurationService,
+ logConf: Option[FaultToleranceConfig],
+ workflowId: Long,
+ emailNotificationEnabled: Boolean,
+ userEmailOpt: Option[String],
+ sessionUri: URI
+) extends SubscriptionManager
+ with LazyLogging {
+
+ private val emailNotificationService = for {
+ email <- userEmailOpt
+ if emailNotificationEnabled
+ } yield new EmailNotificationService(
+ new WorkflowEmailNotifier(
+ workflowId,
+ email,
+ sessionUri
+ )
+ )
+
+ //Receive skip tuple
+ addSubscription(wsInput.subscribe((req: SkipTupleRequest, uidOpt) => {
+ throw new RuntimeException("skipping tuple is temporarily disabled")
+ }))
+
+ // Receive execution state update from Amber
+ addSubscription(client.registerCallback[ExecutionStateUpdate]((evt: ExecutionStateUpdate) => {
+ stateStore.metadataStore.updateState(metadataStore =>
+ updateWorkflowState(evt.state, metadataStore)
+ )
+
+ emailNotificationService.foreach(_.processEmailNotificationIfNeeded(evt.state))
+
+ if (evt.state == COMPLETED) {
+ client.shutdown()
+ stateStore.statsStore.updateState(stats => stats.withEndTimeStamp(System.currentTimeMillis()))
+ }
+ }))
+
+ // Receive Pause
+ addSubscription(wsInput.subscribe((req: WorkflowPauseRequest, uidOpt) => {
+ stateStore.metadataStore.updateState(metadataStore =>
+ updateWorkflowState(PAUSING, metadataStore)
+ )
+ client.controllerInterface.pauseWorkflow(EmptyRequest(), ())
+ }))
+
+ // Receive Resume
+ addSubscription(wsInput.subscribe((req: WorkflowResumeRequest, uidOpt) => {
+ reconfigurationService.performReconfigurationOnResume()
+ stateStore.metadataStore.updateState(metadataStore =>
+ updateWorkflowState(RESUMING, metadataStore)
+ )
+ client.controllerInterface
+ .resumeWorkflow(EmptyRequest(), ())
+ .onSuccess(_ =>
+ stateStore.metadataStore.updateState(metadataStore =>
+ updateWorkflowState(RUNNING, metadataStore)
+ )
+ )
+ }))
+
+ // Receive Kill
+ addSubscription(wsInput.subscribe((req: WorkflowKillRequest, uidOpt) => {
+ client.shutdown()
+ stateStore.statsStore.updateState(stats => stats.withEndTimeStamp(System.currentTimeMillis()))
+ stateStore.metadataStore.updateState(metadataStore =>
+ updateWorkflowState(KILLED, metadataStore)
+ )
+ }))
+
+ // Receive Interaction
+ addSubscription(wsInput.subscribe((req: WorkflowCheckpointRequest, uidOpt) => {
+ assert(
+ logConf.nonEmpty,
+ "Fault tolerance log folder is not established. Unable to take a global checkpoint."
+ )
+ val checkpointId = EmbeddedControlMessageIdentity(s"Checkpoint_${UUID.randomUUID().toString}")
+ val uri = logConf.get.writeTo.resolve(checkpointId.toString)
+ client.controllerInterface.takeGlobalCheckpoint(
+ TakeGlobalCheckpointRequest(estimationOnly = false, checkpointId, uri.toString),
+ ()
+ )
+ }))
+
+ override def unsubscribeAll(): Unit = {
+ super.unsubscribeAll()
+ emailNotificationService.foreach(_.shutdown())
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/ExecutionStatsService.scala b/amber/src/main/scala/org/apache/texera/web/service/ExecutionStatsService.scala
new file mode 100644
index 00000000000..3703a2bf417
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/ExecutionStatsService.scala
@@ -0,0 +1,340 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.google.protobuf.timestamp.Timestamp
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.storage.model.BufferedItemWriter
+import org.apache.texera.amber.core.storage.result.ResultSchema
+import org.apache.texera.amber.core.storage.{DocumentFactory, VFSURIFactory}
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.workflow.WorkflowContext
+import org.apache.texera.amber.core.workflowruntimestate.FatalErrorType.EXECUTION_FAILURE
+import org.apache.texera.amber.core.workflowruntimestate.WorkflowFatalError
+import org.apache.texera.amber.engine.architecture.controller.{
+ ExecutionStateUpdate,
+ ExecutionStatsUpdate,
+ FatalError,
+ RuntimeStatisticsPersist,
+ WorkerAssignmentUpdate,
+ WorkflowRecoveryStatus
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{
+ COMPLETED,
+ FAILED,
+ KILLED
+}
+import org.apache.texera.amber.engine.common.Utils
+import org.apache.texera.amber.engine.common.Utils.maptoStatusCode
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.engine.common.executionruntimestate.{
+ OperatorMetrics,
+ OperatorStatistics,
+ OperatorWorkerMapping
+}
+import org.apache.texera.amber.error.ErrorUtils.{
+ getOperatorFromActorIdOpt,
+ getStackTraceWithAllCauses
+}
+import org.apache.texera.web.SubscriptionManager
+import org.apache.texera.web.model.websocket.event.{
+ ExecutionDurationUpdateEvent,
+ OperatorAggregatedMetrics,
+ OperatorStatisticsUpdateEvent,
+ WorkerAssignmentUpdateEvent
+}
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+import org.apache.texera.web.storage.ExecutionStateStore
+import org.apache.texera.web.storage.ExecutionStateStore.updateWorkflowState
+
+import java.time.Instant
+import java.util.concurrent.Executors
+
+class ExecutionStatsService(
+ client: AmberClient,
+ stateStore: ExecutionStateStore,
+ workflowContext: WorkflowContext
+) extends SubscriptionManager
+ with LazyLogging {
+ private val (metricsPersistThread, runtimeStatsWriter) = {
+ val thread = Executors.newSingleThreadExecutor()
+ val uri = VFSURIFactory.createRuntimeStatisticsURI(
+ workflowContext.workflowId,
+ workflowContext.executionId
+ )
+ val writer = DocumentFactory
+ .createDocument(uri, ResultSchema.runtimeStatisticsSchema)
+ .writer("runtime_statistics")
+ .asInstanceOf[BufferedItemWriter[Tuple]]
+ WorkflowExecutionsResource.updateRuntimeStatsUri(
+ workflowContext.workflowId.id,
+ workflowContext.executionId.id,
+ uri
+ )
+ writer.open()
+ (thread, writer)
+ }
+
+ private var lastPersistedMetrics: Map[String, OperatorMetrics] =
+ Map.empty[String, OperatorMetrics]
+
+ registerCallbacks()
+
+ addSubscription(
+ stateStore.statsStore.registerDiffHandler((oldState, newState) => {
+ // Update operator stats if any operator updates its stat
+ if (newState.operatorInfo.toSet != oldState.operatorInfo.toSet) {
+ Iterable(
+ OperatorStatisticsUpdateEvent(newState.operatorInfo.collect {
+ case x =>
+ val metrics = x._2
+ val inMap = metrics.operatorStatistics.inputMetrics
+ .map(pm => pm.portId.id.toString -> pm.tupleMetrics.count)
+ .toMap
+ val outMap = metrics.operatorStatistics.outputMetrics
+ .map(pm => pm.portId.id.toString -> pm.tupleMetrics.count)
+ .toMap
+
+ val res = OperatorAggregatedMetrics(
+ Utils.aggregatedStateToString(metrics.operatorState),
+ metrics.operatorStatistics.inputMetrics.map(_.tupleMetrics.count).sum,
+ metrics.operatorStatistics.inputMetrics.map(_.tupleMetrics.size).sum,
+ inMap,
+ metrics.operatorStatistics.outputMetrics.map(_.tupleMetrics.count).sum,
+ metrics.operatorStatistics.outputMetrics.map(_.tupleMetrics.size).sum,
+ outMap,
+ metrics.operatorStatistics.numWorkers,
+ metrics.operatorStatistics.dataProcessingTime,
+ metrics.operatorStatistics.controlProcessingTime,
+ metrics.operatorStatistics.idleTime
+ )
+ (x._1, res)
+ })
+ )
+ } else {
+ Iterable.empty
+ }
+ })
+ )
+
+ addSubscription(
+ stateStore.statsStore.registerDiffHandler((oldState, newState) => {
+ // update operators' workers.
+ if (newState.operatorWorkerMapping != oldState.operatorWorkerMapping) {
+ newState.operatorWorkerMapping
+ .map { opToWorkers =>
+ WorkerAssignmentUpdateEvent(opToWorkers.operatorId, opToWorkers.workerIds)
+ }
+ } else {
+ Iterable()
+ }
+ })
+ )
+
+ addSubscription(
+ stateStore.statsStore.registerDiffHandler((oldState, newState) => {
+ // update execution duration.
+ if (
+ newState.startTimeStamp != oldState.startTimeStamp || newState.endTimeStamp != oldState.endTimeStamp
+ ) {
+ if (newState.endTimeStamp != 0) {
+ Iterable(
+ ExecutionDurationUpdateEvent(
+ newState.endTimeStamp - newState.startTimeStamp,
+ isRunning = false
+ )
+ )
+ } else {
+ val currentTime = System.currentTimeMillis()
+ Iterable(
+ ExecutionDurationUpdateEvent(currentTime - newState.startTimeStamp, isRunning = true)
+ )
+ }
+ } else {
+ Iterable()
+ }
+ })
+ )
+
+ private[this] def registerCallbacks(): Unit = {
+ registerCallbackOnWorkflowStatsUpdate()
+ registerCallbackOnWorkerAssignedUpdate()
+ registerCallbackOnWorkflowRecoveryUpdate()
+ registerCallbackOnFatalError()
+ }
+
+ private[this] def registerCallbackOnWorkflowStatsUpdate(): Unit = {
+ // Register callback for UI updates (UI state store update only, no persistence)
+ addSubscription(
+ client
+ .registerCallback[ExecutionStatsUpdate]((evt: ExecutionStatsUpdate) => {
+ stateStore.statsStore.updateState { statsStore =>
+ statsStore.withOperatorInfo(evt.operatorMetrics)
+ }
+ })
+ )
+
+ // Register callback for statistics persistence (persistence only, no UI update)
+ addSubscription(
+ client
+ .registerCallback[RuntimeStatisticsPersist]((evt: RuntimeStatisticsPersist) => {
+ metricsPersistThread.execute(() => {
+ storeRuntimeStatistics(computeStatsDiff(evt.operatorMetrics))
+ lastPersistedMetrics = evt.operatorMetrics
+ })
+ })
+ )
+ }
+
+ addSubscription(
+ client.registerCallback[ExecutionStateUpdate] {
+ case ExecutionStateUpdate(state: WorkflowAggregatedState.Recognized)
+ if Set(COMPLETED, FAILED, KILLED).contains(state) =>
+ logger.info("Workflow execution terminated. Commit runtime statistics.")
+ try {
+ runtimeStatsWriter.close()
+ } catch {
+ case e: Exception =>
+ logger.error("Failed to close runtime statistics writer", e)
+ }
+ case _ =>
+ }
+ )
+
+ private def computeStatsDiff(
+ newMetrics: Map[String, OperatorMetrics]
+ ): Map[String, OperatorMetrics] = {
+ // Default metrics for new operators
+ val defaultMetrics = OperatorMetrics(
+ WorkflowAggregatedState.UNINITIALIZED,
+ OperatorStatistics(Seq.empty, Seq.empty, 0, 0, 0, 0)
+ )
+
+ // Determine new and old keys
+ val newKeys = newMetrics.keySet.diff(lastPersistedMetrics.keySet)
+ val oldKeys = lastPersistedMetrics.keySet.diff(newMetrics.keySet)
+
+ // Update last metrics with default metrics for new keys
+ val updatedLastMetrics = lastPersistedMetrics ++ newKeys.map(_ -> defaultMetrics)
+
+ // Combine new metrics with old metrics for keys that are no longer present
+ val completeMetricsMap = newMetrics ++ oldKeys.map(key => key -> updatedLastMetrics(key))
+
+ // Transform the complete metrics map to ensure consistent structure
+ completeMetricsMap.map {
+ case (key, metrics) =>
+ key -> OperatorMetrics(
+ metrics.operatorState,
+ OperatorStatistics(
+ metrics.operatorStatistics.inputMetrics,
+ metrics.operatorStatistics.outputMetrics,
+ metrics.operatorStatistics.numWorkers,
+ metrics.operatorStatistics.dataProcessingTime,
+ metrics.operatorStatistics.controlProcessingTime,
+ metrics.operatorStatistics.idleTime
+ )
+ )
+ }
+ }
+
+ private def storeRuntimeStatistics(
+ operatorStatistics: scala.collection.immutable.Map[String, OperatorMetrics]
+ ): Unit = {
+ try {
+ operatorStatistics.foreach {
+ case (operatorId, stat) =>
+ val runtimeStats = new Tuple(
+ ResultSchema.runtimeStatisticsSchema,
+ Array(
+ operatorId,
+ new java.sql.Timestamp(System.currentTimeMillis()),
+ stat.operatorStatistics.inputMetrics.map(_.tupleMetrics.count).sum,
+ stat.operatorStatistics.inputMetrics.map(_.tupleMetrics.size).sum,
+ stat.operatorStatistics.outputMetrics.map(_.tupleMetrics.count).sum,
+ stat.operatorStatistics.outputMetrics.map(_.tupleMetrics.size).sum,
+ stat.operatorStatistics.dataProcessingTime,
+ stat.operatorStatistics.controlProcessingTime,
+ stat.operatorStatistics.idleTime,
+ stat.operatorStatistics.numWorkers,
+ maptoStatusCode(stat.operatorState).toInt
+ )
+ )
+ runtimeStatsWriter.putOne(runtimeStats)
+ }
+ } catch {
+ case err: Throwable => logger.error("error occurred when storing runtime statistics", err)
+ }
+ }
+
+ private[this] def registerCallbackOnWorkerAssignedUpdate(): Unit = {
+ addSubscription(
+ client
+ .registerCallback[WorkerAssignmentUpdate]((evt: WorkerAssignmentUpdate) => {
+ stateStore.statsStore.updateState { statsStore =>
+ statsStore.withOperatorWorkerMapping(
+ evt.workerMapping
+ .map({
+ case (opId, workerIds) => OperatorWorkerMapping(opId, workerIds.toSeq)
+ })
+ .toSeq
+ )
+ }
+ })
+ )
+ }
+
+ private[this] def registerCallbackOnWorkflowRecoveryUpdate(): Unit = {
+ addSubscription(
+ client
+ .registerCallback[WorkflowRecoveryStatus]((evt: WorkflowRecoveryStatus) => {
+ stateStore.metadataStore.updateState { metadataStore =>
+ metadataStore.withIsRecovering(evt.isRecovering)
+ }
+ })
+ )
+ }
+
+ private[this] def registerCallbackOnFatalError(): Unit = {
+ addSubscription(
+ client
+ .registerCallback[FatalError]((evt: FatalError) => {
+ client.shutdown()
+ val (operatorId, workerId) = getOperatorFromActorIdOpt(evt.fromActor)
+ stateStore.statsStore.updateState(stats =>
+ stats.withEndTimeStamp(System.currentTimeMillis())
+ )
+ stateStore.metadataStore.updateState { metadataStore =>
+ logger.error("error occurred in execution", evt.e)
+ updateWorkflowState(FAILED, metadataStore).addFatalErrors(
+ WorkflowFatalError(
+ EXECUTION_FAILURE,
+ Timestamp(Instant.now),
+ evt.e.toString,
+ getStackTraceWithAllCauses(evt.e),
+ operatorId,
+ workerId
+ )
+ )
+ }
+ })
+ )
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/ExecutionsMetadataPersistService.scala b/amber/src/main/scala/org/apache/texera/web/service/ExecutionsMetadataPersistService.scala
new file mode 100644
index 00000000000..833c4433328
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/ExecutionsMetadataPersistService.scala
@@ -0,0 +1,100 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity}
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.tables.daos.WorkflowExecutionsDao
+import org.apache.texera.dao.jooq.generated.tables.pojos.WorkflowExecutions
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowVersionResource._
+
+import java.sql.Timestamp
+
+/**
+ * This global object handles inserting a new entry to the DB to store metadata information about every workflow execution
+ * It also updates the entry if an execution status is updated
+ */
+object ExecutionsMetadataPersistService extends LazyLogging {
+ private def context =
+ SqlServer
+ .getInstance()
+ .createDSLContext()
+ private def workflowExecutionsDao =
+ new WorkflowExecutionsDao(
+ context.configuration
+ )
+
+ /**
+ * This method inserts a new entry of a workflow execution in the database and returns the generated eId
+ *
+ * @param workflowId the given workflow
+ * @param uid user id that initiated the execution
+ * @return generated execution ID
+ */
+
+ def insertNewExecution(
+ workflowId: WorkflowIdentity,
+ uid: Option[Integer],
+ executionName: String,
+ environmentVersion: String,
+ computingUnitId: Integer
+ ): ExecutionIdentity = {
+ // first retrieve the latest version of this workflow
+ val vid = getLatestVersion(workflowId.id.toInt)
+ val newExecution = new WorkflowExecutions()
+ if (executionName != "") {
+ newExecution.setName(executionName)
+ }
+ newExecution.setVid(vid)
+ newExecution.setUid(uid.orNull)
+ newExecution.setStartingTime(new Timestamp(System.currentTimeMillis()))
+ newExecution.setEnvironmentVersion(environmentVersion)
+
+ // Set computing unit ID if provided
+ newExecution.setCuid(computingUnitId)
+
+ workflowExecutionsDao.insert(newExecution)
+ ExecutionIdentity(newExecution.getEid.longValue())
+ }
+
+ def tryGetExistingExecution(executionId: ExecutionIdentity): Option[WorkflowExecutions] = {
+ try {
+ Some(workflowExecutionsDao.fetchOneByEid(executionId.id.toInt))
+ } catch {
+ case t: Throwable =>
+ logger.info("Unable to get execution. Error = " + t.getMessage)
+ None
+ }
+ }
+
+ def tryUpdateExistingExecution(
+ executionId: ExecutionIdentity
+ )(updateFunc: WorkflowExecutions => Unit): Unit = {
+ try {
+ val execution = workflowExecutionsDao.fetchOneByEid(executionId.id.toInt)
+ updateFunc(execution)
+ workflowExecutionsDao.update(execution)
+ } catch {
+ case t: Throwable =>
+ logger.info("Unable to update execution. Error = " + t.getMessage)
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/ResultExportService.scala b/amber/src/main/scala/org/apache/texera/web/service/ResultExportService.scala
new file mode 100644
index 00000000000..e4fdc92da94
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/ResultExportService.scala
@@ -0,0 +1,624 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.fasterxml.jackson.core.`type`.TypeReference
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import com.github.tototoshi.csv.CSVWriter
+import org.apache.texera.amber.config.EnvironmentalVariable
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.storage.model.VirtualDocument
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.{OperatorIdentity, WorkflowIdentity}
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.util.ArrowUtils
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector._
+import org.apache.arrow.vector.ipc.ArrowFileWriter
+import org.apache.commons.io.IOUtils
+import org.apache.commons.lang3.StringUtils
+import org.apache.texera.auth.JwtAuth
+import org.apache.texera.auth.JwtAuth.{TOKEN_EXPIRE_TIME_IN_MINUTES, jwtClaims}
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.apache.texera.web.model.http.request.result.{OperatorExportInfo, ResultExportRequest}
+import org.apache.texera.web.model.http.response.result.ResultExportResponse
+import org.apache.texera.web.resource.dashboard.user.workflow.{
+ WorkflowExecutionsResource,
+ WorkflowVersionResource
+}
+import org.apache.texera.web.service.WorkflowExecutionService.getLatestExecutionId
+
+import java.io.{FilterOutputStream, IOException, OutputStream}
+import java.net.{HttpURLConnection, URL, URLEncoder}
+import java.nio.channels.Channels
+import java.nio.charset.StandardCharsets
+import java.time.LocalDateTime
+import java.time.format.DateTimeFormatter
+import java.time.temporal.ChronoUnit
+import java.util.zip.{ZipEntry, ZipOutputStream}
+import javax.ws.rs.WebApplicationException
+import javax.ws.rs.core.{MediaType, Response, StreamingOutput}
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+import scala.util.Using
+
+object Constants {
+ val CHUNK_SIZE = 10
+}
+
+/**
+ * A simple wrapper that ignores 'close()' calls on the underlying stream.
+ * This allows each operator's writer to call close() without ending the entire ZipOutputStream.
+ */
+private class NonClosingOutputStream(os: OutputStream) extends FilterOutputStream(os) {
+ @throws[IOException]
+ override def close(): Unit = {
+ // do not actually close the underlying stream
+ super.flush()
+ // omit super.close()
+ }
+}
+
+object ResultExportService {
+ lazy val fileServiceUploadOneFileToDatasetEndpoint: String =
+ sys.env
+ .getOrElse(
+ EnvironmentalVariable.ENV_FILE_SERVICE_UPLOAD_ONE_FILE_TO_DATASET_ENDPOINT,
+ "http://localhost:9092/api/dataset/did/upload"
+ )
+ .trim
+}
+
+class ResultExportService(workflowIdentity: WorkflowIdentity, computingUnitId: Int) {
+
+ import ResultExportService._
+
+ /**
+ * Export operator results to a dataset and return the result.
+ */
+ def exportToDataset(
+ user: User,
+ request: ResultExportRequest
+ ): Response = {
+ val successMessages = new mutable.ListBuffer[String]()
+ val errorMessages = new mutable.ListBuffer[String]()
+
+ request.operators.foreach { op =>
+ try {
+ val (msgOpt, errOpt) = exportSingleOperatorToDataset(user, request, op)
+ msgOpt.foreach(successMessages += _)
+ errOpt.foreach(errorMessages += _)
+ } catch {
+ case ex: Exception =>
+ errorMessages += s"Error exporting operator $op: ${ex.getMessage}"
+ }
+ }
+
+ var exportResponse: ResultExportResponse = null
+ if (errorMessages.isEmpty) {
+ exportResponse = ResultExportResponse("success", successMessages.mkString("\n"))
+ } else if (successMessages.isEmpty) {
+ exportResponse = ResultExportResponse("error", errorMessages.mkString("\n"))
+ } else {
+ // At least one success, so we consider overall success (with partial possible).
+ exportResponse = ResultExportResponse("success", successMessages.mkString("\n"))
+ }
+
+ Response.ok(exportResponse).build()
+ }
+
+ /**
+ * Export operator results as downloadable files.
+ * If multiple operators are selected, their results are streamed as a ZIP file.
+ * If a single operator is selected, its result is streamed directly.
+ */
+ def exportToLocal(request: ResultExportRequest): Response = {
+ if (request.operators.size > 1) {
+ val (zipStream, zipFileNameOpt) = exportOperatorsAsZip(request)
+ if (zipStream == null) {
+ throw new RuntimeException("Zip stream is null")
+ }
+ val fileName = zipFileNameOpt.getOrElse("operators.zip")
+
+ Response
+ .ok(zipStream, "application/zip")
+ .header("Content-Disposition", s"""attachment; filename="$fileName"""")
+ .build()
+
+ } else {
+ val op = request.operators.head
+ val (streamingOutput, fileNameOpt) = exportOperatorResultAsStream(request, op)
+ if (streamingOutput == null) {
+ throw new RuntimeException("Failed to export operator")
+ }
+ val fileName = fileNameOpt.getOrElse("download.dat")
+
+ Response
+ .ok(streamingOutput, MediaType.APPLICATION_OCTET_STREAM)
+ .header("Content-Disposition", s"""attachment; filename="$fileName"""")
+ .build()
+ }
+ }
+
+ /**
+ * Export a single operator's result and handle different export types.
+ */
+ private def exportSingleOperatorToDataset(
+ user: User,
+ request: ResultExportRequest,
+ operatorRequest: OperatorExportInfo
+ ): (Option[String], Option[String]) = {
+
+ val execIdOpt = getLatestExecutionId(workflowIdentity, computingUnitId)
+ if (execIdOpt.isEmpty) {
+ return (None, Some(s"Workflow ${request.workflowId} has no execution result"))
+ }
+ val operatorDocument = getOperatorDocument(operatorRequest.id, computingUnitId)
+ if (operatorDocument == null || operatorDocument.getCount == 0)
+ return (None, Some(s"No results to export for operator $operatorRequest"))
+
+ val attributeNames =
+ operatorDocument.getRange(0, 1).to(Iterable).head.getSchema.getAttributeNames // small cost
+
+ val writer: OutputStream => Unit = operatorRequest.outputType match {
+ case "csv" => out => streamDocumentAsCSV(operatorDocument, out, Some(attributeNames))
+ case "arrow" => out => streamDocumentAsArrow(operatorDocument, out)
+ case "html" => out => streamDocumentAsHTML(out, operatorDocument)
+ case "data" => out => streamCellData(out, request, operatorDocument)
+ case "parquet" => out => streamDocumentAsParquetZip(operatorDocument, out)
+ case _ => out => streamDocumentAsCSV(operatorDocument, out, Some(attributeNames))
+ }
+
+ saveStreamToDataset(
+ operatorId = operatorRequest.id,
+ user = user,
+ request = request,
+ extension = operatorRequest.outputType,
+ writer = writer
+ )
+ }
+
+ /**
+ * Export a single operator's results as a streaming response (e.g., for download).
+ */
+ def exportOperatorResultAsStream(
+ request: ResultExportRequest,
+ operatorRequest: OperatorExportInfo
+ ): (StreamingOutput, Option[String]) = {
+ val execIdOpt = getLatestExecutionId(workflowIdentity, computingUnitId)
+ if (execIdOpt.isEmpty) {
+ return (null, None)
+ }
+
+ val operatorDocument = getOperatorDocument(operatorRequest.id, computingUnitId)
+ if (operatorDocument == null || operatorDocument.getCount == 0) {
+ return (null, None)
+ }
+
+ val fileName =
+ if (request.filename.isEmpty)
+ generateFileName(
+ request,
+ operatorRequest.id,
+ operatorRequest.outputType
+ )
+ else request.filename
+
+ val streamingOutput: StreamingOutput = (out: OutputStream) => {
+ operatorRequest.outputType match {
+ case "csv" => streamDocumentAsCSV(operatorDocument, out, None)
+ case "arrow" => streamDocumentAsArrow(operatorDocument, out)
+ case "data" => streamCellData(out, request, operatorDocument)
+ case "html" => streamDocumentAsHTML(out, operatorDocument)
+ case "parquet" => streamDocumentAsParquetZip(operatorDocument, out)
+ case _ => streamDocumentAsCSV(operatorDocument, out, None)
+ }
+ }
+
+ (streamingOutput, Some(fileName))
+ }
+
+ /**
+ * Export multiple operators' results as a single ZIP file stream.
+ */
+ def exportOperatorsAsZip(
+ request: ResultExportRequest
+ ): (StreamingOutput, Option[String]) = {
+ val timestamp = LocalDateTime
+ .now()
+ .truncatedTo(ChronoUnit.SECONDS)
+ .format(DateTimeFormatter.ofPattern("yyyy-MM-dd_HH-mm-ss"))
+ val zipFileName = s"${request.workflowName}-$timestamp.zip"
+
+ val execIdOpt = getLatestExecutionId(workflowIdentity, computingUnitId)
+ if (execIdOpt.isEmpty) {
+ throw new WebApplicationException(
+ s"No execution result for workflow ${request.workflowId}"
+ )
+ }
+
+ val streamingOutput: StreamingOutput = new StreamingOutput {
+ override def write(outputStream: OutputStream): Unit = {
+ Using.resource(new ZipOutputStream(outputStream)) { zipOut =>
+ request.operators.foreach { op =>
+ val operatorDocument = getOperatorDocument(op.id, computingUnitId)
+ if (operatorDocument == null || operatorDocument.getCount == 0) {
+ // create an "empty" file for this operator
+ zipOut.putNextEntry(new ZipEntry(s"${op.id}-empty.txt"))
+ val msg = s"Operator ${op.id} has no results"
+ zipOut.write(msg.getBytes(StandardCharsets.UTF_8))
+ zipOut.closeEntry()
+ } else {
+ val operatorFileName = generateFileName(request, op.id, op.outputType)
+
+ zipOut.putNextEntry(new ZipEntry(operatorFileName))
+ val nonClosingStream = new NonClosingOutputStream(zipOut)
+
+ op.outputType match {
+ case "csv" => streamDocumentAsCSV(operatorDocument, nonClosingStream, None)
+ case "arrow" => streamDocumentAsArrow(operatorDocument, nonClosingStream)
+ case "data" => streamCellData(nonClosingStream, request, operatorDocument)
+ case "html" => streamDocumentAsHTML(nonClosingStream, operatorDocument)
+ case "parquet" => streamDocumentAsParquetZip(operatorDocument, nonClosingStream)
+ case _ => streamDocumentAsCSV(operatorDocument, nonClosingStream, None)
+ }
+ zipOut.closeEntry()
+ }
+ }
+ }
+ }
+ }
+
+ (streamingOutput, Some(zipFileName))
+ }
+
+ /**
+ * Streams the entire content of `VirtualDocument` as CSV into `outputStream` in a single pass.
+ */
+ private def streamDocumentAsCSV(
+ doc: VirtualDocument[Tuple],
+ outputStream: OutputStream,
+ maybeHeaders: Option[List[String]]
+ ): Unit = {
+ val totalCount = doc.getCount
+ if (totalCount == 0) {
+ return
+ }
+
+ val iterator = doc.get()
+ if (!iterator.hasNext) {
+ return
+ }
+
+ val csvWriter = CSVWriter.open(outputStream)
+
+ val headers: List[String] = maybeHeaders match {
+ case Some(hdrs) =>
+ hdrs
+ case None =>
+ val firstRow = iterator.next()
+ val inferredHeaders = firstRow.getSchema.getAttributeNames
+
+ csvWriter.writeRow(inferredHeaders)
+ csvWriter.writeRow(firstRow.getFields.toIndexedSeq)
+
+ inferredHeaders
+ }
+
+ if (maybeHeaders.isDefined) {
+ csvWriter.writeRow(headers)
+ }
+
+ val buffer = new ArrayBuffer[Tuple](Constants.CHUNK_SIZE)
+
+ while (iterator.hasNext) {
+ buffer.clear()
+ var count = 0
+
+ while (count < Constants.CHUNK_SIZE && iterator.hasNext) {
+ buffer += iterator.next()
+ count += 1
+ }
+ buffer.foreach { t =>
+ csvWriter.writeRow(t.getFields.toIndexedSeq)
+ }
+ csvWriter.flush()
+ }
+
+ csvWriter.close()
+ }
+
+ /**
+ * Streams the entire content of `VirtualDocument` as Arrow into `outputStream` in a single pass.
+ */
+ private def streamDocumentAsArrow(
+ doc: VirtualDocument[Tuple],
+ outputStream: OutputStream
+ ): Unit = {
+ if (doc.getCount == 0) return
+
+ val allocator = new RootAllocator()
+ Using.Manager { use =>
+ val firstTuple = doc.getRange(0, 1).to(Iterable).head
+ val schema = firstTuple.getSchema
+ val arrowSchema = ArrowUtils.fromTexeraSchema(schema)
+
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ use(root)
+
+ val channel = Channels.newChannel(outputStream)
+ val writer = new ArrowFileWriter(root, null, channel)
+ use(writer)
+ use(allocator)
+
+ writer.start()
+
+ val iterator = doc.get()
+ val buffer = new ArrayBuffer[Tuple](Constants.CHUNK_SIZE)
+
+ while (iterator.hasNext) {
+ buffer.clear()
+ var count = 0
+
+ while (count < Constants.CHUNK_SIZE && iterator.hasNext) {
+ buffer += iterator.next()
+ count += 1
+ }
+
+ if (buffer.nonEmpty) {
+ val currentBatchSize = buffer.size
+
+ for (i <- 0 until currentBatchSize) {
+ val tuple = buffer(i)
+ ArrowUtils.setTexeraTuple(tuple, i, root)
+ }
+
+ root.setRowCount(currentBatchSize)
+ writer.writeBatch()
+
+ root.clear()
+ }
+ }
+
+ writer.end()
+ }
+ }
+
+ /*
+ * Handle streaming HTML result from a visualization operator's result.
+ */
+ private def streamDocumentAsHTML(
+ out: OutputStream,
+ operatorDocument: VirtualDocument[Tuple]
+ ): Unit = {
+ val results: Iterable[Tuple] = operatorDocument.get().to(Iterable)
+ val resHead = results.head
+ val htmlCode = resHead.getField(0).toString
+ out.write(htmlCode.getBytes(StandardCharsets.UTF_8))
+ out.flush()
+ }
+
+ /**
+ * Streams the underlying Parquet files of an Iceberg document into a ZIP archive.
+ * This avoids re-encoding and uses minimal memory and no temporary disk space.
+ */
+ private def streamDocumentAsParquetZip(
+ doc: VirtualDocument[Tuple],
+ outputStream: OutputStream
+ ): Unit = {
+ try {
+ val zipStream = doc.asInputStream()
+ try {
+ IOUtils.copy(zipStream, outputStream)
+ } finally {
+ zipStream.close()
+ }
+ } catch {
+ case e: Exception =>
+ throw e
+ }
+ }
+
+ /*
+ * Handle streaming a single (row, column) from an operator's result.
+ * This is used for the "data" export type, which exports a single field value.
+ */
+ private def streamCellData(
+ out: OutputStream,
+ request: ResultExportRequest,
+ operatorDocument: VirtualDocument[Tuple]
+ ): Unit = {
+ val rowIndex = request.rowIndex
+ val columnIndex = request.columnIndex
+
+ if (rowIndex >= operatorDocument.getCount) {
+ throw new WebApplicationException(
+ s"Invalid rowIndex ($rowIndex). Total rows: ${operatorDocument.getCount}"
+ )
+ }
+
+ val selectedRow = operatorDocument
+ .getRange(rowIndex, rowIndex + 1)
+ .to(Iterable)
+ .headOption
+ .getOrElse(throw new RuntimeException(s"Could not retrieve row at index $rowIndex"))
+
+ if (columnIndex >= selectedRow.getFields.length) {
+ throw new WebApplicationException(
+ s"Invalid columnIndex ($columnIndex). Total columns: ${selectedRow.getFields.length}"
+ )
+ }
+
+ val field: Any = selectedRow.getField(columnIndex)
+ val dataBytes = convertFieldToBytes(field)
+ out.write(dataBytes)
+ }
+
+ /**
+ * Generate the VirtualDocument for one operator's result.
+ * Incorporates the remote code's extra parameter `None` for sub-operator ID.
+ */
+ private def getOperatorDocument(
+ operatorId: String,
+ computingUnitId: Int
+ ): VirtualDocument[Tuple] = {
+ // By now the workflow should finish running
+ // Only supports external port 0 for now. TODO: support multiple ports
+ val storageUri = WorkflowExecutionsResource.getResultUriByLogicalPortId(
+ getLatestExecutionId(workflowIdentity, computingUnitId).get,
+ OperatorIdentity(operatorId),
+ PortIdentity()
+ )
+
+ storageUri
+ .map(uri => DocumentFactory.openDocument(uri)._1.asInstanceOf[VirtualDocument[Tuple]])
+ .orNull
+ }
+
+ private def saveStreamToDataset(
+ operatorId: String,
+ user: User,
+ request: ResultExportRequest,
+ extension: String,
+ writer: OutputStream => Unit
+ ): (Option[String], Option[String]) = {
+ val fileName =
+ if (request.filename.isEmpty) generateFileName(request, operatorId, extension)
+ else request.filename
+
+ try {
+ saveToDatasets(request, user, writer, fileName)
+ (Some(s"$extension export done for operator $operatorId -> file: $fileName"), None)
+ } catch {
+ case ex: Exception =>
+ (None, Some(s"$extension export failed for operator $operatorId: ${ex.getMessage}"))
+ }
+ }
+
+ private def convertFieldToBytes(field: Any): Array[Byte] = {
+ field match {
+ case data: Array[Byte] => data
+ case data: String => data.getBytes(StandardCharsets.UTF_8)
+ case other => other.toString.getBytes(StandardCharsets.UTF_8)
+ }
+ }
+
+ /**
+ * Save the pipedInputStream into the specified datasets as a new dataset version.
+ */
+ private def saveToDatasets(
+ request: ResultExportRequest,
+ user: User,
+ fileWriter: OutputStream => Unit,
+ fileName: String
+ ): Unit = {
+ request.datasetIds.foreach { did =>
+ val encodedFilePath = URLEncoder.encode(fileName, StandardCharsets.UTF_8.name())
+ val message = URLEncoder.encode(
+ s"Export from workflow ${request.workflowName}",
+ StandardCharsets.UTF_8.name()
+ )
+
+ val uploadUrl = s"$fileServiceUploadOneFileToDatasetEndpoint"
+ .replace("did", did.toString) + s"?filePath=$encodedFilePath&message=$message"
+
+ var connection: HttpURLConnection = null
+ try {
+ val url = new URL(uploadUrl)
+ connection = url.openConnection().asInstanceOf[HttpURLConnection]
+ connection.setDoOutput(true)
+ connection.setRequestMethod("POST")
+ connection.setRequestProperty("Content-Type", "application/octet-stream")
+ connection.setRequestProperty(
+ "Authorization",
+ s"Bearer ${JwtAuth.jwtToken(jwtClaims(user, TOKEN_EXPIRE_TIME_IN_MINUTES))}"
+ )
+ connection.setChunkedStreamingMode(0)
+
+ val outputStream = connection.getOutputStream
+ fileWriter(outputStream)
+ outputStream.close()
+
+ val responseCode = connection.getResponseCode
+ if (responseCode != HttpURLConnection.HTTP_OK) {
+ throw new RuntimeException(s"Failed to upload file. Server responded with: $responseCode")
+ }
+ } catch {
+ case e: Exception =>
+ throw new RuntimeException(s"Error uploading file to dataset $did: ${e.getMessage}", e)
+ } finally {
+ if (connection != null) connection.disconnect()
+ }
+ }
+ }
+
+ /**
+ * Generate a file name for an operator's exported file
+ */
+ private def generateFileName(
+ request: ResultExportRequest,
+ operatorId: String,
+ extension: String
+ ): String = {
+ val extensionMatch = extension match {
+ case "parquet" => "zip"
+ case _ => extension
+ }
+
+ val latestVersion =
+ WorkflowVersionResource.getLatestVersion(request.workflowId)
+ val timestamp = LocalDateTime
+ .now()
+ .truncatedTo(ChronoUnit.SECONDS)
+ .format(DateTimeFormatter.ofPattern("yyyy-MM-dd_HH-mm-ss"))
+
+ val rawName =
+ s"${request.workflowName}-op$operatorId-v$latestVersion-$timestamp.$extensionMatch"
+ // remove path separators
+ StringUtils.replaceEach(rawName, Array("/", "\\"), Array("", ""))
+ }
+
+ /**
+ * Parse a JSON string array of operators into a list of OperatorExportInfo objects.
+ */
+ def parseOperators(operatorsJson: String): List[OperatorExportInfo] = {
+ new ObjectMapper()
+ .registerModule(DefaultScalaModule)
+ .readValue(operatorsJson, new TypeReference[List[OperatorExportInfo]] {})
+ }
+
+ /**
+ * Validate an export request by checking if any operators are selected.
+ * Return an error response if none are selected, otherwise None.
+ */
+ def validateExportRequest(request: ResultExportRequest): Option[Response] = {
+ if (request.operators.isEmpty) {
+ Some(
+ Response
+ .status(Response.Status.BAD_REQUEST)
+ .`type`(MediaType.APPLICATION_JSON)
+ .entity(Map("error" -> "No operator selected").asJava)
+ .build()
+ )
+ } else None
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/WorkflowEmailNotifier.scala b/amber/src/main/scala/org/apache/texera/web/service/WorkflowEmailNotifier.scala
new file mode 100644
index 00000000000..af9a26286d4
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/WorkflowEmailNotifier.scala
@@ -0,0 +1,117 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState._
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource
+import org.apache.texera.web.resource.{EmailMessage, GmailResource}
+import org.hibernate.validator.internal.constraintvalidators.hv.EmailValidator
+
+import java.net.URI
+import java.time.format.DateTimeFormatter
+import java.time.{Instant, ZoneOffset}
+
+class WorkflowEmailNotifier(
+ workflowId: Long,
+ userEmail: String,
+ sessionUri: URI
+) extends EmailNotifier
+ with LazyLogging {
+ private val workflowName = WorkflowResource.getWorkflowName(workflowId.toInt)
+ private val emailValidator = new EmailValidator()
+
+ private val TerminalStates: Set[WorkflowAggregatedState] = Set(
+ COMPLETED,
+ PAUSED,
+ FAILED,
+ KILLED
+ )
+
+ override def shouldSendEmail(workflowState: WorkflowAggregatedState): Boolean =
+ TerminalStates.contains(workflowState)
+
+ override def sendStatusEmail(state: WorkflowAggregatedState): Unit = {
+ if (!isValidEmail(userEmail)) {
+ logger.warn(s"Invalid email address: $userEmail")
+ return
+ }
+
+ val emailMessage = createEmailMessage(state)
+
+ try {
+ GmailResource.sendEmail(emailMessage, userEmail)
+ } catch {
+ case e: Exception => println(s"Failed to send email: ${e.getMessage}")
+ }
+ }
+
+ private def isValidEmail(email: String): Boolean = emailValidator.isValid(email, null)
+
+ private def createEmailMessage(state: WorkflowAggregatedState): EmailMessage = {
+ EmailMessage(
+ receiver = userEmail,
+ subject = createEmailSubject(state),
+ content = createEmailContent(state)
+ )
+ }
+
+ private def createEmailSubject(state: WorkflowAggregatedState): String =
+ s"[Texera] Workflow $workflowName ($workflowId) Status: $state"
+
+ private def createEmailContent(state: WorkflowAggregatedState): String = {
+ val timestamp = formatTimestamp(Instant.now())
+ val dashboardUrl = createDashboardUrl()
+
+ s"""
+ |Hello,
+ |
+ |The workflow with the following details has changed its state:
+ |
+ |- Workflow ID: $workflowId
+ |- Workflow Name: $workflowName
+ |- State: $state
+ |- Timestamp: $timestamp
+ |
+ |You can view more details by visiting: $dashboardUrl
+ |
+ |Regards,
+ |Texera Team
+ """.stripMargin.trim
+ }
+
+ private def formatTimestamp(instant: Instant): String =
+ DateTimeFormatter
+ .ofPattern("MMMM d, yyyy, h:mm:ss a '(UTC)'")
+ .withZone(ZoneOffset.UTC)
+ .format(instant)
+
+ private def createDashboardUrl(): String = {
+ val host = sessionUri.getHost
+ val port = sessionUri.getPort
+ val path = s"/dashboard/user/workspace/$workflowId"
+ if (port == -1 || port == 80 || port == 443) {
+ s"http://$host$path"
+ } else {
+ s"http://$host:$port$path"
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/WorkflowExecutionService.scala b/amber/src/main/scala/org/apache/texera/web/service/WorkflowExecutionService.scala
new file mode 100644
index 00000000000..741687e02c9
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/WorkflowExecutionService.scala
@@ -0,0 +1,181 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity}
+import org.apache.texera.amber.core.workflow.WorkflowContext
+import org.apache.texera.amber.engine.architecture.controller.{ControllerConfig, Workflow}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmptyRequest
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState._
+import org.apache.texera.amber.engine.common.Utils
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.engine.common.executionruntimestate.ExecutionMetadataStore
+import org.apache.texera.web.model.websocket.event.{
+ TexeraWebSocketEvent,
+ WorkflowErrorEvent,
+ WorkflowStateEvent
+}
+import org.apache.texera.web.model.websocket.request.WorkflowExecuteRequest
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+import org.apache.texera.web.storage.ExecutionStateStore
+import org.apache.texera.web.storage.ExecutionStateStore.updateWorkflowState
+import org.apache.texera.web.{ComputingUnitMaster, SubscriptionManager, WebsocketInput}
+import org.apache.texera.workflow.WorkflowCompiler
+
+import java.net.URI
+import scala.collection.mutable
+
+object WorkflowExecutionService {
+ def getLatestExecutionId(
+ workflowId: WorkflowIdentity,
+ computingUnitId: Int
+ ): Option[ExecutionIdentity] = {
+ WorkflowExecutionsResource
+ .getLatestExecutionID(workflowId.id.toInt, computingUnitId)
+ .map(eid => new ExecutionIdentity(eid.longValue()))
+ }
+}
+
+class WorkflowExecutionService(
+ controllerConfig: ControllerConfig,
+ val workflowContext: WorkflowContext,
+ resultService: ExecutionResultService,
+ request: WorkflowExecuteRequest,
+ val executionStateStore: ExecutionStateStore,
+ errorHandler: Throwable => Unit,
+ userEmailOpt: Option[String],
+ sessionUri: URI
+) extends SubscriptionManager
+ with LazyLogging {
+
+ workflowContext.workflowSettings = request.workflowSettings
+ val wsInput = new WebsocketInput(errorHandler)
+
+ addSubscription(
+ executionStateStore.metadataStore.registerDiffHandler((oldState, newState) => {
+ val outputEvents = new mutable.ArrayBuffer[TexeraWebSocketEvent]()
+
+ if (newState.state != oldState.state || newState.isRecovering != oldState.isRecovering) {
+ outputEvents.append(createStateEvent(newState))
+ }
+
+ if (newState.fatalErrors != oldState.fatalErrors) {
+ outputEvents.append(WorkflowErrorEvent(newState.fatalErrors))
+ }
+
+ outputEvents
+ })
+ )
+
+ private def createStateEvent(state: ExecutionMetadataStore): WorkflowStateEvent = {
+ if (state.isRecovering && state.state != COMPLETED) {
+ WorkflowStateEvent("Recovering")
+ } else {
+ WorkflowStateEvent(Utils.aggregatedStateToString(state.state))
+ }
+ }
+
+ var workflow: Workflow = _
+
+ // Runtime starts from here:
+ logger.info("Initialing an AmberClient, runtime starting...")
+ var client: AmberClient = _
+ var executionReconfigurationService: ExecutionReconfigurationService = _
+ var executionStatsService: ExecutionStatsService = _
+ var executionRuntimeService: ExecutionRuntimeService = _
+ var executionConsoleService: ExecutionConsoleService = _
+
+ def executeWorkflow(): Unit = {
+ try {
+ workflow = new WorkflowCompiler(workflowContext)
+ .compile(request.logicalPlan)
+ } catch {
+ case err: Throwable =>
+ errorHandler(err)
+ }
+
+ client = ComputingUnitMaster.createAmberRuntime(
+ workflow.context,
+ workflow.physicalPlan,
+ controllerConfig,
+ errorHandler
+ )
+ executionReconfigurationService =
+ new ExecutionReconfigurationService(client, executionStateStore, workflow)
+ executionStatsService = new ExecutionStatsService(client, executionStateStore, workflow.context)
+ executionRuntimeService = new ExecutionRuntimeService(
+ client,
+ executionStateStore,
+ wsInput,
+ executionReconfigurationService,
+ controllerConfig.faultToleranceConfOpt,
+ workflowContext.workflowId.id,
+ request.emailNotificationEnabled,
+ userEmailOpt,
+ sessionUri
+ )
+ executionConsoleService =
+ new ExecutionConsoleService(client, executionStateStore, wsInput, workflow.context)
+
+ logger.info("Starting the workflow execution.")
+ resultService.attachToExecution(
+ workflow.context.executionId,
+ executionStateStore,
+ workflow.physicalPlan,
+ client
+ )
+ executionStateStore.metadataStore.updateState(metadataStore =>
+ updateWorkflowState(READY, metadataStore)
+ .withFatalErrors(Seq.empty)
+ )
+ executionStateStore.statsStore.updateState(stats =>
+ stats.withStartTimeStamp(System.currentTimeMillis())
+ )
+ client.controllerInterface
+ .startWorkflow(EmptyRequest(), ())
+ .onFailure(err => {
+ errorHandler(err)
+ })
+ .onSuccess(resp =>
+ executionStateStore.metadataStore.updateState(metadataStore =>
+ if (metadataStore.state != FAILED) {
+ updateWorkflowState(resp.workflowState, metadataStore)
+ } else {
+ metadataStore
+ }
+ )
+ )
+ }
+
+ override def unsubscribeAll(): Unit = {
+ super.unsubscribeAll()
+ if (client != null) {
+ // runtime created
+ client.shutdown()
+ executionRuntimeService.unsubscribeAll()
+ executionConsoleService.unsubscribeAll()
+ executionStatsService.unsubscribeAll()
+ executionReconfigurationService.unsubscribeAll()
+ }
+
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/service/WorkflowService.scala b/amber/src/main/scala/org/apache/texera/web/service/WorkflowService.scala
new file mode 100644
index 00000000000..aa593cdcc65
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/service/WorkflowService.scala
@@ -0,0 +1,354 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.google.protobuf.timestamp.Timestamp
+import com.typesafe.scalalogging.LazyLogging
+import io.reactivex.rxjava3.disposables.{CompositeDisposable, Disposable}
+import io.reactivex.rxjava3.subjects.BehaviorSubject
+import org.apache.texera.amber.config.ApplicationConfig
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.storage.result.iceberg.OnIceberg
+import org.apache.texera.amber.core.virtualidentity.{
+ EmbeddedControlMessageIdentity,
+ ExecutionIdentity,
+ WorkflowIdentity
+}
+import org.apache.texera.amber.core.workflow.WorkflowContext
+import org.apache.texera.amber.core.workflowruntimestate.FatalErrorType.EXECUTION_FAILURE
+import org.apache.texera.amber.core.workflowruntimestate.WorkflowFatalError
+import org.apache.texera.amber.engine.architecture.controller.ControllerConfig
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{
+ COMPLETED,
+ FAILED
+}
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ FaultToleranceConfig,
+ StateRestoreConfig
+}
+import org.apache.texera.amber.error.ErrorUtils.{
+ getOperatorFromActorIdOpt,
+ getStackTraceWithAllCauses
+}
+import org.apache.texera.dao.jooq.generated.tables.pojos.User
+import org.apache.texera.service.util.LargeBinaryManager
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+import org.apache.texera.web.model.websocket.request.WorkflowExecuteRequest
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource
+import org.apache.texera.web.service.WorkflowService.mkWorkflowStateId
+import org.apache.texera.web.storage.ExecutionStateStore.updateWorkflowState
+import org.apache.texera.web.storage.{ExecutionStateStore, WorkflowStateStore}
+import org.apache.texera.web.{SubscriptionManager, WorkflowLifecycleManager}
+import org.apache.texera.workflow.LogicalPlan
+import play.api.libs.json.Json
+
+import java.net.URI
+import java.time.Instant
+import java.util.concurrent.ConcurrentHashMap
+import scala.jdk.CollectionConverters.IterableHasAsScala
+
+object WorkflowService {
+ private val workflowServiceMapping = new ConcurrentHashMap[String, WorkflowService]()
+ val cleanUpDeadlineInSeconds: Int = ApplicationConfig.executionStateCleanUpInSecs
+
+ def getAllWorkflowServices: Iterable[WorkflowService] = workflowServiceMapping.values().asScala
+
+ def mkWorkflowStateId(workflowId: WorkflowIdentity): String = {
+ workflowId.toString
+ }
+
+ def getOrCreate(
+ workflowId: WorkflowIdentity,
+ computingUnitId: Int,
+ cleanupTimeout: Int = cleanUpDeadlineInSeconds
+ ): WorkflowService = {
+ workflowServiceMapping.compute(
+ mkWorkflowStateId(workflowId),
+ (_, v) => {
+ if (v == null) {
+ new WorkflowService(workflowId, computingUnitId, cleanupTimeout)
+ } else {
+ v
+ }
+ }
+ )
+ }
+}
+
+class WorkflowService(
+ val workflowId: WorkflowIdentity,
+ val computingUnitId: Int,
+ cleanUpTimeout: Int
+) extends SubscriptionManager
+ with LazyLogging {
+
+ // state across execution:
+ private val errorSubject = BehaviorSubject.create[TexeraWebSocketEvent]().toSerialized
+ val stateStore = new WorkflowStateStore()
+ var executionService: BehaviorSubject[WorkflowExecutionService] = BehaviorSubject.create()
+
+ val resultService: ExecutionResultService =
+ new ExecutionResultService(workflowId, computingUnitId, stateStore)
+ val lifeCycleManager: WorkflowLifecycleManager = new WorkflowLifecycleManager(
+ s"workflowId=$workflowId",
+ cleanUpTimeout,
+ () => {
+ // clear the storage resources associated with the latest execution
+ WorkflowExecutionService
+ .getLatestExecutionId(workflowId, computingUnitId)
+ .foreach(eid => {
+ clearExecutionResources(eid)
+ })
+ WorkflowService.workflowServiceMapping.remove(mkWorkflowStateId(workflowId))
+ if (executionService.getValue != null) {
+ // shutdown client
+ executionService.getValue.client.shutdown()
+ }
+ unsubscribeAll()
+ }
+ )
+
+ var lastCompletedLogicalPlan: Option[LogicalPlan] = Option.empty
+
+ executionService.subscribe { executionService: WorkflowExecutionService =>
+ {
+ executionService.executionStateStore.metadataStore.registerDiffHandler {
+ (oldState, newState) =>
+ {
+ if (oldState.state != COMPLETED && newState.state == COMPLETED) {
+ lastCompletedLogicalPlan = Option.apply(executionService.workflow.logicalPlan)
+ }
+ Iterable.empty
+ }
+ }
+ }
+ }
+
+ def connect(onNext: TexeraWebSocketEvent => Unit): Disposable = {
+ lifeCycleManager.increaseUserCount()
+ val subscriptions = stateStore.getAllStores
+ .map(_.getWebsocketEventObservable)
+ .map(evtPub =>
+ evtPub.subscribe { evts: Iterable[TexeraWebSocketEvent] => evts.foreach(onNext) }
+ )
+ .toSeq
+ val errorSubscription = errorSubject.subscribe { evt: TexeraWebSocketEvent => onNext(evt) }
+ new CompositeDisposable(subscriptions :+ errorSubscription: _*)
+ }
+
+ def connectToExecution(onNext: TexeraWebSocketEvent => Unit): Disposable = {
+ val localDisposable = new CompositeDisposable()
+ val disposable = executionService.subscribe { execService: WorkflowExecutionService =>
+ localDisposable.clear() // Clears previous subscriptions safely
+ val subscriptions = execService.executionStateStore.getAllStores
+ .map(_.getWebsocketEventObservable)
+ .map(evtPub =>
+ evtPub.subscribe { events: Iterable[TexeraWebSocketEvent] => events.foreach(onNext) }
+ )
+ .toSeq
+ localDisposable.addAll(subscriptions: _*)
+ }
+ // Note: this new CompositeDisposable is necessary. DO NOT OPTIMIZE.
+ new CompositeDisposable(localDisposable, disposable)
+ }
+
+ def disconnect(): Unit = {
+ lifeCycleManager.decreaseUserCount(
+ Option(executionService.getValue).map(_.executionStateStore.metadataStore.getState.state)
+ )
+ }
+
+ private[this] def createWorkflowContext(): WorkflowContext = {
+ new WorkflowContext(workflowId)
+ }
+
+ def initExecutionService(
+ req: WorkflowExecuteRequest,
+ userOpt: Option[User],
+ sessionUri: URI
+ ): Unit = {
+
+ if (executionService.hasValue) {
+ executionService.getValue.unsubscribeAll()
+ }
+
+ val (uidOpt, userEmailOpt) = userOpt.map(user => (user.getUid, user.getEmail)).unzip
+
+ val workflowContext: WorkflowContext = createWorkflowContext()
+ var controllerConf = ControllerConfig.default
+
+ // clean up results from previous run
+ val previousExecutionId =
+ WorkflowExecutionService.getLatestExecutionId(workflowId, req.computingUnitId)
+ previousExecutionId.foreach(eid => {
+ clearExecutionResources(eid)
+ }) // TODO: change this behavior after enabling cache.
+
+ workflowContext.executionId = ExecutionsMetadataPersistService.insertNewExecution(
+ workflowContext.workflowId,
+ uidOpt,
+ req.executionName,
+ convertToJson(req.engineVersion),
+ req.computingUnitId
+ )
+
+ if (ApplicationConfig.faultToleranceLogRootFolder.isDefined) {
+ val writeLocation = ApplicationConfig.faultToleranceLogRootFolder.get.resolve(
+ s"${workflowContext.workflowId}/${workflowContext.executionId}/"
+ )
+ ExecutionsMetadataPersistService.tryUpdateExistingExecution(workflowContext.executionId) {
+ execution => execution.setLogLocation(writeLocation.toString)
+ }
+ controllerConf = controllerConf.copy(faultToleranceConfOpt =
+ Some(FaultToleranceConfig(writeTo = writeLocation))
+ )
+ }
+ if (req.replayFromExecution.isDefined) {
+ val replayInfo = req.replayFromExecution.get
+ ExecutionsMetadataPersistService
+ .tryGetExistingExecution(ExecutionIdentity(replayInfo.eid))
+ .foreach { execution =>
+ val readLocation = new URI(execution.getLogLocation)
+ controllerConf = controllerConf.copy(stateRestoreConfOpt =
+ Some(
+ StateRestoreConfig(
+ readFrom = readLocation,
+ replayDestination = EmbeddedControlMessageIdentity(replayInfo.interaction)
+ )
+ )
+ )
+ }
+ }
+
+ val executionStateStore = new ExecutionStateStore()
+ // assign execution id to find the execution from DB in case the constructor fails.
+ executionStateStore.metadataStore.updateState(state =>
+ state.withExecutionId(workflowContext.executionId)
+ )
+ val errorHandler: Throwable => Unit = { t =>
+ {
+ val fromActorOpt = t match {
+ case ex: WorkflowRuntimeException =>
+ ex.relatedWorkerId
+ case other =>
+ None
+ }
+ val (operatorId, workerId) = getOperatorFromActorIdOpt(fromActorOpt)
+ logger.error("error during execution", t)
+ executionStateStore.statsStore.updateState(stats =>
+ stats.withEndTimeStamp(System.currentTimeMillis())
+ )
+ executionStateStore.metadataStore.updateState { metadataStore =>
+ updateWorkflowState(FAILED, metadataStore).addFatalErrors(
+ WorkflowFatalError(
+ EXECUTION_FAILURE,
+ Timestamp(Instant.now),
+ t.toString,
+ getStackTraceWithAllCauses(t),
+ operatorId,
+ workerId
+ )
+ )
+ }
+ }
+ }
+ try {
+ val execution = new WorkflowExecutionService(
+ controllerConf,
+ workflowContext,
+ resultService,
+ req,
+ executionStateStore,
+ errorHandler,
+ userEmailOpt,
+ sessionUri
+ )
+ lifeCycleManager.registerCleanUpOnStateChange(executionStateStore)
+ executionService.onNext(execution)
+ execution.executeWorkflow()
+ } catch {
+ case e: Throwable => errorHandler(e)
+ }
+
+ }
+
+ def convertToJson(frontendVersion: String): String = {
+ val environmentVersionMap = Map(
+ "engine_version" -> Json.toJson(frontendVersion)
+ )
+ Json.stringify(Json.toJson(environmentVersionMap))
+ }
+
+ override def unsubscribeAll(): Unit = {
+ super.unsubscribeAll()
+ Option(executionService.getValue).foreach(_.unsubscribeAll())
+ resultService.unsubscribeAll()
+ }
+
+ /**
+ * Cleans up all resources associated with a workflow execution.
+ *
+ * This method performs resource cleanup in the following sequence:
+ * 1. Retrieves all document URIs associated with the execution
+ * 2. Clears URI references from the execution registry
+ * 3. Safely clears all result and console message documents
+ * 4. Expires Iceberg snapshots for runtime statistics
+ * 5. Deletes large binaries from MinIO
+ *
+ * @param eid The execution identity to clean up resources for
+ */
+ private def clearExecutionResources(eid: ExecutionIdentity): Unit = {
+ // Retrieve URIs for all resources associated with this execution
+ val resultUris = WorkflowExecutionsResource.getResultUrisByExecutionId(eid)
+ val consoleMessagesUris = WorkflowExecutionsResource.getConsoleMessagesUriByExecutionId(eid)
+
+ // Remove references from registry first
+ WorkflowExecutionsResource.deleteConsoleMessageAndExecutionResultUris(eid)
+
+ // Clean up all result and console message documents
+ (resultUris ++ consoleMessagesUris).foreach { uri =>
+ try DocumentFactory.openDocument(uri)._1.clear()
+ catch {
+ case error: Throwable =>
+ logger.debug(s"Error processing document at $uri: ${error.getMessage}")
+ }
+ }
+
+ // Expire any Iceberg snapshots for runtime statistics
+ WorkflowExecutionsResource.getRuntimeStatsUriByExecutionId(eid).foreach { uri =>
+ try {
+ DocumentFactory.openDocument(uri)._1 match {
+ case iceberg: OnIceberg => iceberg.expireSnapshots()
+ case other =>
+ logger.error(
+ s"Cannot expire snapshots: document from URI [$uri] is of type ${other.getClass.getName}. " +
+ s"Expected an instance of ${classOf[OnIceberg].getName}."
+ )
+ }
+ } catch {
+ case error: Throwable =>
+ logger.debug(s"Error processing document at $uri: ${error.getMessage}")
+ }
+ }
+ // Delete large binaries
+ LargeBinaryManager.deleteAllObjects()
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/storage/ExecutionReconfigurationStore.scala b/amber/src/main/scala/org/apache/texera/web/storage/ExecutionReconfigurationStore.scala
new file mode 100644
index 00000000000..946a189ae9f
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/storage/ExecutionReconfigurationStore.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.storage
+
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.core.workflow.PhysicalOp
+import org.apache.texera.amber.operator.StateTransferFunc
+
+case class ExecutionReconfigurationStore(
+ currentReconfigId: Option[String] = None,
+ unscheduledReconfigurations: List[(PhysicalOp, Option[StateTransferFunc])] = List(),
+ completedReconfigurations: Set[ActorVirtualIdentity] = Set()
+)
diff --git a/amber/src/main/scala/org/apache/texera/web/storage/ExecutionStateStore.scala b/amber/src/main/scala/org/apache/texera/web/storage/ExecutionStateStore.scala
new file mode 100644
index 00000000000..654acbbefd1
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/storage/ExecutionStateStore.scala
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.storage
+
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.common.Utils.maptoStatusCode
+import org.apache.texera.amber.engine.common.executionruntimestate.{
+ ExecutionBreakpointStore,
+ ExecutionConsoleStore,
+ ExecutionMetadataStore,
+ ExecutionStatsStore
+}
+import org.apache.texera.web.service.ExecutionsMetadataPersistService
+
+import java.sql.Timestamp
+
+object ExecutionStateStore {
+
+ // Update the state of the specified execution if user system is enabled.
+ // Update the execution only from backend
+ def updateWorkflowState(
+ state: WorkflowAggregatedState,
+ metadataStore: ExecutionMetadataStore
+ ): ExecutionMetadataStore = {
+ ExecutionsMetadataPersistService.tryUpdateExistingExecution(metadataStore.executionId) {
+ execution =>
+ execution.setStatus(maptoStatusCode(state))
+ execution.setLastUpdateTime(new Timestamp(System.currentTimeMillis()))
+ }
+ metadataStore.withState(state)
+ }
+}
+
+// states that within one execution.
+class ExecutionStateStore {
+ val statsStore = new StateStore(ExecutionStatsStore())
+ val metadataStore = new StateStore(ExecutionMetadataStore())
+ val consoleStore = new StateStore(ExecutionConsoleStore())
+ val breakpointStore = new StateStore(ExecutionBreakpointStore())
+ val reconfigurationStore = new StateStore(ExecutionReconfigurationStore())
+
+ def getAllStores: Iterable[StateStore[_]] = {
+ Iterable(statsStore, consoleStore, breakpointStore, metadataStore, reconfigurationStore)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/storage/StateStore.scala b/amber/src/main/scala/org/apache/texera/web/storage/StateStore.scala
new file mode 100644
index 00000000000..eb49e97b96f
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/storage/StateStore.scala
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.storage
+
+import io.reactivex.rxjava3.core.{Observable, Single}
+import io.reactivex.rxjava3.disposables.Disposable
+import io.reactivex.rxjava3.subjects.BehaviorSubject
+import org.apache.texera.amber.engine.common.Utils.withLock
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+
+import java.util
+import java.util.concurrent.locks.ReentrantLock
+import scala.collection.mutable
+
+class StateStore[T](defaultState: T) {
+
+ private val stateSubject = BehaviorSubject.createDefault(defaultState)
+ private val serializedSubject = stateSubject.toSerialized
+ private implicit val lock: ReentrantLock = new ReentrantLock()
+ private val diffHandlers = new mutable.ArrayBuffer[(T, T) => Iterable[TexeraWebSocketEvent]]
+ private val diffSubject = serializedSubject
+ .startWith(Single.just(defaultState))
+ .buffer(2, 1)
+ .filter(states => states.get(0) != states.get(1))
+ .map[Iterable[TexeraWebSocketEvent]] { states: util.List[T] =>
+ withLock {
+ diffHandlers.flatMap(f => f(states.get(0), states.get(1)))
+ }
+ }
+
+ def getState: T = stateSubject.getValue
+
+ def updateState(func: T => T): Unit = {
+ withLock {
+ val newState = func(stateSubject.getValue)
+ serializedSubject.onNext(newState)
+ }
+ }
+
+ def registerDiffHandler(handler: (T, T) => Iterable[TexeraWebSocketEvent]): Disposable = {
+ withLock {
+ diffHandlers.append(handler)
+ }
+ Disposable.fromAction { () =>
+ withLock {
+ diffHandlers -= handler
+ }
+ }
+ }
+
+ def getWebsocketEventObservable: Observable[Iterable[TexeraWebSocketEvent]] =
+ diffSubject.onTerminateDetach()
+
+ def getStateObservable: Observable[T] = serializedSubject.onTerminateDetach()
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/web/storage/WorkflowStateStore.scala b/amber/src/main/scala/org/apache/texera/web/storage/WorkflowStateStore.scala
new file mode 100644
index 00000000000..f207dfb85f4
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/web/storage/WorkflowStateStore.scala
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.storage
+
+import org.apache.texera.amber.core.storage.result.WorkflowResultStore
+
+// states that across executions.
+class WorkflowStateStore {
+ val resultStore = new StateStore(WorkflowResultStore())
+
+ def getAllStores: Iterable[StateStore[_]] = {
+ Iterable(resultStore)
+ }
+
+}
diff --git a/amber/src/main/scala/org/apache/texera/workflow/LogicalLink.scala b/amber/src/main/scala/org/apache/texera/workflow/LogicalLink.scala
new file mode 100644
index 00000000000..5bbc9164b72
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/workflow/LogicalLink.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.workflow
+
+import com.fasterxml.jackson.annotation.{JsonCreator, JsonProperty}
+import org.apache.texera.amber.core.virtualidentity.OperatorIdentity
+import org.apache.texera.amber.core.workflow.PortIdentity
+
+case class LogicalLink(
+ @JsonProperty("fromOpId") fromOpId: OperatorIdentity,
+ fromPortId: PortIdentity,
+ @JsonProperty("toOpId") toOpId: OperatorIdentity,
+ toPortId: PortIdentity
+) {
+ @JsonCreator
+ def this(
+ @JsonProperty("fromOpId") fromOpId: String,
+ fromPortId: PortIdentity,
+ @JsonProperty("toOpId") toOpId: String,
+ toPortId: PortIdentity
+ ) = {
+ this(OperatorIdentity(fromOpId), fromPortId, OperatorIdentity(toOpId), toPortId)
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/workflow/LogicalPlan.scala b/amber/src/main/scala/org/apache/texera/workflow/LogicalPlan.scala
new file mode 100644
index 00000000000..974d17f40a4
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/workflow/LogicalPlan.scala
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.workflow
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.storage.FileResolver
+import org.apache.texera.amber.core.virtualidentity.OperatorIdentity
+import org.apache.texera.amber.operator.LogicalOp
+import org.apache.texera.amber.operator.source.scan.ScanSourceOpDesc
+import org.apache.texera.web.model.websocket.request.LogicalPlanPojo
+import org.jgrapht.graph.DirectedAcyclicGraph
+import org.jgrapht.util.SupplierUtil
+
+import java.util
+import scala.collection.mutable.ArrayBuffer
+import scala.util.{Failure, Success, Try}
+
+object LogicalPlan {
+
+ private def toJgraphtDAG(
+ operatorList: List[LogicalOp],
+ links: List[LogicalLink]
+ ): DirectedAcyclicGraph[OperatorIdentity, LogicalLink] = {
+ val workflowDag =
+ new DirectedAcyclicGraph[OperatorIdentity, LogicalLink](
+ null, // vertexSupplier
+ SupplierUtil.createSupplier(classOf[LogicalLink]), // edgeSupplier
+ false, // weighted
+ true // allowMultipleEdges
+ )
+ operatorList.foreach(op => workflowDag.addVertex(op.operatorIdentifier))
+ links.foreach(l =>
+ workflowDag.addEdge(
+ l.fromOpId,
+ l.toOpId,
+ l
+ )
+ )
+ workflowDag
+ }
+
+ def apply(
+ pojo: LogicalPlanPojo
+ ): LogicalPlan = {
+ LogicalPlan(pojo.operators, pojo.links)
+ }
+}
+
+case class LogicalPlan(
+ operators: List[LogicalOp],
+ links: List[LogicalLink]
+) extends LazyLogging {
+
+ private lazy val operatorMap: Map[OperatorIdentity, LogicalOp] =
+ operators.map(op => (op.operatorIdentifier, op)).toMap
+
+ private lazy val jgraphtDag: DirectedAcyclicGraph[OperatorIdentity, LogicalLink] =
+ LogicalPlan.toJgraphtDAG(operators, links)
+
+ def getTopologicalOpIds: util.Iterator[OperatorIdentity] = jgraphtDag.iterator()
+
+ def getOperator(opId: OperatorIdentity): LogicalOp = operatorMap(opId)
+
+ def getTerminalOperatorIds: List[OperatorIdentity] =
+ operatorMap.keys
+ .filter(op => jgraphtDag.outDegreeOf(op) == 0)
+ .toList
+
+ def getUpstreamLinks(opId: OperatorIdentity): List[LogicalLink] = {
+ links.filter(l => l.toOpId == opId)
+ }
+
+ /**
+ * Resolve all user-given filename for the scan source operators to URIs, and call op.setFileUri to set the URi
+ *
+ * @param errorList if given, put errors during resolving to it
+ */
+ def resolveScanSourceOpFileName(
+ errorList: Option[ArrayBuffer[(OperatorIdentity, Throwable)]]
+ ): Unit = {
+ operators.foreach {
+ case operator @ (scanOp: ScanSourceOpDesc) =>
+ Try {
+ // Resolve file path for ScanSourceOpDesc
+ val fileName = scanOp.fileName.getOrElse(throw new RuntimeException("no input file name"))
+ val fileUri = FileResolver.resolve(fileName) // Convert to URI
+
+ // Set the URI in the ScanSourceOpDesc
+ scanOp.setResolvedFileName(fileUri)
+ } match {
+ case Success(_) => // Successfully resolved and set the file URI
+
+ case Failure(err) =>
+ logger.error("Error resolving file path for ScanSourceOpDesc", err)
+ errorList match {
+ case Some(errList) =>
+ errList.append((operator.operatorIdentifier, err))
+ case None =>
+ // Throw the error if no errorList is provided
+ throw err
+ }
+ }
+
+ case _ => // Skip non-ScanSourceOpDesc operators
+ }
+ }
+}
diff --git a/amber/src/main/scala/org/apache/texera/workflow/WorkflowCompiler.scala b/amber/src/main/scala/org/apache/texera/workflow/WorkflowCompiler.scala
new file mode 100644
index 00000000000..b93aa3e4db3
--- /dev/null
+++ b/amber/src/main/scala/org/apache/texera/workflow/WorkflowCompiler.scala
@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.workflow
+
+import com.typesafe.scalalogging.LazyLogging
+import org.apache.texera.amber.core.virtualidentity.OperatorIdentity
+import org.apache.texera.amber.core.workflow._
+import org.apache.texera.amber.engine.architecture.controller.Workflow
+import org.apache.texera.web.model.websocket.request.LogicalPlanPojo
+
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters.IteratorHasAsScala
+import scala.util.{Failure, Success, Try}
+
+class WorkflowCompiler(
+ context: WorkflowContext
+) extends LazyLogging {
+
+ /**
+ * Function to expand logical plan to physical plan
+ * @return the expanded physical plan and a set of output ports that need storage
+ */
+ private def expandLogicalPlan(
+ logicalPlan: LogicalPlan,
+ logicalOpsToViewResult: List[String],
+ errorList: Option[ArrayBuffer[(OperatorIdentity, Throwable)]]
+ ): (PhysicalPlan, Set[GlobalPortIdentity]) = {
+ val terminalLogicalOps = logicalPlan.getTerminalOperatorIds
+ val logicalOpsNeedingStorage =
+ (terminalLogicalOps ++ logicalOpsToViewResult.map(OperatorIdentity(_))).toSet
+ var physicalPlan = PhysicalPlan(operators = Set.empty, links = Set.empty)
+ val outputPortsNeedingStorage: mutable.HashSet[GlobalPortIdentity] = mutable.HashSet()
+
+ logicalPlan.getTopologicalOpIds.asScala.foreach(logicalOpId =>
+ Try {
+ val logicalOp = logicalPlan.getOperator(logicalOpId)
+
+ val subPlan = logicalOp.getPhysicalPlan(context.workflowId, context.executionId)
+ subPlan
+ .topologicalIterator()
+ .map(subPlan.getOperator)
+ .foreach({ physicalOp =>
+ {
+ val externalLinks = logicalPlan
+ .getUpstreamLinks(logicalOp.operatorIdentifier)
+ .filter(link => physicalOp.inputPorts.contains(link.toPortId))
+ .flatMap { link =>
+ physicalPlan
+ .getPhysicalOpsOfLogicalOp(link.fromOpId)
+ .find(_.outputPorts.contains(link.fromPortId))
+ .map(fromOp =>
+ PhysicalLink(fromOp.id, link.fromPortId, physicalOp.id, link.toPortId)
+ )
+ }
+
+ val internalLinks = subPlan.getUpstreamPhysicalLinks(physicalOp.id)
+
+ // Add the operator to the physical plan
+ physicalPlan = physicalPlan.addOperator(physicalOp.propagateSchema())
+
+ // Add all the links to the physical plan
+ physicalPlan = (externalLinks ++ internalLinks)
+ .foldLeft(physicalPlan) { (plan, link) => plan.addLink(link) }
+
+ // **Check for Python-based operator errors during code generation**
+ if (physicalOp.isPythonBased) {
+ val code = physicalOp.getCode
+ val exceptionPattern = """#EXCEPTION DURING CODE GENERATION:\s*(.*)""".r
+
+ exceptionPattern.findFirstMatchIn(code).foreach { matchResult =>
+ val errorMessage = matchResult.group(1).trim
+ val error =
+ new RuntimeException(s"Operator is not configured properly: $errorMessage")
+
+ errorList match {
+ case Some(list) => list.append((logicalOpId, error)) // Store error and continue
+ case None => throw error // Throw immediately if no error list is provided
+ }
+ }
+ }
+ }
+ })
+
+ // convert logical operators needing storage to output ports needing storage
+ subPlan
+ .topologicalIterator()
+ .filter(opId => logicalOpsNeedingStorage.contains(opId.logicalOpId))
+ .map(physicalPlan.getOperator)
+ .foreach { physicalOp =>
+ physicalOp.outputPorts
+ .filterNot(_._1.internal)
+ .foreach {
+ case (outputPortId, _) =>
+ outputPortsNeedingStorage += GlobalPortIdentity(
+ opId = physicalOp.id,
+ portId = outputPortId
+ )
+ }
+ }
+ } match {
+ case Success(_) =>
+
+ case Failure(err) =>
+ errorList match {
+ case Some(list) => list.append((logicalOpId, err))
+ case None => throw err
+ }
+ }
+ )
+ (physicalPlan, outputPortsNeedingStorage.toSet)
+ }
+
+ /**
+ * Compile a workflow to physical plan, along with the schema propagation result and error(if any)
+ *
+ * Comparing to WorkflowCompilingService's compiler, which is used solely for workflow editing,
+ * This compile is used before executing the workflow.
+ *
+ * TODO: we should consider merge this compile with WorkflowCompilingService's compile
+ * @param logicalPlanPojo the pojo parsed from workflow str provided by user
+ * @return Workflow, containing the physical plan, logical plan and workflow context
+ */
+ def compile(
+ logicalPlanPojo: LogicalPlanPojo
+ ): Workflow = {
+ // 1. convert the pojo to logical plan
+ val logicalPlan: LogicalPlan = LogicalPlan(logicalPlanPojo)
+
+ // 2. resolve the file name in each scan source operator
+ logicalPlan.resolveScanSourceOpFileName(None)
+
+ // 3. expand the logical plan to the physical plan, and get a set of output ports that need storage
+ val (physicalPlan, outputPortsNeedingStorage) =
+ expandLogicalPlan(logicalPlan, logicalPlanPojo.opsToViewResult, None)
+
+ context.workflowSettings = context.workflowSettings.copy(
+ outputPortsNeedingStorage = outputPortsNeedingStorage
+ )
+
+ Workflow(context, logicalPlan, physicalPlan)
+ }
+}
diff --git a/amber/src/test/java/org/apache/texera/web/resource/dashboard/user/dataset/GitVersionControlLocalFileStorageSpec.java b/amber/src/test/java/org/apache/texera/web/resource/dashboard/user/dataset/GitVersionControlLocalFileStorageSpec.java
new file mode 100644
index 00000000000..a07410b1450
--- /dev/null
+++ b/amber/src/test/java/org/apache/texera/web/resource/dashboard/user/dataset/GitVersionControlLocalFileStorageSpec.java
@@ -0,0 +1,222 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.dataset;
+
+import org.apache.texera.amber.core.storage.util.dataset.GitVersionControlLocalFileStorage;
+import org.apache.texera.amber.core.storage.util.dataset.PhysicalFileNode;
+import org.eclipse.jgit.api.errors.GitAPIException;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+public class GitVersionControlLocalFileStorageSpec {
+
+ private Path testRepoPath;
+
+ private List testRepoMasterCommitHashes;
+ private final String testFile1Name = "testFile1.txt";
+
+ private final String testFile2Name = "testFile2.txt";
+ private final String testDirectoryName = "testDir";
+
+ private final String testFile1ContentV1 = "This is a test file1 v1";
+ private final String testFile1ContentV2 = "This is a test file1 v2";
+ private final String testFile1ContentV3 = "This is a test file1 v3";
+
+ private final String testFile2Content = "This is a test file2 in the testDir";
+
+ private void writeFileToRepo(Path filePath, String fileContent) throws IOException, GitAPIException {
+ try (ByteArrayInputStream input = new ByteArrayInputStream(fileContent.getBytes())) {
+ GitVersionControlLocalFileStorage.writeFileToRepo(testRepoPath, filePath, input);
+ }
+ }
+
+ @Before
+ public void setUp() throws IOException, GitAPIException {
+ // Create a temporary directory for the repository
+ testRepoPath = Files.createTempDirectory("testRepo");
+ GitVersionControlLocalFileStorage.initRepo(testRepoPath);
+
+ Path file1Path = testRepoPath.resolve(testFile1Name);
+ // Version 1
+ String v1Hash = GitVersionControlLocalFileStorage.withCreateVersion(
+ testRepoPath,
+ "v1",
+ () -> {
+ try {
+ writeFileToRepo(file1Path, testFile1ContentV1);
+ } catch (IOException | GitAPIException e) {
+ throw new RuntimeException(e);
+ }
+ });
+
+ String v2Hash = GitVersionControlLocalFileStorage.withCreateVersion(
+ testRepoPath,
+ "v2",
+ () -> {
+ try {
+ writeFileToRepo(file1Path, testFile1ContentV2);
+ } catch (IOException | GitAPIException e) {
+ throw new RuntimeException(e);
+ }
+ });
+
+ // Version 3
+ String v3Hash = GitVersionControlLocalFileStorage.withCreateVersion(
+ testRepoPath,
+ "v3",
+ () -> {
+ try {
+ writeFileToRepo(file1Path, testFile1ContentV3);
+ } catch (IOException | GitAPIException e) {
+ throw new RuntimeException(e);
+ }
+ });
+
+ testRepoMasterCommitHashes = new ArrayList() {{
+ add(v1Hash);
+ add(v2Hash);
+ add(v3Hash);
+ }};
+ }
+
+ @After
+ public void tearDown() throws IOException {
+ // Clean up the test repository directory
+ GitVersionControlLocalFileStorage.deleteRepo(testRepoPath);
+ }
+
+ @Test
+ public void testFileContentAcrossVersions() throws IOException, GitAPIException {
+ // File path for the test file
+ Path filePath = testRepoPath.resolve(testFile1Name);
+
+ // testRepoMasterCommitHashes is populated in chronological order: v1, v2, v3
+ // Retrieve and compare file content for version 1
+ ByteArrayOutputStream outputV1 = new ByteArrayOutputStream();
+ GitVersionControlLocalFileStorage.retrieveFileContentOfVersion(testRepoPath, testRepoMasterCommitHashes.get(0), filePath, outputV1);
+ String retrievedContentV1 = outputV1.toString();
+ Assert.assertEquals(
+ "Content for version 1 does not match",
+ testFile1ContentV1,
+ retrievedContentV1);
+
+ // Retrieve and compare file content for version 2
+ ByteArrayOutputStream outputV2 = new ByteArrayOutputStream();
+ GitVersionControlLocalFileStorage.retrieveFileContentOfVersion(testRepoPath, testRepoMasterCommitHashes.get(1), filePath, outputV2);
+ String retrievedContentV2 = outputV2.toString();
+ Assert.assertEquals(
+ "Content for version 2 does not match",
+ testFile1ContentV2,
+ retrievedContentV2);
+
+ // Retrieve and compare file content for version 3
+ ByteArrayOutputStream outputV3 = new ByteArrayOutputStream();
+ GitVersionControlLocalFileStorage.retrieveFileContentOfVersion(testRepoPath, testRepoMasterCommitHashes.get(2), filePath, outputV3);
+ String retrievedContentV3 = outputV3.toString();
+ Assert.assertEquals(
+ "Content for version 3 does not match",
+ testFile1ContentV3,
+ retrievedContentV3);
+ }
+
+ @Test
+ public void testFileTreeRetrieval() throws Exception {
+ // File path for the test file
+ Path file1Path = testRepoPath.resolve(testFile1Name);
+ PhysicalFileNode file1Node = new PhysicalFileNode(testRepoPath, file1Path, Files.size(file1Path));
+ Set physicalFileNodes = new HashSet() {{
+ add(file1Node);
+ }};
+
+ // first retrieve the latest version's file tree
+ Assert.assertEquals("File Tree should match",
+ physicalFileNodes,
+ GitVersionControlLocalFileStorage.retrieveRootFileNodesOfVersion(testRepoPath, testRepoMasterCommitHashes.get(testRepoMasterCommitHashes.size() - 1)));
+
+ // now we add a new file testDir/testFile2.txt
+ Path testDirPath = testRepoPath.resolve(testDirectoryName);
+ Path file2Path = testDirPath.resolve(testFile2Name);
+
+ String v4Hash = GitVersionControlLocalFileStorage.withCreateVersion(testRepoPath, "v4", () -> {
+ try {
+ writeFileToRepo(file2Path, testFile2Content);
+ } catch (IOException | GitAPIException e) {
+ throw new RuntimeException(e);
+ }
+ });
+ testRepoMasterCommitHashes.add(v4Hash);
+
+ PhysicalFileNode dirNode = new PhysicalFileNode(testRepoPath, testDirPath, 0); // Directories typically have size 0
+ dirNode.addChildNode(new PhysicalFileNode(testRepoPath, file2Path, Files.size(file2Path)));
+ // update the expected fileNodes
+ physicalFileNodes.add(dirNode);
+
+ // check the file tree
+ Assert.assertEquals(
+ "File Tree should match",
+ physicalFileNodes,
+ GitVersionControlLocalFileStorage.retrieveRootFileNodesOfVersion(testRepoPath, v4Hash));
+
+ // now we delete the file1, check the filetree
+ String v5Hash = GitVersionControlLocalFileStorage.withCreateVersion(testRepoPath, "v5", () -> {
+ try {
+ GitVersionControlLocalFileStorage.removeFileFromRepo(testRepoPath, file1Path);
+ } catch (IOException | GitAPIException e) {
+ throw new RuntimeException(e);
+ }
+ });
+
+ physicalFileNodes.remove(file1Node);
+ Assert.assertEquals(
+ "File1 should be gone",
+ physicalFileNodes,
+ GitVersionControlLocalFileStorage.retrieveRootFileNodesOfVersion(testRepoPath, v5Hash)
+ );
+
+ }
+
+ @Test
+ public void testUncommittedCheckAndRecoverToLatest() throws Exception {
+ Path tempFilePath = testRepoPath.resolve("tempFile");
+ String content = "some random content";
+ writeFileToRepo(tempFilePath, content);
+
+ Assert.assertTrue(
+ "There should be some uncommitted changes",
+ GitVersionControlLocalFileStorage.hasUncommittedChanges(testRepoPath));
+
+ GitVersionControlLocalFileStorage.discardUncommittedChanges(testRepoPath);
+
+ Assert.assertFalse("There should be no uncommitted changes",
+ GitVersionControlLocalFileStorage.hasUncommittedChanges(testRepoPath));
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/TrivialControlSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/TrivialControlSpec.scala
new file mode 100644
index 00000000000..79726f7fbf0
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/TrivialControlSpec.scala
@@ -0,0 +1,194 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control
+
+import org.apache.pekko.actor.{ActorRef, ActorSystem, PoisonPill, Props}
+import org.apache.pekko.testkit.{TestKit, TestProbe}
+import io.grpc.MethodDescriptor
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.{
+ GetActorRef,
+ NetworkAck,
+ NetworkMessage,
+ RegisterActorRef
+}
+import org.apache.texera.amber.engine.architecture.control.utils.TrivialControlTester
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ IntResponse,
+ ReturnInvocation,
+ StringResponse
+}
+import org.apache.texera.amber.engine.architecture.rpc.testerservice.RPCTesterGrpc._
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.scalatest.wordspec.AnyWordSpecLike
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
+
+import scala.collection.mutable
+import scala.concurrent.duration._
+
+class TrivialControlSpec
+ extends TestKit(ActorSystem("TrivialControlSpec"))
+ with AnyWordSpecLike
+ with BeforeAndAfterEach
+ with BeforeAndAfterAll {
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ def testControl[T](
+ numActors: Int,
+ eventPairs: ((MethodDescriptor[_, _], ControlRequest), T)*
+ ): Unit = {
+ val (events, expectedValues) = eventPairs.unzip
+ val (probe, idMap) = setUp(numActors, events: _*)
+ var flag = 0
+ while (flag < expectedValues.length) {
+ probe.receiveOne(10.seconds) match {
+ case null =>
+ throw new AssertionError(
+ s"timeout: received $flag of ${expectedValues.length} expected returns"
+ )
+ case GetActorRef(id, replyTo) =>
+ replyTo.foreach { actor =>
+ actor ! RegisterActorRef(id, idMap(id))
+ }
+ case NetworkMessage(
+ msgID,
+ workflowMsg @ WorkflowFIFOMessage(_, _, ReturnInvocation(id, returnValue))
+ ) =>
+ probe.sender() ! NetworkAck(
+ msgID,
+ getInMemSize(workflowMsg),
+ 0L // no queued credit
+ )
+ assert(returnValue.asInstanceOf[T] == expectedValues(id.toInt))
+ flag += 1
+ case _ =>
+ //skip
+ }
+ }
+ idMap.foreach { x =>
+ x._2 ! PoisonPill
+ }
+ }
+
+ def setUp(
+ numActors: Int,
+ cmd: (MethodDescriptor[_, _], ControlRequest)*
+ ): (TestProbe, mutable.HashMap[ActorVirtualIdentity, ActorRef]) = {
+ val probe = TestProbe()
+ val idMap = mutable.HashMap[ActorVirtualIdentity, ActorRef]()
+ for (i <- 0 until numActors) {
+ val id = ActorVirtualIdentity(s"$i")
+ val ref =
+ probe.childActorOf(Props(new TrivialControlTester(id)))
+ idMap(id) = ref
+ }
+ idMap(CONTROLLER) = probe.ref
+ var seqNum = 0
+ cmd.foreach {
+ case (methodName, msg) =>
+ probe.send(
+ idMap(ActorVirtualIdentity("0")),
+ NetworkMessage(
+ seqNum,
+ WorkflowFIFOMessage(
+ ChannelIdentity(CONTROLLER, ActorVirtualIdentity("0"), isControl = true),
+ seqNum,
+ ControlInvocation(
+ methodName,
+ msg,
+ AsyncRPCContext(CONTROLLER, ActorVirtualIdentity("0")),
+ seqNum
+ )
+ )
+ )
+ )
+ seqNum += 1
+ }
+ (probe, idMap)
+ }
+
+ "testers" should {
+
+ "execute Ping Pong" in {
+ testControl(2, ((METHOD_SEND_PING, Ping(1, 5, ActorVirtualIdentity("1"))), IntResponse(5)))
+ }
+
+ "execute Ping Pong 2 times" in {
+ testControl(
+ 2,
+ ((METHOD_SEND_PING, Ping(1, 4, ActorVirtualIdentity("1"))), IntResponse(4)),
+ ((METHOD_SEND_PING, Ping(10, 13, ActorVirtualIdentity("1"))), IntResponse(13))
+ )
+ }
+
+ "execute Chain" in {
+ testControl(
+ 10,
+ (
+ (METHOD_SEND_CHAIN, Chain((1 to 9).map(i => ActorVirtualIdentity(i.toString)))),
+ StringResponse("9")
+ )
+ )
+ }
+
+ "execute Collect" in {
+ testControl(
+ 4,
+ (
+ (METHOD_SEND_COLLECT, Collect((1 to 3).map(i => ActorVirtualIdentity(i.toString)))),
+ StringResponse("finished")
+ )
+ )
+ }
+
+ "execute RecursiveCall" in {
+ testControl(1, ((METHOD_SEND_RECURSION, Recursion(0)), StringResponse("0")))
+ }
+
+ "execute MultiCall" in {
+ testControl(
+ 10,
+ (
+ (METHOD_SEND_MULTI_CALL, MultiCall((1 to 9).map(i => ActorVirtualIdentity(i.toString)))),
+ StringResponse("finished")
+ )
+ )
+ }
+
+ "execute NestedCall" in {
+ testControl(1, ((METHOD_SEND_NESTED, Nested(5)), StringResponse("Hello World!")))
+ }
+
+ "execute ErrorCall" in {
+ assertThrows[RuntimeException] {
+ testControl(1, ((METHOD_SEND_ERROR_COMMAND, ErrorCommand()), ()))
+ }
+
+ }
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/ChainHandler.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/ChainHandler.scala
new file mode 100644
index 00000000000..9e7eff617d9
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/ChainHandler.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+
+trait ChainHandler {
+ this: TesterAsyncRPCHandlerInitializer =>
+
+ override def sendChain(request: Chain, ctx: AsyncRPCContext): Future[StringResponse] = {
+ println(s"chained $myID")
+ if (request.nexts.isEmpty) {
+ Future(StringResponse(myID.name))
+ } else {
+ getProxy.sendChain(Chain(request.nexts.drop(1)), mkContext(request.nexts.head)).map { x =>
+ println(s"chain returns from $x")
+ x
+ }
+ }
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/CollectHandler.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/CollectHandler.scala
new file mode 100644
index 00000000000..f36ca3c2b7b
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/CollectHandler.scala
@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+
+import scala.util.Random
+
+trait CollectHandler {
+ this: TesterAsyncRPCHandlerInitializer =>
+
+ override def sendCollect(request: Collect, ctx: AsyncRPCContext): Future[StringResponse] = {
+ println(s"start collecting numbers.")
+ val p = Future.collect(
+ request.workers.indices.map(i =>
+ getProxy.sendGenerateNumber(GenerateNumber(), mkContext(request.workers(i)))
+ )
+ )
+ p.map { res =>
+ println(s"collected: ${res.mkString(" ")}")
+ StringResponse("finished")
+ }
+ }
+
+ override def sendGenerateNumber(
+ request: GenerateNumber,
+ ctx: AsyncRPCContext
+ ): Future[IntResponse] = {
+ IntResponse(Random.nextInt())
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/ErrorHandler.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/ErrorHandler.scala
new file mode 100644
index 00000000000..a1e6b94d05e
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/ErrorHandler.scala
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+
+trait ErrorHandler {
+ this: TesterAsyncRPCHandlerInitializer =>
+
+ override def sendErrorCommand(
+ request: ErrorCommand,
+ ctx: AsyncRPCContext
+ ): Future[StringResponse] = {
+ throw new RuntimeException("this is an EXPECTED exception for testing")
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/MultiCallHandler.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/MultiCallHandler.scala
new file mode 100644
index 00000000000..21400b83f2c
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/MultiCallHandler.scala
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+
+trait MultiCallHandler {
+ this: TesterAsyncRPCHandlerInitializer =>
+
+ override def sendMultiCall(request: MultiCall, ctx: AsyncRPCContext): Future[StringResponse] = {
+ getProxy
+ .sendChain(Chain(request.seq), myID)
+ .flatMap(x => getProxy.sendRecursion(Recursion(1), mkContext(ActorVirtualIdentity(x.value))))
+ .flatMap(ret => getProxy.sendCollect(Collect(request.seq.take(3)), myID))
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/NestedHandler.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/NestedHandler.scala
new file mode 100644
index 00000000000..31aa70d4470
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/NestedHandler.scala
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+
+trait NestedHandler {
+ this: TesterAsyncRPCHandlerInitializer =>
+
+ override def sendNested(request: Nested, ctx: AsyncRPCContext): Future[StringResponse] = {
+ getProxy
+ .sendPass(Pass("Hello"), myID)
+ .flatMap(ret => getProxy.sendPass(Pass(ret.value + " "), myID))
+ .flatMap(ret => getProxy.sendPass(Pass(ret.value + "World!"), myID))
+ }
+
+ override def sendPass(request: Pass, ctx: AsyncRPCContext): Future[StringResponse] = {
+ StringResponse(request.value)
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/PingPongHandler.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/PingPongHandler.scala
new file mode 100644
index 00000000000..5ea7bf298e7
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/PingPongHandler.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+
+trait PingPongHandler {
+ this: TesterAsyncRPCHandlerInitializer =>
+
+ override def sendPing(ping: Ping, ctx: AsyncRPCContext): Future[IntResponse] = {
+ println(s"${ping.i} ping")
+ if (ping.i < ping.end) {
+ getProxy.sendPong(Pong(ping.i + 1, ping.end, myID), ping.to).map { ret: IntResponse =>
+ println(s"${ping.i} ping replied with value ${ret.value}!")
+ ret
+ }
+ } else {
+ Future(ping.i)
+ }
+ }
+
+ override def sendPong(pong: Pong, ctx: AsyncRPCContext): Future[IntResponse] = {
+ println(s"${pong.i} pong")
+ if (pong.i < pong.end) {
+ getProxy.sendPing(Ping(pong.i + 1, pong.end, myID), pong.to).map { ret: IntResponse =>
+ println(s"${pong.i} pong replied with value ${ret.value}!")
+ ret
+ }
+ } else {
+ Future(pong.i)
+ }
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/RecursionHandler.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/RecursionHandler.scala
new file mode 100644
index 00000000000..7d2b7a38e22
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/RecursionHandler.scala
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+
+trait RecursionHandler {
+ this: TesterAsyncRPCHandlerInitializer =>
+
+ override def sendRecursion(r: Recursion, ctx: AsyncRPCContext): Future[StringResponse] = {
+ if (r.i < 5) {
+ println(r.i)
+ getProxy.sendRecursion(Recursion(r.i + 1), myID).map { res =>
+ println(res)
+ r.i.toString
+ }
+ } else {
+ Future(r.i.toString)
+ }
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/TesterAsyncRPCHandlerInitializer.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/TesterAsyncRPCHandlerInitializer.scala
new file mode 100644
index 00000000000..24033da6bbf
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/TesterAsyncRPCHandlerInitializer.scala
@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
+import org.apache.texera.amber.engine.architecture.control.utils.TrivialControlTester.ControlTesterRPCClient
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.AsyncRPCContext
+import org.apache.texera.amber.engine.architecture.rpc.testerservice.RPCTesterFs2Grpc
+import org.apache.texera.amber.engine.common.rpc.{AsyncRPCHandlerInitializer, AsyncRPCServer}
+
+class TesterAsyncRPCHandlerInitializer(
+ val myID: ActorVirtualIdentity,
+ source: ControlTesterRPCClient,
+ receiver: AsyncRPCServer
+) extends AsyncRPCHandlerInitializer(source, receiver)
+ with RPCTesterFs2Grpc[Future, AsyncRPCContext]
+ with PingPongHandler
+ with ChainHandler
+ with MultiCallHandler
+ with CollectHandler
+ with NestedHandler
+ with RecursionHandler
+ with ErrorHandler {
+ def getProxy: RPCTesterFs2Grpc[Future, AsyncRPCContext] = source.getProxy
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/TrivialControlTester.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/TrivialControlTester.scala
new file mode 100644
index 00000000000..ca86338d368
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/control/utils/TrivialControlTester.scala
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.control.utils
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkAck
+import org.apache.texera.amber.engine.architecture.common.{AmberProcessor, WorkflowActor}
+import org.apache.texera.amber.engine.architecture.control.utils.TrivialControlTester.ControlTesterRPCClient
+import org.apache.texera.amber.engine.architecture.messaginglayer.{
+ NetworkInputGateway,
+ NetworkOutputGateway
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.AsyncRPCContext
+import org.apache.texera.amber.engine.architecture.rpc.testerservice.RPCTesterFs2Grpc
+import org.apache.texera.amber.engine.common.CheckpointState
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataPayload,
+ DirectControlMessagePayload,
+ WorkflowFIFOMessage
+}
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient
+
+object TrivialControlTester {
+ class ControlTesterRPCClient(
+ inputGateway: NetworkInputGateway,
+ outputGateway: NetworkOutputGateway,
+ actorId: ActorVirtualIdentity
+ ) extends AsyncRPCClient(inputGateway, outputGateway, actorId) {
+ val getProxy: RPCTesterFs2Grpc[Future, AsyncRPCContext] =
+ AsyncRPCClient
+ .createProxy[RPCTesterFs2Grpc[Future, AsyncRPCContext]](createPromise, outputGateway)
+ }
+}
+
+class TrivialControlTester(
+ id: ActorVirtualIdentity
+) extends WorkflowActor(replayLogConfOpt = None, actorId = id) {
+ val ap = new AmberProcessor(
+ id,
+ {
+ case Left(value) => ???
+ case Right(value) => transferService.send(value)
+ }
+ ) {
+ override val asyncRPCClient = new ControlTesterRPCClient(inputGateway, outputGateway, id)
+ }
+ val initializer =
+ new TesterAsyncRPCHandlerInitializer(ap.actorId, ap.asyncRPCClient, ap.asyncRPCServer)
+
+ override def handleInputMessage(id: Long, workflowMsg: WorkflowFIFOMessage): Unit = {
+ val channel = ap.inputGateway.getChannel(workflowMsg.channelId)
+ channel.acceptMessage(workflowMsg)
+ while (channel.isEnabled && channel.hasMessage) {
+ val msg = channel.take
+ msg.payload match {
+ case payload: DirectControlMessagePayload => ap.processDCM(msg.channelId, payload)
+ case _: DataPayload => ???
+ case _ => ???
+ }
+ }
+ sender() ! NetworkAck(id, getInMemSize(workflowMsg), getQueuedCredit(workflowMsg.channelId))
+ }
+
+ /** flow-control */
+ override def getQueuedCredit(channelId: ChannelIdentity): Long = 0L
+
+ override def preStart(): Unit = {
+ transferService.initialize()
+ }
+
+ override def handleBackpressure(isBackpressured: Boolean): Unit = {}
+
+ override def initState(): Unit = {}
+
+ override def loadFromCheckpoint(chkpt: CheckpointState): Unit = {}
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/ControllerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/ControllerSpec.scala
new file mode 100644
index 00000000000..f3f8ca5dd0d
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/ControllerSpec.scala
@@ -0,0 +1,334 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.pekko.actor.{ActorSystem, Props}
+import org.apache.pekko.testkit.{ImplicitSender, TestKit}
+import org.apache.pekko.util.Timeout
+import org.apache.texera.amber.clustering.SingleNodeListener
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpecLike
+
+import scala.concurrent.ExecutionContextExecutor
+import scala.concurrent.duration._
+
+class ControllerSpec
+ extends TestKit(ActorSystem("ControllerSpec"))
+ with ImplicitSender
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll {
+
+ implicit val timeout: Timeout = Timeout(5.seconds)
+ implicit val executionContext: ExecutionContextExecutor = system.dispatcher
+
+ override def beforeAll(): Unit = {
+ system.actorOf(Props[SingleNodeListener](), "cluster-info")
+ }
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ // private val logicalPlan1 =
+ // """{
+ // |"operators":[
+ // |{"tableName":"D:\\large_input.csv","operatorId":"Scan","operatorType":"LocalScanSource","delimiter":","},
+ // |{"attributeName":0,"keyword":"Asia","operatorId":"KeywordSearch","operatorType":"KeywordMatcher"},
+ // |{"operatorId":"Count","operatorType":"Aggregation"},
+ // |{"operatorId":"Sink","operatorType":"Sink"}],
+ // |"links":[
+ // |{"origin":"Scan","destination":"KeywordSearch"},
+ // |{"origin":"KeywordSearch","destination":"Count"},
+ // |{"origin":"Count","destination":"Sink"}]
+ // |}""".stripMargin
+ //
+ // private val logicalPlan2 =
+ // """{
+ // |"operators":[
+ // |{"tableName":"D:\\large_input.csv","operatorId":"Scan","operatorType":"LocalScanSource","delimiter":","},
+ // |{"operatorId":"Count","operatorType":"Aggregation"},
+ // |{"operatorId":"Sink","operatorType":"Sink"}],
+ // |"links":[
+ // |{"origin":"Scan","destination":"Count"},
+ // |{"origin":"Count","destination":"Sink"}]
+ // |}""".stripMargin
+ //
+ // private val logicalPlan3 =
+ // """{
+ // |"operators":[
+ // |{"tableName":"D:\\test.txt","operatorId":"Scan","operatorType":"LocalScanSource","delimiter":"|"},
+ // |{"attributeName":15,"keyword":"package","operatorId":"KeywordSearch","operatorType":"KeywordMatcher"},
+ // |{"operatorId":"Count","operatorType":"Aggregation"},
+ // |{"operatorId":"Sink","operatorType":"Sink"}],
+ // |"links":[
+ // |{"origin":"Scan","destination":"KeywordSearch"},
+ // |{"origin":"KeywordSearch","destination":"Count"},
+ // |{"origin":"Count","destination":"Sink"}]
+ // |}""".stripMargin
+ //
+ // private val logicalPlan4 =
+ // """{
+ // |"operators":[
+ // |{"tableName":"D:\\test.txt","operatorId":"Scan1","operatorType":"LocalScanSource","delimiter":"|","indicesToKeep":null},
+ // |{"tableName":"D:\\test.txt","operatorId":"Scan2","operatorType":"LocalScanSource","delimiter":"|","indicesToKeep":null},
+ // |{"attributeName":15,"keyword":"package","operatorId":"KeywordSearch","operatorType":"KeywordMatcher"},
+ // |{"operatorId":"Join","operatorType":"HashJoin","innerTableIndex":0,"outerTableIndex":0},
+ // |{"operatorId":"Count","operatorType":"Aggregation"},
+ // |{"operatorId":"Sink","operatorType":"Sink"}],
+ // |"links":[
+ // |{"origin":"Scan1","destination":"KeywordSearch"},
+ // |{"origin":"KeywordSearch","destination":"Join"},
+ // |{"origin":"Scan2","destination":"Join"},
+ // |{"origin":"Join","destination":"Count"},
+ // |{"origin":"Count","destination":"Sink"}]
+ // |}""".stripMargin
+ //
+ // "A controller" should "be able to set and trigger count breakpoint in the workflow1" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan1))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(30.seconds, ReportState(ControllerState.Ready))
+ // controller ! PassBreakpointTo("KeywordSearch", new CountGlobalBreakpoint("break1", 100000))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // var isCompleted = false
+ // parent.receiveWhile(30.seconds, 10.seconds) {
+ // case ReportState(ControllerState.Paused) =>
+ // controller ! Resume
+ // case ReportState(ControllerState.Completed) =>
+ // isCompleted = true
+ // case _ =>
+ // }
+ // assert(isCompleted)
+ // parent.ref ! PoisonPill
+ // }
+ //
+ // "A controller" should "execute the workflow1 normally" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan1))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(30.seconds, ReportState(ControllerState.Ready))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // parent.expectMsg(1.minute, ReportState(ControllerState.Completed))
+ // parent.ref ! PoisonPill
+ // }
+ //
+ // "A controller" should "execute the workflow3 normally" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan3))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(30.seconds, ReportState(ControllerState.Ready))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // parent.expectMsg(1.minute, ReportState(ControllerState.Completed))
+ // parent.ref ! PoisonPill
+ // }
+ //
+ // "A controller" should "execute the workflow2 normally" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan2))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(ReportState(ControllerState.Ready))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // parent.expectMsg(1.minute, ReportState(ControllerState.Completed))
+ // parent.ref ! PoisonPill
+ // }
+ //
+ // "A controller" should "be able to pause/resume the workflow1" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan1))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(ReportState(ControllerState.Ready))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // controller ! Pause
+ // parent.expectMsg(ReportState(ControllerState.Pausing))
+ // parent.expectMsg(ReportState(ControllerState.Paused))
+ // controller ! Resume
+ // parent.expectMsg(ReportState(ControllerState.Resuming))
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // controller ! Pause
+ // parent.expectMsg(ReportState(ControllerState.Pausing))
+ // parent.expectMsg(ReportState(ControllerState.Paused))
+ // controller ! Resume
+ // parent.expectMsg(ReportState(ControllerState.Resuming))
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // controller ! Pause
+ // parent.expectMsg(ReportState(ControllerState.Pausing))
+ // parent.expectMsg(ReportState(ControllerState.Paused))
+ // controller ! Resume
+ // parent.expectMsg(ReportState(ControllerState.Resuming))
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // controller ! Pause
+ // parent.expectMsg(ReportState(ControllerState.Pausing))
+ // parent.expectMsg(ReportState(ControllerState.Paused))
+ // controller ! Resume
+ // parent.expectMsg(ReportState(ControllerState.Resuming))
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // parent.expectMsg(1.minute, ReportState(ControllerState.Completed))
+ // parent.ref ! PoisonPill
+ // }
+
+ // "A controller" should "be able to modify the logic after pausing the workflow1" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan1))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(30.seconds, ReportState(ControllerState.Ready))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // Thread.sleep(300)
+ // controller ! Pause
+ // parent.expectMsg(ReportState(ControllerState.Pausing))
+ // parent.expectMsg(ReportState(ControllerState.Paused))
+ // controller ! ModifyLogic(
+ // new KeywordSearchMetadata(
+ // OperatorTag("sample", "KeywordSearch"),
+ // Constants.currentWorkerNum,
+ // 0,
+ // "asia"
+ // )
+ // )
+ // parent.expectMsg(Ack)
+ // Thread.sleep(10000)
+ // controller ! Resume
+ // parent.expectMsg(ReportState(ControllerState.Resuming))
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // parent.expectMsg(1.minute, ReportState(ControllerState.Completed))
+ // parent.ref ! PoisonPill
+ // }
+
+ // "A controller" should "be able to set and trigger conditional breakpoint in the workflow1" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan1))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(30.seconds, ReportState(ControllerState.Ready))
+ // controller ! PassBreakpointTo(
+ // "KeywordSearch",
+ // new ConditionalGlobalBreakpoint("break2", x => x.getString(8).toInt == 9884)
+ // )
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // var isCompleted = false
+ // parent.receiveWhile(30.seconds, 10.seconds) {
+ // case ReportState(ControllerState.Paused) =>
+ // controller ! Resume
+ // case ReportState(ControllerState.Completed) =>
+ // isCompleted = true
+ // case _ =>
+ // }
+ // assert(isCompleted)
+ // parent.ref ! PoisonPill
+ // }
+ //
+ // "A controller" should "be able to set and trigger count breakpoint on complete in the workflow1" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan1))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(30.seconds, ReportState(ControllerState.Ready))
+ // controller ! PassBreakpointTo("KeywordSearch", new CountGlobalBreakpoint("break1", 146017))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // var isCompleted = false
+ // parent.receiveWhile(30.seconds, 10.seconds) {
+ // case ReportState(ControllerState.Paused) =>
+ // controller ! Resume
+ // case ReportState(ControllerState.Completed) =>
+ // isCompleted = true
+ // case _ =>
+ // }
+ // assert(isCompleted)
+ // parent.ref ! PoisonPill
+ // }
+ //
+ // "A controller" should "be able to pause/resume with conditional breakpoint in the workflow1" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan1))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(30.seconds, ReportState(ControllerState.Ready))
+ // controller ! PassBreakpointTo(
+ // "KeywordSearch",
+ // new ConditionalGlobalBreakpoint("break2", x => x.getString(8).toInt == 9884)
+ // )
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // val random = new Random()
+ // for (i <- 0 until 100) {
+ // if (random.nextBoolean()) {
+ // controller ! Pause
+ // } else {
+ // controller ! Resume
+ // }
+ // }
+ // controller ! Resume
+ // var isCompleted = false
+ // parent.receiveWhile(30.seconds, 10.seconds) {
+ // case ReportState(ControllerState.Paused) =>
+ // controller ! Resume
+ // case ReportState(ControllerState.Completed) =>
+ // isCompleted = true
+ // case _ =>
+ // }
+ // assert(isCompleted)
+ // parent.ref ! PoisonPill
+ // }
+ //
+ // "A controller" should "be able to pause/resume with count breakpoint in the workflow1" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan1))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(30.seconds, ReportState(ControllerState.Ready))
+ // controller ! PassBreakpointTo("KeywordSearch", new CountGlobalBreakpoint("break1", 100000))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // val random = new Random()
+ // for (i <- 0 until 100) {
+ // if (random.nextBoolean()) {
+ // controller ! Pause
+ // } else {
+ // controller ! Resume
+ // }
+ // }
+ // controller ! Resume
+ // var isCompleted = false
+ // parent.receiveWhile(30.seconds, 10.seconds) {
+ // case ReportState(ControllerState.Paused) =>
+ // controller ! Resume
+ // case ReportState(ControllerState.Completed) =>
+ // isCompleted = true
+ // case _ =>
+ // }
+ // assert(isCompleted)
+ // parent.ref ! PoisonPill
+ // }
+ //
+ // "A controller" should "execute the workflow4 normally" in {
+ // val parent = TestProbe()
+ // val controller = parent.childActorOf(CONTROLLER.props(logicalPlan4))
+ // controller ! AckedControllerInitialization
+ // parent.expectMsg(ReportState(ControllerState.Ready))
+ // controller ! Start
+ // parent.expectMsg(ReportState(ControllerState.Running))
+ // parent.expectMsg(1.minute, ReportState(ControllerState.Completed))
+ // parent.ref ! PoisonPill
+ // }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/WorkflowSchedulerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/WorkflowSchedulerSpec.scala
new file mode 100644
index 00000000000..ac7358b438f
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/WorkflowSchedulerSpec.scala
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller
+
+import org.apache.texera.amber.core.workflow.{PortIdentity, WorkflowContext}
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.engine.e2e.TestUtils.buildWorkflow
+import org.apache.texera.amber.operator.TestOperators
+import org.apache.texera.workflow.LogicalLink
+import org.scalatest.flatspec.AnyFlatSpec
+
+class WorkflowSchedulerSpec extends AnyFlatSpec {
+
+ private def buildHeaderlessCsvKeywordWorkflow() = {
+ val csvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ buildWorkflow(
+ List(csvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(0),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(0)
+ )
+ ),
+ new WorkflowContext()
+ )
+ }
+
+ "WorkflowScheduler.updateSchedule" should "populate the schedule and physicalPlan fields" in {
+ val workflow = buildHeaderlessCsvKeywordWorkflow()
+ val scheduler = new WorkflowScheduler(workflow.context, CONTROLLER)
+
+ assert(scheduler.getSchedule == null)
+ assert(scheduler.physicalPlan == null)
+
+ scheduler.updateSchedule(workflow.physicalPlan)
+
+ assert(scheduler.getSchedule != null)
+ assert(scheduler.physicalPlan != null)
+ assert(scheduler.getSchedule.getRegions.nonEmpty)
+ }
+
+ it should "include every workflow operator in some region of the produced schedule" in {
+ val workflow = buildHeaderlessCsvKeywordWorkflow()
+ val scheduler = new WorkflowScheduler(workflow.context, CONTROLLER)
+ scheduler.updateSchedule(workflow.physicalPlan)
+
+ val operatorsInSchedule = scheduler.getSchedule.getRegions
+ .flatMap(_.getOperators.map(_.id.logicalOpId))
+ .toSet
+ val operatorsInPlan = scheduler.physicalPlan.operators.map(_.id.logicalOpId)
+
+ assert(operatorsInPlan.subsetOf(operatorsInSchedule))
+ }
+
+ "WorkflowScheduler.getNextRegions" should "exhaust the schedule and then return an empty set" in {
+ val workflow = buildHeaderlessCsvKeywordWorkflow()
+ val scheduler = new WorkflowScheduler(workflow.context, CONTROLLER)
+ scheduler.updateSchedule(workflow.physicalPlan)
+
+ val pulledLevels = Iterator
+ .continually(scheduler.getNextRegions)
+ .takeWhile(_.nonEmpty)
+ .toList
+
+ assert(pulledLevels.nonEmpty)
+ assert(scheduler.getNextRegions.isEmpty)
+ }
+
+ it should "yield region sets that together cover every region in the schedule" in {
+ val workflow = buildHeaderlessCsvKeywordWorkflow()
+ val scheduler = new WorkflowScheduler(workflow.context, CONTROLLER)
+ scheduler.updateSchedule(workflow.physicalPlan)
+
+ val expectedRegions = scheduler.getSchedule.getRegions.toSet
+ val pulledRegions = Iterator
+ .continually(scheduler.getNextRegions)
+ .takeWhile(_.nonEmpty)
+ .flatten
+ .toSet
+
+ assert(pulledRegions == expectedRegions)
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/execution/ExecutionUtilsSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/execution/ExecutionUtilsSpec.scala
new file mode 100644
index 00000000000..cf07c228438
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/execution/ExecutionUtilsSpec.scala
@@ -0,0 +1,340 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.worker.statistics.{
+ PortTupleMetricsMapping,
+ TupleMetrics
+}
+import org.apache.texera.amber.engine.common.executionruntimestate.{
+ OperatorMetrics,
+ OperatorStatistics
+}
+import org.scalatest.flatspec.AnyFlatSpec
+
+class ExecutionUtilsSpec extends AnyFlatSpec {
+
+ // Sentinel labels used as the generic T for ExecutionUtils.aggregateStates.
+ private val Completed = "completed"
+ private val Terminated = "terminated"
+ private val Running = "running"
+ private val Uninitialized = "uninitialized"
+ private val Paused = "paused"
+ private val Ready = "ready"
+
+ private def aggregate(states: String*): WorkflowAggregatedState =
+ ExecutionUtils.aggregateStates(
+ states,
+ Completed,
+ Terminated,
+ Running,
+ Uninitialized,
+ Paused,
+ Ready
+ )
+
+ "ExecutionUtils.aggregateStates" should "return UNINITIALIZED for an empty input" in {
+ assert(aggregate() == WorkflowAggregatedState.UNINITIALIZED)
+ }
+
+ it should "return COMPLETED when every state is the completed sentinel" in {
+ assert(aggregate(Completed, Completed) == WorkflowAggregatedState.COMPLETED)
+ }
+
+ it should "return COMPLETED when every state is the terminated sentinel" in {
+ assert(aggregate(Terminated, Terminated) == WorkflowAggregatedState.COMPLETED)
+ }
+
+ it should "return RUNNING when any state is the running sentinel" in {
+ assert(aggregate(Completed, Running, Paused) == WorkflowAggregatedState.RUNNING)
+ }
+
+ it should "return UNINITIALIZED when remaining non-completed states are all uninitialized" in {
+ assert(
+ aggregate(Completed, Uninitialized, Uninitialized) ==
+ WorkflowAggregatedState.UNINITIALIZED
+ )
+ }
+
+ it should "return PAUSED when remaining non-completed states are all paused" in {
+ assert(aggregate(Completed, Paused, Paused) == WorkflowAggregatedState.PAUSED)
+ }
+
+ it should "return RUNNING when remaining non-completed states are all ready" in {
+ // Note: an all-ready aggregate maps to RUNNING by current contract.
+ assert(aggregate(Completed, Ready, Ready) == WorkflowAggregatedState.RUNNING)
+ }
+
+ it should "return UNKNOWN when remaining non-completed states are mixed" in {
+ assert(aggregate(Completed, Paused, Ready) == WorkflowAggregatedState.UNKNOWN)
+ }
+
+ // Anti / boundary cases — make sure unexpected inputs cannot smuggle in a wrong
+ // state, and that branch precedence is what the contract claims.
+
+ it should "return UNKNOWN when completed and terminated are mixed (neither forall branch matches)" in {
+ // Both `forall(_ == completed)` and `forall(_ == terminated)` fail, no running
+ // sentinel is present, and the non-completed remainder is purely terminated —
+ // which is none of uninitialized / paused / ready, so the result must be
+ // UNKNOWN rather than COMPLETED.
+ assert(aggregate(Completed, Terminated) == WorkflowAggregatedState.UNKNOWN)
+ }
+
+ it should "give running precedence over completed and terminated" in {
+ assert(aggregate(Completed, Running) == WorkflowAggregatedState.RUNNING)
+ assert(aggregate(Terminated, Running) == WorkflowAggregatedState.RUNNING)
+ assert(aggregate(Running) == WorkflowAggregatedState.RUNNING)
+ }
+
+ it should "report PAUSED / UNINITIALIZED / RUNNING even when no completed sentinel is present" in {
+ assert(aggregate(Paused, Paused) == WorkflowAggregatedState.PAUSED)
+ assert(aggregate(Uninitialized, Uninitialized) == WorkflowAggregatedState.UNINITIALIZED)
+ // All-ready (no completed) maps to RUNNING, same as the with-completed case above.
+ assert(aggregate(Ready, Ready) == WorkflowAggregatedState.RUNNING)
+ }
+
+ it should "fall back to UNKNOWN when input contains values matching none of the sentinels" in {
+ // Defensive: a stray label that is not any of the six sentinels must not be
+ // silently classified as completed or running.
+ assert(aggregate("not-a-real-state") == WorkflowAggregatedState.UNKNOWN)
+ assert(aggregate(Completed, "not-a-real-state") == WorkflowAggregatedState.UNKNOWN)
+ }
+
+ // -- aggregatePortMetrics -----------------------------------------------
+
+ "ExecutionUtils.aggregatePortMetrics" should "return empty when given no mappings" in {
+ assert(ExecutionUtils.aggregatePortMetrics(Iterable.empty).isEmpty)
+ }
+
+ it should "preserve a single mapping" in {
+ val mapping = PortTupleMetricsMapping(PortIdentity(0), TupleMetrics(3, 30))
+ assert(ExecutionUtils.aggregatePortMetrics(List(mapping)) == Seq(mapping))
+ }
+
+ it should "sum count and size across mappings on the same port" in {
+ val portId = PortIdentity(0)
+ val a = PortTupleMetricsMapping(portId, TupleMetrics(3, 30))
+ val b = PortTupleMetricsMapping(portId, TupleMetrics(5, 50))
+ val result = ExecutionUtils.aggregatePortMetrics(List(a, b))
+ assert(result == Seq(PortTupleMetricsMapping(portId, TupleMetrics(8, 80))))
+ }
+
+ it should "group mappings by port id when ports differ" in {
+ val a = PortTupleMetricsMapping(PortIdentity(0), TupleMetrics(1, 10))
+ val b = PortTupleMetricsMapping(PortIdentity(1), TupleMetrics(2, 20))
+ val result = ExecutionUtils.aggregatePortMetrics(List(a, b)).toSet
+ assert(result == Set(a, b))
+ }
+
+ it should "sum more than two mappings on the same port without losing any" in {
+ val portId = PortIdentity(0)
+ val mappings = List(
+ PortTupleMetricsMapping(portId, TupleMetrics(1, 10)),
+ PortTupleMetricsMapping(portId, TupleMetrics(2, 20)),
+ PortTupleMetricsMapping(portId, TupleMetrics(4, 40))
+ )
+ assert(
+ ExecutionUtils.aggregatePortMetrics(mappings) ==
+ Seq(PortTupleMetricsMapping(portId, TupleMetrics(7, 70)))
+ )
+ }
+
+ it should "sum independently per port when multiple ports each have multiple mappings" in {
+ val port0 = PortIdentity(0)
+ val port1 = PortIdentity(1)
+ val mappings = List(
+ PortTupleMetricsMapping(port0, TupleMetrics(1, 10)),
+ PortTupleMetricsMapping(port1, TupleMetrics(3, 30)),
+ PortTupleMetricsMapping(port0, TupleMetrics(2, 20)),
+ PortTupleMetricsMapping(port1, TupleMetrics(4, 40))
+ )
+ val result = ExecutionUtils.aggregatePortMetrics(mappings).toSet
+ assert(
+ result == Set(
+ PortTupleMetricsMapping(port0, TupleMetrics(3, 30)),
+ PortTupleMetricsMapping(port1, TupleMetrics(7, 70))
+ )
+ )
+ }
+
+ it should "preserve a zero-count, zero-size mapping rather than dropping it" in {
+ val mapping = PortTupleMetricsMapping(PortIdentity(0), TupleMetrics(0, 0))
+ assert(ExecutionUtils.aggregatePortMetrics(List(mapping)) == Seq(mapping))
+ }
+
+ // -- aggregateMetrics ---------------------------------------------------
+
+ private def metricsWith(
+ state: WorkflowAggregatedState,
+ input: Seq[PortTupleMetricsMapping] = Seq.empty,
+ output: Seq[PortTupleMetricsMapping] = Seq.empty,
+ numWorkers: Int = 0,
+ dataTime: Long = 0,
+ controlTime: Long = 0,
+ idleTime: Long = 0
+ ): OperatorMetrics =
+ OperatorMetrics(
+ state,
+ OperatorStatistics(input, output, numWorkers, dataTime, controlTime, idleTime)
+ )
+
+ "ExecutionUtils.aggregateMetrics" should "return UNINITIALIZED defaults when given no metrics" in {
+ val result = ExecutionUtils.aggregateMetrics(Iterable.empty)
+ assert(result.operatorState == WorkflowAggregatedState.UNINITIALIZED)
+ assert(result.operatorStatistics.inputMetrics.isEmpty)
+ assert(result.operatorStatistics.outputMetrics.isEmpty)
+ assert(result.operatorStatistics.numWorkers == 0)
+ assert(result.operatorStatistics.dataProcessingTime == 0)
+ assert(result.operatorStatistics.controlProcessingTime == 0)
+ assert(result.operatorStatistics.idleTime == 0)
+ }
+
+ it should "sum scalar statistics and merge per-port metrics across operators" in {
+ val portIn = PortIdentity(0)
+ val portOut = PortIdentity(0)
+ val left = metricsWith(
+ WorkflowAggregatedState.RUNNING,
+ input = Seq(PortTupleMetricsMapping(portIn, TupleMetrics(2, 20))),
+ output = Seq(PortTupleMetricsMapping(portOut, TupleMetrics(1, 10))),
+ numWorkers = 1,
+ dataTime = 100,
+ controlTime = 5,
+ idleTime = 1
+ )
+ val right = metricsWith(
+ WorkflowAggregatedState.RUNNING,
+ input = Seq(PortTupleMetricsMapping(portIn, TupleMetrics(3, 30))),
+ output = Seq(PortTupleMetricsMapping(portOut, TupleMetrics(4, 40))),
+ numWorkers = 2,
+ dataTime = 200,
+ controlTime = 10,
+ idleTime = 2
+ )
+
+ val result = ExecutionUtils.aggregateMetrics(List(left, right))
+
+ assert(result.operatorState == WorkflowAggregatedState.RUNNING)
+ assert(
+ result.operatorStatistics.inputMetrics ==
+ Seq(PortTupleMetricsMapping(portIn, TupleMetrics(5, 50)))
+ )
+ assert(
+ result.operatorStatistics.outputMetrics ==
+ Seq(PortTupleMetricsMapping(portOut, TupleMetrics(5, 50)))
+ )
+ assert(result.operatorStatistics.numWorkers == 3)
+ assert(result.operatorStatistics.dataProcessingTime == 300)
+ assert(result.operatorStatistics.controlProcessingTime == 15)
+ assert(result.operatorStatistics.idleTime == 3)
+ }
+
+ it should "filter out internal ports when aggregating port metrics" in {
+ val publicPort = PortIdentity(0)
+ val internalPort = PortIdentity(1, internal = true)
+ val metrics = metricsWith(
+ WorkflowAggregatedState.RUNNING,
+ input = Seq(
+ PortTupleMetricsMapping(publicPort, TupleMetrics(1, 10)),
+ PortTupleMetricsMapping(internalPort, TupleMetrics(99, 990))
+ ),
+ output = Seq(PortTupleMetricsMapping(internalPort, TupleMetrics(7, 70)))
+ )
+
+ val result = ExecutionUtils.aggregateMetrics(List(metrics))
+
+ assert(
+ result.operatorStatistics.inputMetrics ==
+ Seq(PortTupleMetricsMapping(publicPort, TupleMetrics(1, 10)))
+ )
+ assert(result.operatorStatistics.outputMetrics.isEmpty)
+ }
+
+ it should "preserve a single operator's statistics (modulo internal-port filtering)" in {
+ val portIn = PortIdentity(0)
+ val portOut = PortIdentity(0)
+ val single = metricsWith(
+ WorkflowAggregatedState.RUNNING,
+ input = Seq(PortTupleMetricsMapping(portIn, TupleMetrics(2, 20))),
+ output = Seq(PortTupleMetricsMapping(portOut, TupleMetrics(3, 30))),
+ numWorkers = 4,
+ dataTime = 50,
+ controlTime = 6,
+ idleTime = 1
+ )
+
+ val result = ExecutionUtils.aggregateMetrics(List(single))
+
+ assert(result.operatorState == WorkflowAggregatedState.RUNNING)
+ assert(
+ result.operatorStatistics.inputMetrics ==
+ Seq(PortTupleMetricsMapping(portIn, TupleMetrics(2, 20)))
+ )
+ assert(
+ result.operatorStatistics.outputMetrics ==
+ Seq(PortTupleMetricsMapping(portOut, TupleMetrics(3, 30)))
+ )
+ assert(result.operatorStatistics.numWorkers == 4)
+ assert(result.operatorStatistics.dataProcessingTime == 50)
+ assert(result.operatorStatistics.controlProcessingTime == 6)
+ assert(result.operatorStatistics.idleTime == 1)
+ }
+
+ it should "report RUNNING when at least one operator is running and the rest are completed" in {
+ val running = metricsWith(WorkflowAggregatedState.RUNNING)
+ val completed = metricsWith(WorkflowAggregatedState.COMPLETED)
+
+ val result = ExecutionUtils.aggregateMetrics(List(running, completed))
+
+ assert(result.operatorState == WorkflowAggregatedState.RUNNING)
+ }
+
+ it should "report COMPLETED when every operator is completed" in {
+ val completedA = metricsWith(WorkflowAggregatedState.COMPLETED, numWorkers = 1)
+ val completedB = metricsWith(WorkflowAggregatedState.COMPLETED, numWorkers = 2)
+
+ val result = ExecutionUtils.aggregateMetrics(List(completedA, completedB))
+
+ assert(result.operatorState == WorkflowAggregatedState.COMPLETED)
+ assert(result.operatorStatistics.numWorkers == 3)
+ }
+
+ it should "tolerate operators with empty per-port stats while summing scalars" in {
+ val withStats = metricsWith(
+ WorkflowAggregatedState.RUNNING,
+ input = Seq(PortTupleMetricsMapping(PortIdentity(0), TupleMetrics(1, 10))),
+ numWorkers = 1,
+ dataTime = 5
+ )
+ val empty = metricsWith(WorkflowAggregatedState.RUNNING, numWorkers = 2, dataTime = 7)
+
+ val result = ExecutionUtils.aggregateMetrics(List(withStats, empty))
+
+ assert(result.operatorState == WorkflowAggregatedState.RUNNING)
+ assert(
+ result.operatorStatistics.inputMetrics ==
+ Seq(PortTupleMetricsMapping(PortIdentity(0), TupleMetrics(1, 10)))
+ )
+ assert(result.operatorStatistics.outputMetrics.isEmpty)
+ assert(result.operatorStatistics.numWorkers == 3)
+ assert(result.operatorStatistics.dataProcessingTime == 12)
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkflowExecutionSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkflowExecutionSpec.scala
new file mode 100644
index 00000000000..94285e7f8c5
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkflowExecutionSpec.scala
@@ -0,0 +1,167 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.controller.execution
+
+import org.apache.texera.amber.core.executor.OpExecInitInfo
+import org.apache.texera.amber.core.virtualidentity.{
+ ExecutionIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity,
+ WorkflowIdentity
+}
+import org.apache.texera.amber.core.workflow.PhysicalOp
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.architecture.scheduling.{Region, RegionIdentity}
+import org.scalatest.flatspec.AnyFlatSpec
+
+class WorkflowExecutionSpec extends AnyFlatSpec {
+
+ private def physicalOpId(opId: String): PhysicalOpIdentity =
+ PhysicalOpIdentity(OperatorIdentity(opId), "main")
+
+ private def op(opId: String): PhysicalOp =
+ PhysicalOp(
+ physicalOpId(opId),
+ WorkflowIdentity(0),
+ ExecutionIdentity(0),
+ OpExecInitInfo.Empty
+ )
+
+ /** A region with no ports — its `RegionExecution.getState` defaults to COMPLETED. */
+ private def region(regionId: Long, opId: String): Region =
+ Region(RegionIdentity(regionId), Set(op(opId)), Set.empty)
+
+ "WorkflowExecution.initRegionExecution" should "create a new RegionExecution for the given region" in {
+ val we = WorkflowExecution()
+ val r = region(1, "a")
+
+ val regionExecution = we.initRegionExecution(r)
+
+ assert(regionExecution.region == r)
+ assert(we.getRegionExecution(r.id) eq regionExecution)
+ }
+
+ it should "throw when called twice for the same region id" in {
+ val we = WorkflowExecution()
+ val r = region(1, "a")
+ we.initRegionExecution(r)
+
+ assertThrows[AssertionError] {
+ we.initRegionExecution(r)
+ }
+ }
+
+ "WorkflowExecution.hasRegionExecution" should "be false before init and true after" in {
+ val we = WorkflowExecution()
+ val r = region(1, "a")
+
+ assert(!we.hasRegionExecution(r.id))
+ we.initRegionExecution(r)
+ assert(we.hasRegionExecution(r.id))
+ }
+
+ "WorkflowExecution.getRegionExecution" should "throw NoSuchElementException for an unknown region id" in {
+ val we = WorkflowExecution()
+ assertThrows[NoSuchElementException] {
+ we.getRegionExecution(RegionIdentity(99))
+ }
+ }
+
+ "WorkflowExecution.getAllRegionExecutions" should "preserve the insertion order of region executions" in {
+ val we = WorkflowExecution()
+ val r0 = region(0, "a")
+ val r1 = region(1, "b")
+ val r2 = region(2, "c")
+
+ val e0 = we.initRegionExecution(r0)
+ val e1 = we.initRegionExecution(r1)
+ val e2 = we.initRegionExecution(r2)
+
+ assert(we.getAllRegionExecutions.toList == List(e0, e1, e2))
+ }
+
+ "WorkflowExecution.restartRegionExecution" should "behave like a fresh init when no prior region execution exists" in {
+ val we = WorkflowExecution()
+ val r = region(1, "a")
+
+ val regionExecution = we.restartRegionExecution(r)
+
+ assert(we.hasRegionExecution(r.id))
+ assert(we.getRegionExecution(r.id) eq regionExecution)
+ }
+
+ it should "replace an existing completed region execution with a fresh one" in {
+ val we = WorkflowExecution()
+ val r = region(1, "a")
+ val original = we.initRegionExecution(r)
+ assert(original.isCompleted)
+
+ val replacement = we.restartRegionExecution(r)
+
+ assert(replacement ne original)
+ assert(we.getRegionExecution(r.id) eq replacement)
+ }
+
+ "WorkflowExecution.getRunningRegionExecutions" should "exclude completed region executions" in {
+ val we = WorkflowExecution()
+ val r = region(1, "a")
+ val regionExecution = we.initRegionExecution(r)
+ assert(regionExecution.isCompleted)
+
+ assert(we.getRunningRegionExecutions.toList.isEmpty)
+ }
+
+ "WorkflowExecution.getState" should "return UNINITIALIZED when no regions have been initialized" in {
+ val we = WorkflowExecution()
+ assert(we.getState == WorkflowAggregatedState.UNINITIALIZED)
+ assert(!we.isCompleted)
+ }
+
+ it should "return COMPLETED when every initialized region is completed" in {
+ val we = WorkflowExecution()
+ we.initRegionExecution(region(0, "a"))
+ we.initRegionExecution(region(1, "b"))
+
+ assert(we.getState == WorkflowAggregatedState.COMPLETED)
+ assert(we.isCompleted)
+ }
+
+ "WorkflowExecution.getLatestOperatorExecutionOption" should "return None when no operator execution exists for the id" in {
+ val we = WorkflowExecution()
+ we.initRegionExecution(region(0, "a"))
+
+ assert(we.getLatestOperatorExecutionOption(physicalOpId("never-initialized")).isEmpty)
+ }
+
+ it should "return the latest matching operator execution across regions" in {
+ val we = WorkflowExecution()
+ val regionA = we.initRegionExecution(region(0, "a"))
+ val regionB = we.initRegionExecution(region(1, "b"))
+
+ val olderExecution = regionA.initOperatorExecution(physicalOpId("a"))
+ val newerExecution = regionB.initOperatorExecution(physicalOpId("a"))
+
+ val result = we.getLatestOperatorExecutionOption(physicalOpId("a"))
+ // Use reference identity: OperatorExecution is a no-field case class so
+ // instances are structurally equal; only `eq` distinguishes them.
+ assert(result.exists(_ eq newerExecution))
+ assert(!result.exists(_ eq olderExecution))
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkInputGatewaySpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkInputGatewaySpec.scala
new file mode 100644
index 00000000000..04f4f000455
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkInputGatewaySpec.scala
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.tuple.{AttributeType, Schema, TupleLike}
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.common.ambermessage.{DataFrame, WorkflowFIFOMessage}
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.flatspec.AnyFlatSpec
+
+class NetworkInputGatewaySpec extends AnyFlatSpec with MockFactory {
+
+ private val fakeReceiverID = ActorVirtualIdentity("testReceiver")
+ private val fakeSenderID = ActorVirtualIdentity("testSender")
+ private val channelId = ChannelIdentity(fakeSenderID, fakeReceiverID, isControl = false)
+ private val payloads = (0 until 4).map { i =>
+ DataFrame(
+ Array(
+ TupleLike(i) enforceSchema Schema().add("field1", AttributeType.INTEGER)
+ )
+ )
+ }.toArray
+ private val messages = (0 until 4).map { i =>
+ WorkflowFIFOMessage(channelId, i, payloads(i))
+ }.toArray
+
+ "network input port" should "output payload in FIFO order" in {
+ val inputPort = new NetworkInputGateway(fakeReceiverID)
+ Array(2, 0, 1, 3).foreach { i =>
+ inputPort.getChannel(channelId).acceptMessage(messages(i))
+ }
+
+ (0 until 4).foreach { i =>
+ val msg = inputPort.getChannel(channelId).take
+ assert(msg.sequenceNumber == i)
+ }
+
+ }
+
+ "network input port" should "de-duplicate payload" in {
+ val inputPort = new NetworkInputGateway(fakeReceiverID)
+ Array(2, 2, 2, 2, 2, 2, 0, 1, 1, 3, 3).foreach { i =>
+ inputPort.getChannel(channelId).acceptMessage(messages(i))
+ }
+ (0 until 4).foreach { i =>
+ val msg = inputPort.getChannel(channelId).take
+ assert(msg.sequenceNumber == i)
+ }
+
+ assert(!inputPort.getChannel(channelId).hasMessage)
+
+ }
+
+ "network input port" should "keep unordered messages" in {
+ val inputPort = new NetworkInputGateway(fakeReceiverID)
+ Array(3, 2, 1).foreach { i =>
+ inputPort.getChannel(channelId).acceptMessage(messages(i))
+ }
+ assert(!inputPort.getChannel(channelId).hasMessage)
+ inputPort.getChannel(channelId).acceptMessage(messages(0))
+ assert(inputPort.getChannel(channelId).hasMessage)
+ (0 until 4).foreach { i =>
+ val msg = inputPort.getChannel(channelId).take
+ assert(msg.sequenceNumber == i)
+ }
+ assert(!inputPort.getChannel(channelId).hasMessage)
+
+ }
+
+ "network input port" should "remove control channel by sender" in {
+ val inputPort = new NetworkInputGateway(fakeReceiverID)
+ val controlChannelId = ChannelIdentity(fakeSenderID, fakeReceiverID, isControl = true)
+ inputPort.getChannel(controlChannelId)
+
+ assert(inputPort.getAllControlChannels.size == 1)
+
+ inputPort.removeControlChannel(fakeSenderID)
+
+ assert(inputPort.getAllControlChannels.isEmpty)
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManagerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManagerSpec.scala
new file mode 100644
index 00000000000..405d9247068
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManagerSpec.scala
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import com.softwaremill.macwire.wire
+import org.apache.texera.amber.core.tuple.{AttributeType, Schema, TupleLike}
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity
+}
+import org.apache.texera.amber.core.workflow.{PhysicalLink, PortIdentity}
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.OneToOnePartitioning
+import org.apache.texera.amber.engine.common.ambermessage._
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.flatspec.AnyFlatSpec
+
+class OutputManagerSpec extends AnyFlatSpec with MockFactory {
+ private val mockHandler =
+ mock[WorkflowFIFOMessage => Unit]
+ private val identifier = ActorVirtualIdentity("batch producer mock")
+ private val mockDataOutputPort = // scalafix:ok; need it for wiring purpose
+ new NetworkOutputGateway(identifier, mockHandler)
+ var counter: Int = 0
+ val schema: Schema = Schema()
+ .add("field1", AttributeType.INTEGER)
+ .add("field2", AttributeType.INTEGER)
+ .add("field3", AttributeType.INTEGER)
+ .add("field4", AttributeType.INTEGER)
+ .add("field5", AttributeType.STRING)
+ .add("field6", AttributeType.DOUBLE)
+
+ def physicalOpId(): PhysicalOpIdentity = {
+ counter += 1
+ PhysicalOpIdentity(OperatorIdentity("" + counter), "" + counter)
+ }
+
+ def mkDataMessage(
+ to: ActorVirtualIdentity,
+ from: ActorVirtualIdentity,
+ seq: Long,
+ payload: DataPayload
+ ): WorkflowFIFOMessage = {
+ WorkflowFIFOMessage(ChannelIdentity(from, to, isControl = false), seq, payload)
+ }
+
+ "OutputManager" should "aggregate tuples and output" in {
+ val outputManager = wire[OutputManager]
+ val mockPortId = PortIdentity()
+ outputManager.addPort(mockPortId, schema, None)
+
+ val tuples = Array.fill(21)(
+ TupleLike(1, 2, 3, 4, "5", 9.8).enforceSchema(schema)
+ )
+ val fakeID = ActorVirtualIdentity("testReceiver")
+ inSequence {
+ (mockHandler.apply _).expects(
+ mkDataMessage(fakeID, identifier, 0, DataFrame(tuples.slice(0, 10)))
+ )
+ (mockHandler.apply _).expects(
+ mkDataMessage(fakeID, identifier, 1, DataFrame(tuples.slice(10, 20)))
+ )
+ (mockHandler.apply _).expects(
+ mkDataMessage(fakeID, identifier, 2, DataFrame(tuples.slice(20, 21)))
+ )
+ }
+ val fakeLink = PhysicalLink(physicalOpId(), mockPortId, physicalOpId(), mockPortId)
+ val fakeReceiver =
+ Array[ChannelIdentity](ChannelIdentity(identifier, fakeID, isControl = false))
+
+ outputManager.addPartitionerWithPartitioning(
+ fakeLink,
+ OneToOnePartitioning(10, fakeReceiver.toSeq)
+ )
+ tuples.foreach { t =>
+ outputManager.passTupleToDownstream(TupleLike(t.getFields).enforceSchema(schema), None)
+ }
+ outputManager.flush()
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/RangeBasedShuffleSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/RangeBasedShuffleSpec.scala
new file mode 100644
index 00000000000..2f906f59732
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/RangeBasedShuffleSpec.scala
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.messaginglayer
+
+import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema, Tuple}
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitioners.RangeBasedShufflePartitioner
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.RangeBasedShufflePartitioning
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.flatspec.AnyFlatSpec
+
+class RangeBasedShuffleSpec extends AnyFlatSpec with MockFactory {
+ val identifier = ActorVirtualIdentity("batch producer mock")
+ val fakeID1: ActorVirtualIdentity = ActorVirtualIdentity("rec1")
+ val fakeID2: ActorVirtualIdentity = ActorVirtualIdentity("rec2")
+ val fakeID3: ActorVirtualIdentity = ActorVirtualIdentity("rec3")
+ val fakeID4: ActorVirtualIdentity = ActorVirtualIdentity("rec4")
+ val fakeID5: ActorVirtualIdentity = ActorVirtualIdentity("rec5")
+
+ val attr: Attribute = new Attribute("Attr1", AttributeType.INTEGER)
+ val schema: Schema = Schema().add(attr)
+ val partitioning: RangeBasedShufflePartitioning =
+ RangeBasedShufflePartitioning(
+ 400,
+ List(
+ ChannelIdentity(identifier, fakeID1, isControl = false),
+ ChannelIdentity(identifier, fakeID2, isControl = false),
+ ChannelIdentity(identifier, fakeID3, isControl = false),
+ ChannelIdentity(identifier, fakeID4, isControl = false),
+ ChannelIdentity(identifier, fakeID5, isControl = false)
+ ),
+ Seq("Attr1"),
+ -400,
+ 600
+ )
+
+ val partitioner: RangeBasedShufflePartitioner = RangeBasedShufflePartitioner(partitioning)
+
+ "RangeBasedShuffleSpec" should "return 0 when value is less than rangeMin" in {
+ val tuple = Tuple.builder(schema).add(attr, -600).build()
+ val idx = partitioner.getBucketIndex(tuple)
+ assert(idx.next() == 0)
+ }
+
+ "RangeBasedShuffleSpec" should "return last receiver when value is more than rangeMax" in {
+ val tuple = Tuple.builder(schema).add(attr, 800).build()
+ val idx = partitioner.getBucketIndex(tuple)
+ assert(idx.next() == 4)
+ }
+
+ "RangeBasedShuffleSpec" should "find index correctly" in {
+ var tuple = Tuple.builder(schema).add(attr, -400).build()
+ var idx = partitioner.getBucketIndex(tuple)
+ assert(idx.next() == 0)
+
+ tuple = Tuple.builder(schema).add(attr, -200).build()
+ idx = partitioner.getBucketIndex(tuple)
+ assert(idx.next() == 0)
+
+ tuple = Tuple.builder(schema).add(attr, -199).build()
+ idx = partitioner.getBucketIndex(tuple)
+ assert(idx.next() == 1)
+ }
+
+ "RangeBasedShuffleSpec" should "handle different data types correctly" in {
+ var tuple = Tuple.builder(schema).add(attr, -90).build()
+ var idx = partitioner.getBucketIndex(tuple)
+ assert(idx.next() == 1)
+
+ val partitioning2: RangeBasedShufflePartitioning =
+ RangeBasedShufflePartitioning(
+ 400,
+ List(
+ ChannelIdentity(identifier, fakeID1, isControl = false),
+ ChannelIdentity(identifier, fakeID2, isControl = false),
+ ChannelIdentity(identifier, fakeID3, isControl = false),
+ ChannelIdentity(identifier, fakeID4, isControl = false),
+ ChannelIdentity(identifier, fakeID5, isControl = false)
+ ),
+ Seq("Attr2"),
+ -400,
+ 600
+ )
+
+ val partitioner2: RangeBasedShufflePartitioner = RangeBasedShufflePartitioner(partitioning2)
+ val doubleAttr: Attribute = new Attribute("Attr2", AttributeType.DOUBLE)
+ val doubleSchema: Schema = Schema().add(doubleAttr)
+ tuple = Tuple.builder(doubleSchema).add(doubleAttr, -90.5).build()
+ idx = partitioner2.getBucketIndex(tuple)
+ assert(idx.next() == 1)
+
+ val partitioning3: RangeBasedShufflePartitioning =
+ RangeBasedShufflePartitioning(
+ 400,
+ List(
+ ChannelIdentity(identifier, fakeID1, isControl = false),
+ ChannelIdentity(identifier, fakeID2, isControl = false),
+ ChannelIdentity(identifier, fakeID3, isControl = false),
+ ChannelIdentity(identifier, fakeID4, isControl = false),
+ ChannelIdentity(identifier, fakeID5, isControl = false)
+ ),
+ Seq("Attr3"),
+ -400,
+ 600
+ )
+
+ val partitioner3: RangeBasedShufflePartitioner = RangeBasedShufflePartitioner(partitioning3)
+ val longAttr: Attribute = new Attribute("Attr3", AttributeType.LONG)
+ val longSchema: Schema = Schema().add(longAttr)
+ tuple = Tuple.builder(longSchema).add(longAttr, -90L).build()
+ idx = partitioner3.getBucketIndex(tuple)
+ assert(idx.next() == 1)
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala
new file mode 100644
index 00000000000..1ca0572dca6
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+//package org.apache.texera.amber.engine.architecture.pythonworker
+//
+//import org.apache.pekko.actor.{ActorRef, ActorSystem, Props}
+//import org.apache.pekko.testkit.{ImplicitSender, TestActorRef, TestKit}
+//import org.apache.texera.amber.clustering.SingleNodeListener
+//import org.apache.texera.amber.engine.architecture.common.WorkflowActor.{NetworkAck, NetworkMessage}
+//import org.apache.texera.amber.engine.architecture.pythonworker.promisehandlers.InitializeOperatorLogicHandler.InitializeOperatorLogic
+//import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.OneToOnePartitioning
+//import org.apache.texera.amber.engine.architecture.worker.controlcommands.LinkOrdinal
+//import org.apache.texera.amber.engine.architecture.worker.promisehandlers.AddPartitioningHandler.AddPartitioning
+//import org.apache.texera.amber.engine.architecture.worker.promisehandlers.OpenOperatorHandler.OpenOperator
+//import org.apache.texera.amber.engine.architecture.worker.promisehandlers.UpdateInputLinkingHandler.UpdateInputLinking
+//import org.apache.texera.amber.engine.common.Constants
+//import org.apache.texera.amber.engine.common.ambermessage.{
+// ChannelIdentity,
+// ControlPayload,
+// DataFrame,
+// DataPayload,
+// EndOfUpstream,
+// WorkflowFIFOMessage
+//}
+//import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.{ControlInvocation, ReturnInvocation}
+//import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+//import org.apache.texera.amber.core.virtualidentity.{
+// ActorVirtualIdentity,
+// PhysicalLink,
+// PhysicalLink,
+// OperatorIdentity
+//}
+//import org.apache.texera.amber.engine.e2e.TestOperators
+//import org.apache.texera.workflow.common.tuple.Tuple
+//import org.apache.texera.workflow.common.tuple.schema.{Attribute, AttributeType, Schema}
+//import org.scalamock.scalatest.MockFactory
+//import org.scalatest.BeforeAndAfterAll
+//import org.scalatest.flatspec.AnyFlatSpecLike
+//
+//import scala.concurrent.duration.DurationInt
+//
+//class PythonWorkflowWorkerSpec
+// extends TestKit(ActorSystem("PythonWorkerSpec"))
+// with ImplicitSender
+// with AnyFlatSpecLike
+// with BeforeAndAfterAll
+// with MockFactory {
+//
+// override def beforeAll: Unit = {
+// system.actorOf(Props[SingleNodeListener], "cluster-info")
+// }
+// override def afterAll: Unit = {
+// TestKit.shutdownActorSystem(system)
+// }
+// private val identifier1 = ActorVirtualIdentity("worker-1")
+// private val identifier2 = ActorVirtualIdentity("worker-2")
+// private val operatorIdentity = OperatorIdentity("testWorkflow", "testOperator")
+// private val layerId1 =
+// PhysicalLink(operatorIdentity.workflow, operatorIdentity.operator, "1st-layer")
+// private val layerId2 =
+// PhysicalLink(operatorIdentity.workflow, operatorIdentity.operator, "2nd-layer")
+// private val pythonOp = TestOperators.pythonOpDesc()
+// private val link = PhysicalLink(layerId1, 0, layerId2, 0)
+// private val schema = Schema
+// .newBuilder()
+// .add(new Attribute("text", AttributeType.STRING))
+// .build()
+// private val initialization = InitializeOperatorLogic(
+// pythonOp.code,
+// isSource = false,
+// Seq(LinkOrdinal(link, 0)),
+// Seq(LinkOrdinal(link, 0)),
+// schema
+// )
+//
+// def sendControlToWorker(
+// worker: ActorRef,
+// controls: Array[ControlInvocation],
+// beginSeqNum: Long = 0
+// ): Unit = {
+// var seq = beginSeqNum
+// controls.foreach { ctrl =>
+// worker ! NetworkMessage(
+// seq,
+// WorkflowFIFOMessage(ChannelIdentity(CONTROLLER, identifier1, true), seq, ctrl)
+// )
+// val received = receiveWhile(3.seconds) {
+// case NetworkAck(id, credits) =>
+// // pass
+// case NetworkMessage(id, fifoPayload) =>
+// fifoPayload.payload.asInstanceOf[ControlPayload] match {
+// case ControlInvocation(commandID, command) => assert(commandID == seq)
+// case ReturnInvocation(originalCommandID, controlReturn) =>
+// assert(originalCommandID == seq)
+// case _ => ???
+// }
+// worker ! NetworkAck(id, Constants.unprocessedBatchesSizeLimitInBytesPerWorkerPair)
+// }
+// seq += 1
+// }
+// }
+//
+// def mkWorker: ActorRef = TestActorRef(new PythonWorkflowWorker(identifier1))
+//
+// "python worker" should "start" in {
+// val worker = mkWorker
+// sendControlToWorker(worker, Array(ControlInvocation(0, initialization)))
+// }
+//
+// "python worker" should "process data" in {
+// val worker = mkWorker
+// sendControlToWorker(worker, Array(ControlInvocation(0, initialization)))
+// val mockPolicy = OneToOnePartitioning(1, Array(identifier2))
+// val openControl = ControlInvocation(1, OpenOperator())
+// val invocation = ControlInvocation(2, AddPartitioning(link, mockPolicy))
+// val updateInputLinking = ControlInvocation(3, UpdateInputLinking(identifier2, link))
+// sendControlToWorker(worker, Array(openControl, invocation, updateInputLinking), 1)
+// worker ! NetworkMessage(
+// 4,
+// WorkflowFIFOMessage(
+// ChannelIdentity(identifier2, identifier1, false),
+// 0,
+// DataFrame(
+// Array(
+// Tuple
+// .newBuilder(schema)
+// .add("text", AttributeType.STRING, "123")
+// .build()
+// )
+// )
+// )
+// )
+// expectMsgClass(classOf[NetworkAck])
+// val data = receiveOne(30.seconds)
+// assert(data.asInstanceOf[NetworkMessage].internalMessage.payload.isInstanceOf[DataFrame])
+// }
+//
+// "python worker" should "process data and receive end marker" in {
+// val worker = mkWorker
+// sendControlToWorker(worker, Array(ControlInvocation(0, initialization)))
+// val mockPolicy = OneToOnePartitioning(100, Array(identifier2))
+// val openControl = ControlInvocation(1, OpenOperator())
+// val invocation = ControlInvocation(2, AddPartitioning(link, mockPolicy))
+// val updateInputLinking = ControlInvocation(3, UpdateInputLinking(identifier2, link))
+// sendControlToWorker(worker, Array(openControl, invocation, updateInputLinking), 1)
+// worker ! NetworkMessage(
+// 4,
+// WorkflowFIFOMessage(
+// ChannelIdentity(identifier2, identifier1, false),
+// 0,
+// DataFrame(
+// (0 until 100)
+// .map(_ =>
+// Tuple
+// .newBuilder(schema)
+// .add("text", AttributeType.STRING, "123")
+// .build()
+// )
+// .toArray
+// )
+// )
+// )
+// expectMsgClass(classOf[NetworkAck])
+// val data = receiveOne(30.seconds)
+// assert(data.asInstanceOf[NetworkMessage].internalMessage.payload.isInstanceOf[DataFrame])
+// worker ! NetworkMessage(
+// 5,
+// WorkflowFIFOMessage(
+// ChannelIdentity(identifier2, identifier1, false),
+// 1,
+// EndOfUpstream()
+// )
+// )
+// expectMsgClass(classOf[NetworkAck])
+// receiveWhile(10.seconds) {
+// case NetworkMessage(id, fifoPayload) =>
+// fifoPayload.payload match {
+// case payload: ControlPayload => //skip
+// case payload: DataPayload => assert(payload.isInstanceOf[EndOfUpstream])
+// case _ => ???
+// }
+// }
+// }
+//
+//}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/CostBasedScheduleGeneratorSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/CostBasedScheduleGeneratorSpec.scala
new file mode 100644
index 00000000000..7d5227c36bc
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/CostBasedScheduleGeneratorSpec.scala
@@ -0,0 +1,518 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.core.workflow.{
+ ExecutionMode,
+ PortIdentity,
+ WorkflowContext,
+ WorkflowSettings
+}
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.engine.e2e.TestUtils.buildWorkflow
+import org.apache.texera.amber.operator.TestOperators
+import org.apache.texera.workflow.LogicalLink
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.flatspec.AnyFlatSpec
+
+import scala.jdk.CollectionConverters._
+
+class CostBasedScheduleGeneratorSpec extends AnyFlatSpec with MockFactory {
+
+ "CostBasedRegionPlanGenerator" should "finish bottom-up search using different pruning techniques with correct number of states explored in csv->->filter->join->filter2 workflow" in {
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+ val keywordOpDesc2 = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ keywordOpDesc,
+ joinOpDesc,
+ keywordOpDesc2
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ ),
+ LogicalLink(
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc2.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val globalSearchNoPruningResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).bottomUpSearch(globalSearch = true, oChains = false, oCleanEdges = false, oEarlyStop = false)
+
+ // Should have explored all possible states (2^4 states)
+ assert(globalSearchNoPruningResult.numStatesExplored == 16)
+
+ val globalSearchOChainsResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).bottomUpSearch(globalSearch = true, oCleanEdges = false, oEarlyStop = false)
+
+ // By applying pruning based on Chains alone, it should skip 10 (8 + 2) states. 8 states where CSV->Build is
+ // materialized should be skipped because this edge is in the same chain as another blocking edge.
+ // Of the remaining states, 2 more states where both CSV->KeywordFilter and KeywordFilter->Probe are materialized
+ // should be skipped because these two edges are in the same chain.
+ assert(globalSearchOChainsResult.numStatesExplored == 6)
+
+ val globalSearchOCleanEdgesResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).bottomUpSearch(globalSearch = true, oChains = false, oEarlyStop = false)
+
+ // By applying pruning based on Clean edges (bridges) alone, it should skip 8 states. There is one clean edge
+ // in the DAG (Probe->Keyword2) and the 8 states where this edge is materialized should be skipped.
+ assert(globalSearchOCleanEdgesResult.numStatesExplored == 8)
+
+ val globalSearchOEarlyStopResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).bottomUpSearch(globalSearch = true, oChains = false, oCleanEdges = false)
+
+ // By applying pruning based on Early Stop alone, only 6 states that are not descendants of a schedulable states
+ // should be explored.
+ assert(globalSearchOEarlyStopResult.numStatesExplored == 6)
+
+ val globalSearchAllPruningEnabledResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).bottomUpSearch(globalSearch = true)
+
+ // By combining all pruning techniques, only 3 states should be visited (1 state where both CSV->KeywordFilter and
+ // KeywordFilter->Probe are pipelined, and two states where only one of CSV->KeywordFilter or KeywordFilter->Probe
+ // is materialized. The other two edges should always be pipelined.)
+ assert(globalSearchAllPruningEnabledResult.numStatesExplored == 3)
+
+ }
+
+ "CostBasedRegionPlanGenerator" should "finish top-down search using different pruning techniques with correct number of states explored in csv->->filter->join->filter2 workflow" in {
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+ val keywordOpDesc2 = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ keywordOpDesc,
+ joinOpDesc,
+ keywordOpDesc2
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ ),
+ LogicalLink(
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc2.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val globalSearchNoPruningResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).topDownSearch(globalSearch = true, oChains = false, oCleanEdges = false)
+
+ // Should have explored all possible states (2^4 states)
+ assert(globalSearchNoPruningResult.numStatesExplored == 16)
+
+ val globalSearchOChainsResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).topDownSearch(globalSearch = true, oCleanEdges = false)
+
+ // By applying pruning based on Chains alone, it should start with a state where CSV->Build is pipelined because
+ // this edge is in the same chain as another blocking edge. That reduces the search space to 8 states.
+ assert(globalSearchOChainsResult.numStatesExplored == 8)
+
+ val globalSearchOCleanEdgesResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).topDownSearch(globalSearch = true, oChains = false)
+
+ // By applying pruning based on Clean Edges (bridges) alone, it should start with a state where Probe->Keyword2 is
+ // pipelined because this edge is a clean edge. That reduces the search space to 8 states.
+ assert(globalSearchOCleanEdgesResult.numStatesExplored == 8)
+
+ val globalSearchAllPruningEnabledResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).topDownSearch(globalSearch = true)
+
+ // By combining both pruning techniques, the search should start with a state where both CSV->Build and
+ // Probe->Keyword2 are pipelined, reducing the search space to 4 states.
+ assert(globalSearchAllPruningEnabledResult.numStatesExplored == 4)
+
+ }
+
+ // MATERIALIZED ExecutionMode tests - each operator should be a separate region
+ "CostBasedRegionPlanGenerator" should "create separate region for each operator in MATERIALIZED mode for simple csv workflow" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val materializedContext = new WorkflowContext(
+ workflowSettings = WorkflowSettings(
+ dataTransferBatchSize = 400,
+ executionMode = ExecutionMode.MATERIALIZED
+ )
+ )
+ val workflow = buildWorkflow(
+ List(csvOpDesc),
+ List(),
+ materializedContext
+ )
+
+ val scheduleGenerator = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ )
+ val result = scheduleGenerator.getFullyMaterializedSearchState
+
+ // Should only explore 1 state (fully materialized)
+ assert(result.numStatesExplored == 1)
+
+ // Each physical operator should be in its own region
+ val regions = result.regionDAG.vertexSet().asScala
+ val numPhysicalOps = workflow.physicalPlan.operators.size
+ assert(regions.size == numPhysicalOps, s"Expected $numPhysicalOps regions, got ${regions.size}")
+
+ // Each region should contain exactly 1 operator
+ regions.foreach { region =>
+ assert(
+ region.getOperators.size == 1,
+ s"Expected region to have 1 operator, got ${region.getOperators.size}"
+ )
+ }
+ }
+
+ "CostBasedRegionPlanGenerator" should "create separate region for each operator in MATERIALIZED mode for csv->keyword workflow" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val materializedContext = new WorkflowContext(
+ workflowSettings = WorkflowSettings(
+ dataTransferBatchSize = 400,
+ executionMode = ExecutionMode.MATERIALIZED
+ )
+ )
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ materializedContext
+ )
+
+ val scheduleGenerator = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ )
+ val result = scheduleGenerator.getFullyMaterializedSearchState
+
+ // Should only explore 1 state (fully materialized)
+ assert(result.numStatesExplored == 1)
+
+ // Each physical operator should be in its own region
+ val regions = result.regionDAG.vertexSet().asScala
+ val numPhysicalOps = workflow.physicalPlan.operators.size
+ assert(regions.size == numPhysicalOps, s"Expected $numPhysicalOps regions, got ${regions.size}")
+
+ // Each region should contain exactly 1 operator
+ regions.foreach { region =>
+ assert(
+ region.getOperators.size == 1,
+ s"Expected region to have 1 operator, got ${region.getOperators.size}"
+ )
+ }
+
+ // All links should be materialized (represented as region links)
+ val numRegionLinks = result.regionDAG.edgeSet().asScala.size
+ val numPhysicalLinks = workflow.physicalPlan.links.size
+ assert(
+ numRegionLinks == numPhysicalLinks,
+ s"Expected $numPhysicalLinks region links, got $numRegionLinks"
+ )
+ }
+
+ "CostBasedRegionPlanGenerator" should "create separate region for each operator in MATERIALIZED mode for csv->keyword->count workflow" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val countOpDesc = TestOperators.aggregateAndGroupByDesc(
+ "Region",
+ org.apache.texera.amber.operator.aggregate.AggregationFunction.COUNT,
+ List[String]()
+ )
+ val materializedContext = new WorkflowContext(
+ workflowSettings = WorkflowSettings(
+ dataTransferBatchSize = 400,
+ executionMode = ExecutionMode.MATERIALIZED
+ )
+ )
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc, countOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ countOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ materializedContext
+ )
+
+ val scheduleGenerator = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ )
+ val result = scheduleGenerator.getFullyMaterializedSearchState
+
+ // Should only explore 1 state (fully materialized)
+ assert(result.numStatesExplored == 1)
+
+ // Each physical operator should be in its own region
+ val regions = result.regionDAG.vertexSet().asScala
+ val numPhysicalOps = workflow.physicalPlan.operators.size
+ assert(regions.size == numPhysicalOps, s"Expected $numPhysicalOps regions, got ${regions.size}")
+
+ // Each region should contain exactly 1 operator
+ regions.foreach { region =>
+ assert(
+ region.getOperators.size == 1,
+ s"Expected region to have 1 operator, got ${region.getOperators.size}"
+ )
+ }
+
+ // All links should be materialized (represented as region links)
+ val numRegionLinks = result.regionDAG.edgeSet().asScala.size
+ val numPhysicalLinks = workflow.physicalPlan.links.size
+ assert(
+ numRegionLinks == numPhysicalLinks,
+ s"Expected $numPhysicalLinks region links, got $numRegionLinks"
+ )
+ }
+
+ "CostBasedRegionPlanGenerator" should "create separate region for each operator in MATERIALIZED mode for join workflow" in {
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val headerlessCsvOpDesc2 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+ val materializedContext = new WorkflowContext(
+ workflowSettings = WorkflowSettings(
+ dataTransferBatchSize = 400,
+ executionMode = ExecutionMode.MATERIALIZED
+ )
+ )
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ headerlessCsvOpDesc2,
+ joinOpDesc
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc2.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ )
+ ),
+ materializedContext
+ )
+
+ val scheduleGenerator = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ )
+ val result = scheduleGenerator.getFullyMaterializedSearchState
+
+ // Should only explore 1 state (fully materialized)
+ assert(result.numStatesExplored == 1)
+
+ // Each physical operator should be in its own region
+ val regions = result.regionDAG.vertexSet().asScala
+ val numPhysicalOps = workflow.physicalPlan.operators.size
+ assert(regions.size == numPhysicalOps, s"Expected $numPhysicalOps regions, got ${regions.size}")
+
+ // Each region should contain exactly 1 operator
+ regions.foreach { region =>
+ assert(
+ region.getOperators.size == 1,
+ s"Expected region to have 1 operator, got ${region.getOperators.size}"
+ )
+ }
+
+ // All links should be materialized (represented as region links)
+ val numRegionLinks = result.regionDAG.edgeSet().asScala.size
+ val numPhysicalLinks = workflow.physicalPlan.links.size
+ assert(
+ numRegionLinks == numPhysicalLinks,
+ s"Expected $numPhysicalLinks region links, got $numRegionLinks"
+ )
+ }
+
+ "CostBasedRegionPlanGenerator" should "create separate region for each operator in MATERIALIZED mode for complex csv->->filter->join->filter2 workflow" in {
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+ val keywordOpDesc2 = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val materializedContext = new WorkflowContext(
+ workflowSettings = WorkflowSettings(
+ dataTransferBatchSize = 400,
+ executionMode = ExecutionMode.MATERIALIZED
+ )
+ )
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ keywordOpDesc,
+ joinOpDesc,
+ keywordOpDesc2
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ ),
+ LogicalLink(
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc2.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ materializedContext
+ )
+
+ val scheduleGenerator = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ )
+ val result = scheduleGenerator.getFullyMaterializedSearchState
+
+ // Should only explore 1 state (fully materialized)
+ assert(result.numStatesExplored == 1)
+
+ // Each physical operator should be in its own region
+ val regions = result.regionDAG.vertexSet().asScala
+ val numPhysicalOps = workflow.physicalPlan.operators.size
+ assert(regions.size == numPhysicalOps, s"Expected $numPhysicalOps regions, got ${regions.size}")
+
+ // Each region should contain exactly 1 operator
+ regions.foreach { region =>
+ assert(
+ region.getOperators.size == 1,
+ s"Expected region to have 1 operator, got ${region.getOperators.size}"
+ )
+ }
+
+ // All links should be materialized (represented as region links)
+ val numRegionLinks = result.regionDAG.edgeSet().asScala.size
+ val numPhysicalLinks = workflow.physicalPlan.links.size
+ assert(
+ numRegionLinks == numPhysicalLinks,
+ s"Expected $numPhysicalLinks region links, got $numRegionLinks"
+ )
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/DefaultCostEstimatorSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/DefaultCostEstimatorSpec.scala
new file mode 100644
index 00000000000..d1b2595cbb5
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/DefaultCostEstimatorSpec.scala
@@ -0,0 +1,393 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.core.storage.model.{BufferedItemWriter, VirtualDocument}
+import org.apache.texera.amber.core.storage.result.ResultSchema
+import org.apache.texera.amber.core.storage.{DocumentFactory, VFSURIFactory}
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity}
+import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PortIdentity, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.scheduling.resourcePolicies.{
+ DefaultResourceAllocator,
+ ExecutionClusterInfo
+}
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.engine.e2e.TestUtils.buildWorkflow
+import org.apache.texera.amber.operator.TestOperators
+import org.apache.texera.amber.operator.aggregate.{AggregateOpDesc, AggregationFunction}
+import org.apache.texera.amber.operator.keywordSearch.KeywordSearchOpDesc
+import org.apache.texera.amber.operator.source.scan.csv.CSVScanSourceOpDesc
+import org.apache.texera.dao.MockTexeraDB
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.daos._
+import org.apache.texera.dao.jooq.generated.tables.pojos._
+import org.apache.texera.workflow.LogicalLink
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
+
+import java.net.URI
+import java.sql.Timestamp
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+class DefaultCostEstimatorSpec
+ extends AnyFlatSpec
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with MockTexeraDB {
+
+ private val headerlessCsvOpDesc: CSVScanSourceOpDesc =
+ TestOperators.headerlessSmallCsvScanOpDesc()
+ private val keywordOpDesc: KeywordSearchOpDesc =
+ TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ private val groupByOpDesc: AggregateOpDesc =
+ TestOperators.aggregateAndGroupByDesc("column-1", AggregationFunction.COUNT, List[String]())
+
+ private val testUser: User = {
+ val user = new User
+ user.setUid(Integer.valueOf(1))
+ user.setName("test_user")
+ user.setRole(UserRoleEnum.ADMIN)
+ user.setPassword("123")
+ user.setEmail("test_user@test.com")
+ user
+ }
+
+ private val testWorkflowEntry: Workflow = {
+ val workflow = new Workflow
+ workflow.setName("test workflow")
+ workflow.setWid(Integer.valueOf(1))
+ workflow.setContent("test workflow content")
+ workflow.setDescription("test description")
+ workflow
+ }
+
+ private val testWorkflowVersionEntry: WorkflowVersion = {
+ val workflowVersion = new WorkflowVersion
+ workflowVersion.setWid(Integer.valueOf(1))
+ workflowVersion.setVid(Integer.valueOf(1))
+ workflowVersion.setContent("test version content")
+ workflowVersion
+ }
+
+ private val testWorkflowExecutionEntry: WorkflowExecutions = {
+ val workflowExecution = new WorkflowExecutions
+ workflowExecution.setEid(Integer.valueOf(1))
+ workflowExecution.setVid(Integer.valueOf(1))
+ workflowExecution.setUid(Integer.valueOf(1))
+ workflowExecution.setStatus(3.toByte)
+ workflowExecution.setEnvironmentVersion("test engine")
+ workflowExecution
+ }
+
+ private var uri: URI = _
+ private var writer: BufferedItemWriter[Tuple] = _
+ private var document: VirtualDocument[_] = _
+
+ override protected def beforeEach(): Unit = {
+ initializeDBAndReplaceDSLContext()
+ uri = VFSURIFactory.createRuntimeStatisticsURI(
+ WorkflowIdentity(testWorkflowEntry.getWid.longValue()),
+ ExecutionIdentity(testWorkflowExecutionEntry.getEid.longValue())
+ )
+ document = DocumentFactory.createDocument(uri, ResultSchema.runtimeStatisticsSchema)
+ writer = document
+ .writer(s"runtime_statistics_${testWorkflowExecutionEntry.getEid.longValue()}")
+ .asInstanceOf[BufferedItemWriter[Tuple]]
+ writer.open()
+ }
+
+ override protected def afterEach(): Unit = {
+ document.clear()
+ shutdownDB()
+ }
+
+ "DefaultCostEstimator" should "use fallback method when no past statistics are available" in {
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc.operatorIdentifier,
+ PortIdentity(0),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(0)
+ )
+ ),
+ new WorkflowContext()
+ )
+ val resourceAllocator =
+ new DefaultResourceAllocator(
+ workflow.physicalPlan,
+ new ExecutionClusterInfo(),
+ workflow.context.workflowSettings
+ )
+
+ val costEstimator = new DefaultCostEstimator(
+ workflow.context,
+ resourceAllocator,
+ CONTROLLER
+ )
+ val ports = workflow.physicalPlan.operators.flatMap(op =>
+ op.inputPorts.keys
+ .map(inputPortId => GlobalPortIdentity(op.id, inputPortId, input = true))
+ .toSet ++ op.outputPorts.keys
+ .map(outputPortId => GlobalPortIdentity(op.id, outputPortId))
+ .toSet
+ )
+
+ val region = Region(
+ id = RegionIdentity(0),
+ physicalOps = workflow.physicalPlan.operators,
+ physicalLinks = workflow.physicalPlan.links,
+ ports = ports
+ )
+
+ val (_, costOfRegion) = costEstimator.allocateResourcesAndEstimateCost(region, 1)
+
+ assert(costOfRegion == 0)
+ }
+
+ "DefaultCostEstimator" should "use the latest successful execution to estimate cost when available" in {
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc.operatorIdentifier,
+ PortIdentity(0),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(0)
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val userDao = new UserDao(getDSLContext.configuration())
+ val workflowDao = new WorkflowDao(getDSLContext.configuration())
+ val workflowExecutionsDao = new WorkflowExecutionsDao(getDSLContext.configuration())
+ val workflowVersionDao = new WorkflowVersionDao(getDSLContext.configuration())
+
+ userDao.insert(testUser)
+ workflowDao.insert(testWorkflowEntry)
+ workflowVersionDao.insert(testWorkflowVersionEntry)
+ testWorkflowExecutionEntry.setRuntimeStatsUri(uri.toString)
+ workflowExecutionsDao.insert(testWorkflowExecutionEntry)
+
+ val headerlessCsvOpRuntimeStatistics = new Tuple(
+ ResultSchema.runtimeStatisticsSchema,
+ Array(
+ headerlessCsvOpDesc.operatorIdentifier.id,
+ new Timestamp(System.currentTimeMillis()),
+ 0L,
+ 0L,
+ 0L,
+ 0L,
+ 100L,
+ 100L,
+ 0L,
+ 1,
+ 0
+ )
+ )
+ val keywordOpRuntimeStatistics = new Tuple(
+ ResultSchema.runtimeStatisticsSchema,
+ Array(
+ keywordOpDesc.operatorIdentifier.id,
+ new Timestamp(System.currentTimeMillis()),
+ 0L,
+ 0L,
+ 0L,
+ 0L,
+ 300L,
+ 300L,
+ 0L,
+ 1,
+ 0
+ )
+ )
+
+ writer.putOne(headerlessCsvOpRuntimeStatistics)
+ writer.putOne(keywordOpRuntimeStatistics)
+ writer.close()
+
+ val resourceAllocator =
+ new DefaultResourceAllocator(
+ workflow.physicalPlan,
+ new ExecutionClusterInfo(),
+ workflow.context.workflowSettings
+ )
+
+ val costEstimator = new DefaultCostEstimator(
+ workflow.context,
+ resourceAllocator,
+ CONTROLLER
+ )
+
+ val ports = workflow.physicalPlan.operators.flatMap(op =>
+ op.inputPorts.keys
+ .map(inputPortId => GlobalPortIdentity(op.id, inputPortId, input = true))
+ .toSet ++ op.outputPorts.keys
+ .map(outputPortId => GlobalPortIdentity(op.id, outputPortId))
+ .toSet
+ )
+
+ val region = Region(
+ id = RegionIdentity(0),
+ physicalOps = workflow.physicalPlan.operators,
+ physicalLinks = workflow.physicalPlan.links,
+ ports = ports
+ )
+
+ val (_, costOfRegion) = costEstimator.allocateResourcesAndEstimateCost(region, 1)
+
+ assert(costOfRegion != 0)
+ }
+
+ "DefaultCostEstimator" should "use correctly estimate costs in a search" in {
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc, groupByOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc.operatorIdentifier,
+ PortIdentity(0),
+ groupByOpDesc.operatorIdentifier,
+ PortIdentity(0)
+ ),
+ LogicalLink(
+ groupByOpDesc.operatorIdentifier,
+ PortIdentity(0),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(0)
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val userDao = new UserDao(getDSLContext.configuration())
+ val workflowDao = new WorkflowDao(getDSLContext.configuration())
+ val workflowExecutionsDao = new WorkflowExecutionsDao(getDSLContext.configuration())
+ val workflowVersionDao = new WorkflowVersionDao(getDSLContext.configuration())
+
+ userDao.insert(testUser)
+ workflowDao.insert(testWorkflowEntry)
+ workflowVersionDao.insert(testWorkflowVersionEntry)
+ testWorkflowExecutionEntry.setRuntimeStatsUri(uri.toString)
+ workflowExecutionsDao.insert(testWorkflowExecutionEntry)
+
+ val headerlessCsvOpRuntimeStatistics = new Tuple(
+ ResultSchema.runtimeStatisticsSchema,
+ Array(
+ headerlessCsvOpDesc.operatorIdentifier.id,
+ new Timestamp(System.currentTimeMillis()),
+ 0L,
+ 0L,
+ 0L,
+ 0L,
+ 100L,
+ 100L,
+ 0L,
+ 1,
+ 0
+ )
+ )
+ val groupByOpRuntimeStatistics = new Tuple(
+ ResultSchema.runtimeStatisticsSchema,
+ Array(
+ groupByOpDesc.operatorIdentifier.id,
+ new Timestamp(System.currentTimeMillis()),
+ 0L,
+ 0L,
+ 0L,
+ 0L,
+ 1000L,
+ 1000L,
+ 0L,
+ 1,
+ 0
+ )
+ )
+ val keywordOpRuntimeStatistics = new Tuple(
+ ResultSchema.runtimeStatisticsSchema,
+ Array(
+ keywordOpDesc.operatorIdentifier.id,
+ new Timestamp(System.currentTimeMillis()),
+ 0L,
+ 0L,
+ 0L,
+ 0L,
+ 300L,
+ 300L,
+ 0L,
+ 1,
+ 0
+ )
+ )
+
+ writer.putOne(headerlessCsvOpRuntimeStatistics)
+ writer.putOne(groupByOpRuntimeStatistics)
+ writer.putOne(keywordOpRuntimeStatistics)
+ writer.close()
+
+ // Should contain two regions, one with CSV->localAgg->globalAgg, another with keyword
+ val searchResult = new CostBasedScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan,
+ CONTROLLER
+ ).bottomUpSearch()
+
+ val groupByRegion =
+ searchResult.regionDAG.vertexSet().asScala.filter(region => region.physicalOps.size == 3).head
+ val keywordRegion =
+ searchResult.regionDAG.vertexSet().asScala.filter(region => region.physicalOps.size == 1).head
+
+ val resourceAllocator =
+ new DefaultResourceAllocator(
+ workflow.physicalPlan,
+ new ExecutionClusterInfo(),
+ workflow.context.workflowSettings
+ )
+
+ val costEstimator = new DefaultCostEstimator(
+ workflow.context,
+ resourceAllocator,
+ CONTROLLER
+ )
+
+ val (_, groupByRegionCost) = costEstimator.allocateResourcesAndEstimateCost(groupByRegion, 1)
+
+ val groupByOperatorCost = (groupByOpRuntimeStatistics.getField(6).asInstanceOf[Long] +
+ groupByOpRuntimeStatistics.getField(7).asInstanceOf[Long]) / 1e9
+
+ // The cost of the first region should be the cost of the GroupBy operator (note the two physical operators for
+ // the GroupBy logical operator have the same cost because we use logical operator in the statistics.
+ // The GroupBy operator has a longer running time.
+ assert(groupByRegionCost == groupByOperatorCost)
+
+ val (_, keywordRegionCost) = costEstimator.allocateResourcesAndEstimateCost(keywordRegion, 1)
+
+ val keywordOperatorCost = (keywordOpRuntimeStatistics.getField(6).asInstanceOf[Long] +
+ keywordOpRuntimeStatistics.getField(7).asInstanceOf[Long]) / 1e9
+
+ // The cost of the second region should be the cost of the keyword operator.
+ assert(keywordRegionCost == keywordOperatorCost)
+
+ // The cost of the region plan should be the sum of region costs
+ assert(searchResult.cost == groupByRegionCost + keywordRegionCost)
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/ExpansionGreedyScheduleGeneratorSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/ExpansionGreedyScheduleGeneratorSpec.scala
new file mode 100644
index 00000000000..e720b9c6cb5
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/ExpansionGreedyScheduleGeneratorSpec.scala
@@ -0,0 +1,337 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.core.virtualidentity.OperatorIdentity
+import org.apache.texera.amber.core.workflow.{PortIdentity, WorkflowContext}
+import org.apache.texera.amber.engine.e2e.TestUtils.buildWorkflow
+import org.apache.texera.amber.operator.TestOperators
+import org.apache.texera.amber.operator.split.SplitOpDesc
+import org.apache.texera.amber.operator.udf.python.{
+ DualInputPortsPythonUDFOpDescV2,
+ PythonUDFOpDescV2
+}
+import org.apache.texera.workflow.LogicalLink
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.flatspec.AnyFlatSpec
+
+@deprecated("This greedy schedule generator test will be removed in the future.")
+class ExpansionGreedyScheduleGeneratorSpec extends AnyFlatSpec with MockFactory {
+
+ "RegionPlanGenerator" should "correctly find regions in headerlessCsv->keyword workflow" in {
+ val headerlessCsvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc.operatorIdentifier,
+ PortIdentity(0),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(0)
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val (schedule, _) = new ExpansionGreedyScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan
+ ).generate()
+
+ // Assuming each level only has one region
+ val regionList = schedule.toList.map(level => level.head)
+ assert(regionList.size == 1)
+
+ regionList.zip(Iterator(2)).foreach {
+ case (region, opCount) =>
+ assert(region.getOperators.size == opCount)
+ }
+
+ regionList.zip(Iterator(1)).foreach {
+ case (region, linkCount) =>
+ assert(region.getLinks.size == linkCount)
+ }
+
+ regionList.zip(Iterator(3)).foreach {
+ case (region, portCount) =>
+ assert(region.getPorts.size == portCount)
+ }
+ }
+
+ "RegionPlanGenerator" should "correctly find regions in csv->(csv->)->join workflow" in {
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val headerlessCsvOpDesc2 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ headerlessCsvOpDesc2,
+ joinOpDesc
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc2.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val (schedule, _) = new ExpansionGreedyScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan
+ ).generate()
+
+ // Assuming each level only has one region
+ val regionList = schedule.toList.map(level => level.head)
+ assert(regionList.size == 2)
+
+ regionList.zip(Iterator(2, 2)).foreach {
+ case (region, opCount) =>
+ assert(region.getOperators.size == opCount)
+ }
+
+ regionList.zip(Iterator(1, 1)).foreach {
+ case (region, linkCount) =>
+ assert(region.getLinks.size == linkCount)
+ }
+
+ regionList.zip(Iterator(3, 4)).foreach {
+ case (region, portCount) =>
+ assert(region.getPorts.size == portCount)
+ }
+
+ // The fist region should be the build region
+ assert(
+ regionList.head.getOperators
+ .map(_.id)
+ .exists(physicalOpId =>
+ OperatorIdentity(physicalOpId.logicalOpId.id) == headerlessCsvOpDesc1.operatorIdentifier
+ )
+ )
+
+ // The second region should be the probe region
+ assert(
+ regionList(1).getOperators
+ .map(_.id)
+ .exists(physicalOpId =>
+ OperatorIdentity(physicalOpId.logicalOpId.id) == headerlessCsvOpDesc2.operatorIdentifier
+ )
+ )
+
+ }
+
+ "RegionPlanGenerator" should "correctly find regions in csv->->filter->join workflow" in {
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ keywordOpDesc,
+ joinOpDesc
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val (schedule, _) = new ExpansionGreedyScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan
+ ).generate()
+
+ // Assuming each level only has one region
+ val regionList = schedule.toList.map(level => level.head)
+ assert(regionList.size == 2)
+
+ regionList.zip(Iterator(3, 1)).foreach {
+ case (region, opCount) =>
+ assert(region.getOperators.size == opCount)
+ }
+
+ regionList.zip(Iterator(2, 0)).foreach {
+ case (region, linkCount) =>
+ assert(region.getLinks.size == linkCount)
+ }
+
+ regionList.zip(Iterator(5, 3)).foreach {
+ case (region, portCount) =>
+ assert(region.getPorts.size == portCount)
+ }
+ }
+//
+ "RegionPlanGenerator" should "correctly find regions in buildcsv->probecsv->hashjoin->hashjoin workflow" in {
+ val buildCsv = TestOperators.headerlessSmallCsvScanOpDesc()
+ val probeCsv = TestOperators.smallCsvScanOpDesc()
+ val hashJoin1 = TestOperators.joinOpDesc("column-1", "Region")
+ val hashJoin2 = TestOperators.joinOpDesc("column-2", "Country")
+ val workflow = buildWorkflow(
+ List(
+ buildCsv,
+ probeCsv,
+ hashJoin1,
+ hashJoin2
+ ),
+ List(
+ LogicalLink(
+ buildCsv.operatorIdentifier,
+ PortIdentity(),
+ hashJoin1.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ probeCsv.operatorIdentifier,
+ PortIdentity(),
+ hashJoin1.operatorIdentifier,
+ PortIdentity(1)
+ ),
+ LogicalLink(
+ buildCsv.operatorIdentifier,
+ PortIdentity(),
+ hashJoin2.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ hashJoin1.operatorIdentifier,
+ PortIdentity(),
+ hashJoin2.operatorIdentifier,
+ PortIdentity(1)
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val (schedule, _) = new ExpansionGreedyScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan
+ ).generate()
+
+ // Assuming each level only has one region
+ val regionList = schedule.toList.map(level => level.head)
+ assert(regionList.size == 2)
+ regionList.zip(Iterator(3, 3)).foreach {
+ case (region, opCount) =>
+ assert(region.getOperators.size == opCount)
+ }
+
+ regionList.zip(Iterator(2, 2)).foreach {
+ case (region, linkCount) =>
+ assert(region.getLinks.size == linkCount)
+ }
+
+ regionList.zip(Iterator(5, 7)).foreach {
+ case (region, portCount) =>
+ assert(region.getPorts.size == portCount)
+ }
+ }
+
+ "RegionPlanGenerator" should "correctly find regions in csv->split->training-infer workflow" in {
+ val csv = TestOperators.headerlessSmallCsvScanOpDesc()
+ val split = new SplitOpDesc()
+ val training = new PythonUDFOpDescV2()
+ val inference = new DualInputPortsPythonUDFOpDescV2()
+ val workflow = buildWorkflow(
+ List(
+ csv,
+ split,
+ training,
+ inference
+ ),
+ List(
+ LogicalLink(
+ csv.operatorIdentifier,
+ PortIdentity(),
+ split.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ split.operatorIdentifier,
+ PortIdentity(),
+ training.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ training.operatorIdentifier,
+ PortIdentity(),
+ inference.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ split.operatorIdentifier,
+ PortIdentity(1),
+ inference.operatorIdentifier,
+ PortIdentity(1)
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ val (schedule, _) = new ExpansionGreedyScheduleGenerator(
+ workflow.context,
+ workflow.physicalPlan
+ ).generate()
+
+ val regionList = schedule.toList.map(level => level.head)
+ assert(regionList.size == 2)
+ regionList.zip(Iterator(3, 1)).foreach {
+ case (region, opCount) =>
+ assert(region.getOperators.size == opCount)
+ }
+
+ regionList.zip(Iterator(2, 0)).foreach {
+ case (region, linkCount) =>
+ assert(region.getLinks.size == linkCount)
+ }
+
+ regionList.zip(Iterator(6, 3)).foreach {
+ case (region, portCount) =>
+ assert(region.getPorts.size == portCount)
+ }
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionCoordinatorTestSupport.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionCoordinatorTestSupport.scala
new file mode 100644
index 00000000000..facba102415
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionCoordinatorTestSupport.scala
@@ -0,0 +1,241 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import com.twitter.util.{Await, Duration, Future}
+import org.apache.pekko.actor.{Actor, ActorRef, Props}
+import org.apache.pekko.testkit.{TestActorRef, TestKit}
+import org.apache.texera.amber.core.executor.OpExecWithClassName
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity
+}
+import org.apache.texera.amber.core.workflow.PhysicalOp
+import org.apache.texera.amber.core.workflow.WorkflowContext.{
+ DEFAULT_EXECUTION_ID,
+ DEFAULT_WORKFLOW_ID
+}
+import org.apache.texera.amber.engine.architecture.common.{
+ AkkaActorRefMappingService,
+ AkkaActorService,
+ WorkflowActor
+}
+import org.apache.texera.amber.engine.architecture.controller.execution.WorkflowExecution
+import org.apache.texera.amber.engine.architecture.messaginglayer.{
+ NetworkInputGateway,
+ NetworkOutputGateway
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.ControlInvocation
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+import org.apache.texera.amber.engine.architecture.scheduling.config.{
+ OperatorConfig,
+ ResourceConfig,
+ WorkerConfig
+}
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState
+import org.apache.texera.amber.engine.common.CheckpointState
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.util.VirtualIdentityUtils
+
+import scala.collection.mutable
+
+object RegionCoordinatorTestSupport {
+ val InitializeExecutor = "initializeExecutor"
+ val OpenExecutor = "openExecutor"
+ val StartWorker = "startWorker"
+ val EndWorker = "endWorker"
+
+ // Generous deadline for the polling helpers below. Production timing under test (notably the
+ // 200 ms `killRetryDelay` in `RegionExecutionCoordinator`) fits comfortably; the rest is
+ // headroom for slow CI.
+ val testTimeout: Duration = Duration.fromSeconds(5)
+
+ case class WorkerRpcCall(
+ methodName: String,
+ receiver: ActorVirtualIdentity,
+ commandId: Long
+ )
+
+ case class ControllerHarnessFixture(
+ actorService: AkkaActorService,
+ actorRefService: AkkaActorRefMappingService
+ )
+
+ /**
+ * Captures controller-to-worker RPCs at the same boundary used by production
+ * `AsyncRPCClient.workerInterface`.
+ *
+ * Non-termination RPCs are completed immediately because these tests focus on termination
+ * ordering. `endWorker` responses are controlled by `endWorkerResponse`, allowing each test to
+ * hold termination pending, fail an attempt, or allow it to succeed.
+ */
+ class ControllerRpcProbe(endWorkerResponse: WorkerRpcCall => Option[ControlReturn]) {
+ val calls: mutable.ArrayBuffer[WorkerRpcCall] = mutable.ArrayBuffer()
+ val inputGateway = new NetworkInputGateway(CONTROLLER)
+ val outputGateway = new NetworkOutputGateway(CONTROLLER, handleOutput)
+ val asyncRPCClient = new AsyncRPCClient(inputGateway, outputGateway, CONTROLLER)
+
+ def methodTrace: Seq[String] = calls.map(_.methodName).toSeq
+
+ def initializedWorkers: Seq[ActorVirtualIdentity] =
+ calls.filter(_.methodName == InitializeExecutor).map(_.receiver).toSeq
+
+ def startedWorkers: Seq[ActorVirtualIdentity] =
+ calls.filter(_.methodName == StartWorker).map(_.receiver).toSeq
+
+ def endWorkerCalls: Seq[WorkerRpcCall] =
+ calls.filter(_.methodName == EndWorker).toSeq
+
+ def onlyEndWorkerCall: WorkerRpcCall = {
+ assert(endWorkerCalls.size == 1)
+ endWorkerCalls.head
+ }
+
+ def fulfill(call: WorkerRpcCall, returnValue: ControlReturn): Unit = {
+ asyncRPCClient.fulfillPromise(ReturnInvocation(call.commandId, returnValue))
+ }
+
+ private def handleOutput(message: WorkflowFIFOMessage): Unit = {
+ message.payload match {
+ case invocation: ControlInvocation =>
+ recordAndMaybeFulfill(invocation)
+ case _ =>
+ // Client events and stats updates are irrelevant to the coordinator lifecycle assertions.
+ }
+ }
+
+ private def recordAndMaybeFulfill(invocation: ControlInvocation): Unit = {
+ val call = WorkerRpcCall(
+ methodName = invocation.methodName,
+ receiver = invocation.context.receiver,
+ commandId = invocation.commandId
+ )
+ calls += call
+ immediateReturn(call).foreach(fulfill(call, _))
+ }
+
+ private def immediateReturn(call: WorkerRpcCall): Option[ControlReturn] = {
+ call.methodName match {
+ case InitializeExecutor | OpenExecutor =>
+ Some(EmptyReturn())
+ case StartWorker =>
+ Some(WorkerStateResponse(WorkerState.RUNNING))
+ case EndWorker =>
+ endWorkerResponse(call)
+ case other =>
+ throw new AssertionError(s"Unexpected worker RPC in test: $other")
+ }
+ }
+ }
+
+ class IdleActor extends Actor {
+ override def receive: Receive = { case _ => () }
+ }
+
+ class ControllerHarness extends WorkflowActor(None, CONTROLLER) {
+ override def handleInputMessage(id: Long, workflowMsg: WorkflowFIFOMessage): Unit = ()
+
+ override def getQueuedCredit(channelId: ChannelIdentity): Long = 0
+
+ override def handleBackpressure(isBackpressured: Boolean): Unit = ()
+
+ override def initState(): Unit = ()
+
+ override def loadFromCheckpoint(chkpt: CheckpointState): Unit = ()
+ }
+
+ def createSourceOp(logicalOpId: String): PhysicalOp =
+ PhysicalOp.sourcePhysicalOp(
+ PhysicalOpIdentity(OperatorIdentity(logicalOpId), "main"),
+ DEFAULT_WORKFLOW_ID,
+ DEFAULT_EXECUTION_ID,
+ OpExecWithClassName("unused")
+ )
+
+ def createWorkerId(physicalOp: PhysicalOp): ActorVirtualIdentity =
+ VirtualIdentityUtils.createWorkerIdentity(DEFAULT_WORKFLOW_ID, physicalOp.id, 0)
+
+ def createSingleWorkerRegion(
+ regionId: Long,
+ physicalOp: PhysicalOp,
+ workerId: ActorVirtualIdentity
+ ): Region =
+ Region(
+ RegionIdentity(regionId),
+ physicalOps = Set(physicalOp),
+ physicalLinks = Set.empty,
+ resourceConfig = Some(
+ ResourceConfig(
+ operatorConfigs = Map(physicalOp.id -> OperatorConfig(List(WorkerConfig(workerId))))
+ )
+ )
+ )
+
+ def seedReusableWorkerExecution(
+ workflowExecution: WorkflowExecution,
+ seedRegionId: Long,
+ physicalOp: PhysicalOp,
+ workerId: ActorVirtualIdentity
+ ): Unit = {
+ // RegionExecutionCoordinator skips real worker creation when an execution for this operator
+ // already exists.
+ workflowExecution
+ .initRegionExecution(createSingleWorkerRegion(seedRegionId, physicalOp, workerId))
+ .initOperatorExecution(physicalOp.id)
+ .initWorkerExecution(workerId)
+ }
+
+ def await[T](future: Future[T]): T = Await.result(future, testTimeout)
+
+ def waitUntil(condition: => Boolean): Unit = {
+ val deadline = System.nanoTime() + testTimeout.inNanoseconds
+ while (!condition && System.nanoTime() < deadline) {
+ Thread.sleep(20)
+ }
+ assert(condition, s"condition not satisfied within $testTimeout")
+ }
+}
+
+trait RegionCoordinatorTestSupport { self: TestKit =>
+ import RegionCoordinatorTestSupport._
+
+ protected def createControllerHarness(): ControllerHarnessFixture = {
+ val controllerRef = TestActorRef(new ControllerHarness)
+ controllerRef.underlyingActor.actorService.getAvailableNodeAddressesFunc = () =>
+ Array(controllerRef.path.address)
+ ControllerHarnessFixture(
+ actorService = controllerRef.underlyingActor.actorService,
+ actorRefService = controllerRef.underlyingActor.actorRefMappingService
+ )
+ }
+
+ protected def registerLiveWorker(
+ actorRefService: AkkaActorRefMappingService,
+ workerId: ActorVirtualIdentity
+ ): ActorRef = {
+ val workerRef = system.actorOf(Props(new IdleActor), s"worker-${System.nanoTime()}")
+ actorRefService.registerActorRef(workerId, workerRef)
+ workerRef
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinatorSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinatorSpec.scala
new file mode 100644
index 00000000000..8fab3b67fca
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinatorSpec.scala
@@ -0,0 +1,202 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import com.twitter.util.Future
+import org.apache.pekko.actor.ActorSystem
+import org.apache.pekko.testkit.TestKit
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow.PhysicalOp
+import org.apache.texera.amber.engine.architecture.common.AkkaActorRefMappingService
+import org.apache.texera.amber.engine.architecture.controller.ControllerConfig
+import org.apache.texera.amber.engine.architecture.controller.execution.WorkflowExecution
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns._
+import org.apache.texera.amber.engine.architecture.scheduling.RegionCoordinatorTestSupport._
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpecLike
+
+import java.util.concurrent.atomic
+
+/**
+ * Tests the real region-coordination lifecycle around synchronous region kill.
+ *
+ * The tests let the coordinator call the real `AsyncRPCClient.workerInterface`, capture the generated
+ * `ControlInvocation`s at the controller output gateway, and fulfill those RPC promises
+ * explicitly. This keeps the important production behavior under test:
+ *
+ * - regular launch RPCs (`initializeExecutor`, `openExecutor`, `startWorker`) are allowed to
+ * complete immediately;
+ * - `endWorker` can be held pending or failed to model worker-side drain/termination behavior;
+ * - the real coordinator then decides when to remove actor refs, clean control channels, mark
+ * workers terminated, and allow the next region to start.
+ */
+class RegionExecutionCoordinatorSpec
+ extends TestKit(ActorSystem("RegionExecutionCoordinatorSpec", AmberRuntime.akkaConfig))
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll
+ with RegionCoordinatorTestSupport {
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ "RegionExecutionCoordinator" should "send gracefulStop only after EndWorker succeeds" in {
+ val fixture = createSingleRegionFixture(endWorkerResponse = _ => None)
+
+ launchRegion(fixture.coordinator)
+ val completion = requestRegionCompletion(fixture.coordinator)
+
+ assert(
+ fixture.rpcProbe.methodTrace == Seq(InitializeExecutor, OpenExecutor, StartWorker, EndWorker)
+ )
+ assert(completion.poll.isEmpty)
+ assert(!fixture.coordinator.isCompleted)
+ assert(fixture.actorRefService.hasActorRef(fixture.workerId))
+
+ fixture.rpcProbe.fulfill(fixture.rpcProbe.onlyEndWorkerCall, EmptyReturn())
+ await(completion)
+
+ assert(fixture.coordinator.isCompleted)
+ assert(!fixture.actorRefService.hasActorRef(fixture.workerId))
+ assert(workerState(fixture) == WorkerState.TERMINATED)
+ assertControlChannelsAreRemoved(fixture)
+ }
+
+ it should "retry EndWorker failures and delay gracefulStop until a retry succeeds" in {
+ val attempts = new atomic.AtomicInteger(0)
+ val fixture = createSingleRegionFixture(endWorkerResponse =
+ _ =>
+ if (attempts.incrementAndGet() == 1) {
+ Some(transientEndWorkerFailure)
+ } else {
+ None
+ }
+ )
+
+ launchRegion(fixture.coordinator)
+ val completion = requestRegionCompletion(fixture.coordinator)
+
+ waitUntil(fixture.rpcProbe.endWorkerCalls.size >= 2)
+ assert(completion.poll.isEmpty)
+ assert(!fixture.coordinator.isCompleted)
+ assert(fixture.actorRefService.hasActorRef(fixture.workerId))
+
+ fixture.rpcProbe.fulfill(fixture.rpcProbe.endWorkerCalls.last, EmptyReturn())
+ await(completion)
+
+ assert(fixture.coordinator.isCompleted)
+ assert(fixture.rpcProbe.endWorkerCalls.size == 2)
+ assert(!fixture.actorRefService.hasActorRef(fixture.workerId))
+ assert(workerState(fixture) == WorkerState.TERMINATED)
+ }
+
+ private case class SingleRegionFixture(
+ coordinator: RegionExecutionCoordinator,
+ rpcProbe: ControllerRpcProbe,
+ workflowExecution: WorkflowExecution,
+ region: Region,
+ physicalOp: PhysicalOp,
+ workerId: ActorVirtualIdentity,
+ actorRefService: AkkaActorRefMappingService
+ )
+
+ private def createSingleRegionFixture(
+ endWorkerResponse: WorkerRpcCall => Option[ControlReturn]
+ ): SingleRegionFixture = {
+ val physicalOp = createSourceOp("test-op")
+ val workerId = createWorkerId(physicalOp)
+ val region = createSingleWorkerRegion(1, physicalOp, workerId)
+
+ val workflowExecution = WorkflowExecution()
+ seedReusableWorkerExecution(workflowExecution, seedRegionId = 0, physicalOp, workerId)
+ workflowExecution.initRegionExecution(region)
+
+ val rpcProbe = new ControllerRpcProbe(endWorkerResponse)
+ val controller = createControllerHarness()
+ registerLiveWorker(controller.actorRefService, workerId)
+
+ // Seed stale control channels to verify that successful termination removes them.
+ rpcProbe.inputGateway.getChannel(ChannelIdentity(workerId, CONTROLLER, isControl = true))
+ rpcProbe.outputGateway.getSequenceNumber(
+ ChannelIdentity(CONTROLLER, workerId, isControl = true)
+ )
+
+ val coordinator = new RegionExecutionCoordinator(
+ region,
+ isRestart = false,
+ workflowExecution,
+ rpcProbe.asyncRPCClient,
+ ControllerConfig(None, None, None, None),
+ controller.actorService,
+ controller.actorRefService
+ )
+
+ SingleRegionFixture(
+ coordinator = coordinator,
+ rpcProbe = rpcProbe,
+ workflowExecution = workflowExecution,
+ region = region,
+ physicalOp = physicalOp,
+ workerId = workerId,
+ actorRefService = controller.actorRefService
+ )
+ }
+
+ private def launchRegion(coordinator: RegionExecutionCoordinator): Unit = {
+ await(coordinator.syncStatusAndTransitionRegionExecutionPhase())
+ }
+
+ private def requestRegionCompletion(
+ coordinator: RegionExecutionCoordinator
+ ): Future[Unit] = {
+ coordinator.syncStatusAndTransitionRegionExecutionPhase()
+ }
+
+ private def workerState(fixture: SingleRegionFixture): WorkerState =
+ fixture.workflowExecution
+ .getRegionExecution(fixture.region.id)
+ .getOperatorExecution(fixture.physicalOp.id)
+ .getWorkerExecution(fixture.workerId)
+ .getState
+
+ private def assertControlChannelsAreRemoved(fixture: SingleRegionFixture): Unit = {
+ assert(
+ !fixture.rpcProbe.inputGateway.getAllControlChannels.exists(
+ _.channelId == ChannelIdentity(fixture.workerId, CONTROLLER, isControl = true)
+ )
+ )
+ assert(
+ !fixture.rpcProbe.outputGateway.getActiveChannels.exists(
+ _ == ChannelIdentity(CONTROLLER, fixture.workerId, isControl = true)
+ )
+ )
+ }
+
+ private def transientEndWorkerFailure: ControlError =
+ ControlError(
+ errorMessage = "transient EndWorker failure",
+ errorDetails = "",
+ stackTrace = "",
+ language = ErrorLanguage.SCALA
+ )
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionSpec.scala
new file mode 100644
index 00000000000..0aacaaeae29
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionSpec.scala
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.core.executor.OpExecInitInfo
+import org.apache.texera.amber.core.virtualidentity.{
+ ExecutionIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity,
+ WorkflowIdentity
+}
+import org.apache.texera.amber.core.workflow.{
+ GlobalPortIdentity,
+ PhysicalLink,
+ PhysicalOp,
+ PortIdentity
+}
+import org.scalatest.flatspec.AnyFlatSpec
+
+class RegionSpec extends AnyFlatSpec {
+
+ private def physicalOpId(opId: String): PhysicalOpIdentity =
+ PhysicalOpIdentity(OperatorIdentity(opId), "main")
+
+ private def op(opId: String): PhysicalOp =
+ PhysicalOp(
+ physicalOpId(opId),
+ WorkflowIdentity(0),
+ ExecutionIdentity(0),
+ OpExecInitInfo.Empty
+ )
+
+ private def link(fromOp: String, toOp: String): PhysicalLink =
+ PhysicalLink(physicalOpId(fromOp), PortIdentity(0), physicalOpId(toOp), PortIdentity(0))
+
+ "Region" should "expose the physical operators provided at construction" in {
+ val a = op("a")
+ val b = op("b")
+ val region = Region(RegionIdentity(1), Set(a, b), Set.empty)
+
+ assert(region.getOperators == Set(a, b))
+ }
+
+ it should "expose the physical links provided at construction" in {
+ val a = op("a")
+ val b = op("b")
+ val ab = link("a", "b")
+ val region = Region(RegionIdentity(1), Set(a, b), Set(ab))
+
+ assert(region.getLinks == Set(ab))
+ }
+
+ it should "default ports to an empty set" in {
+ val region = Region(RegionIdentity(1), Set(op("a")), Set.empty)
+ assert(region.getPorts.isEmpty)
+ }
+
+ it should "expose the ports provided at construction" in {
+ val portId = GlobalPortIdentity(physicalOpId("a"), PortIdentity(0), input = true)
+ val region = Region(RegionIdentity(1), Set(op("a")), Set.empty, ports = Set(portId))
+ assert(region.getPorts == Set(portId))
+ }
+
+ "Region.getOperator" should "look up a physical operator by id" in {
+ val a = op("a")
+ val b = op("b")
+ val region = Region(RegionIdentity(1), Set(a, b), Set.empty)
+
+ assert(region.getOperator(physicalOpId("a")) == a)
+ assert(region.getOperator(physicalOpId("b")) == b)
+ }
+
+ it should "throw NoSuchElementException for an unknown operator id" in {
+ val region = Region(RegionIdentity(1), Set(op("a")), Set.empty)
+ assertThrows[NoSuchElementException] {
+ region.getOperator(physicalOpId("missing"))
+ }
+ }
+
+ "Region.topologicalIterator" should "yield operators in topological order based on physical links" in {
+ val a = op("a")
+ val b = op("b")
+ val c = op("c")
+ val region = Region(RegionIdentity(1), Set(a, b, c), Set(link("a", "b"), link("b", "c")))
+
+ assert(
+ region.topologicalIterator().toList ==
+ List(physicalOpId("a"), physicalOpId("b"), physicalOpId("c"))
+ )
+ }
+
+ "Region.getSourceOperators" should "treat operators without input ports as sources" in {
+ val a = op("a")
+ val b = op("b")
+ val region = Region(RegionIdentity(1), Set(a, b), Set.empty)
+
+ assert(region.getSourceOperators == Set(a, b))
+ }
+
+ "Region.getStarterOperators" should "match getSourceOperators when no resource config is provided" in {
+ val a = op("a")
+ val b = op("b")
+ val region = Region(RegionIdentity(1), Set(a, b), Set.empty)
+
+ assert(region.getStarterOperators == region.getSourceOperators)
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/SchedulingUtilsSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/SchedulingUtilsSpec.scala
new file mode 100644
index 00000000000..18ee5e88f12
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/SchedulingUtilsSpec.scala
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.texera.amber.core.executor.OpExecInitInfo
+import org.apache.texera.amber.core.virtualidentity.{
+ ExecutionIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity,
+ WorkflowIdentity
+}
+import org.apache.texera.amber.core.workflow.PhysicalOp
+import org.jgrapht.graph.DirectedAcyclicGraph
+import org.scalatest.flatspec.AnyFlatSpec
+
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+class SchedulingUtilsSpec extends AnyFlatSpec {
+
+ private def region(regionId: Long, opId: String): Region = {
+ val physicalOp = PhysicalOp(
+ PhysicalOpIdentity(OperatorIdentity(opId), "main"),
+ WorkflowIdentity(0),
+ ExecutionIdentity(0),
+ OpExecInitInfo.Empty
+ )
+ Region(RegionIdentity(regionId), Set(physicalOp), Set.empty)
+ }
+
+ private def newGraph(): DirectedAcyclicGraph[Region, RegionLink] =
+ new DirectedAcyclicGraph[Region, RegionLink](classOf[RegionLink])
+
+ "SchedulingUtils.replaceVertex" should "replace an isolated vertex with no incident edges" in {
+ val graph = newGraph()
+ val oldVertex = region(1, "a")
+ val newVertex = region(1, "a-prime")
+ graph.addVertex(oldVertex)
+
+ SchedulingUtils.replaceVertex(graph, oldVertex, newVertex)
+
+ assert(!graph.containsVertex(oldVertex))
+ assert(graph.containsVertex(newVertex))
+ assert(graph.edgeSet().isEmpty)
+ }
+
+ it should "rewrite outgoing edges to originate from the new vertex" in {
+ val graph = newGraph()
+ val oldVertex = region(1, "a")
+ val downstream = region(2, "b")
+ val newVertex = region(1, "a-prime")
+ graph.addVertex(oldVertex)
+ graph.addVertex(downstream)
+ graph.addEdge(oldVertex, downstream, RegionLink(oldVertex.id, downstream.id))
+
+ SchedulingUtils.replaceVertex(graph, oldVertex, newVertex)
+
+ assert(!graph.containsVertex(oldVertex))
+ assert(graph.containsVertex(newVertex))
+ val outgoing = graph.outgoingEdgesOf(newVertex).asScala.toList
+ assert(outgoing.size == 1)
+ assert(graph.getEdgeTarget(outgoing.head) == downstream)
+ assert(outgoing.head == RegionLink(newVertex.id, downstream.id))
+ }
+
+ it should "rewrite incoming edges to terminate at the new vertex" in {
+ val graph = newGraph()
+ val upstream = region(0, "u")
+ val oldVertex = region(1, "a")
+ val newVertex = region(1, "a-prime")
+ graph.addVertex(upstream)
+ graph.addVertex(oldVertex)
+ graph.addEdge(upstream, oldVertex, RegionLink(upstream.id, oldVertex.id))
+
+ SchedulingUtils.replaceVertex(graph, oldVertex, newVertex)
+
+ assert(!graph.containsVertex(oldVertex))
+ val incoming = graph.incomingEdgesOf(newVertex).asScala.toList
+ assert(incoming.size == 1)
+ assert(graph.getEdgeSource(incoming.head) == upstream)
+ assert(incoming.head == RegionLink(upstream.id, newVertex.id))
+ }
+
+ it should "preserve both upstream and downstream edges in a chain" in {
+ val graph = newGraph()
+ val upstream = region(0, "u")
+ val oldVertex = region(1, "a")
+ val downstream = region(2, "d")
+ val newVertex = region(1, "a-prime")
+ graph.addVertex(upstream)
+ graph.addVertex(oldVertex)
+ graph.addVertex(downstream)
+ graph.addEdge(upstream, oldVertex, RegionLink(upstream.id, oldVertex.id))
+ graph.addEdge(oldVertex, downstream, RegionLink(oldVertex.id, downstream.id))
+
+ SchedulingUtils.replaceVertex(graph, oldVertex, newVertex)
+
+ assert(graph.vertexSet().asScala.toSet == Set(upstream, newVertex, downstream))
+ assert(
+ graph.edgeSet().asScala.toSet ==
+ Set(
+ RegionLink(upstream.id, newVertex.id),
+ RegionLink(newVertex.id, downstream.id)
+ )
+ )
+ }
+
+ it should "leave the graph unchanged when old and new vertices are equal" in {
+ val graph = newGraph()
+ val upstream = region(0, "u")
+ val vertex = region(1, "a")
+ val downstream = region(2, "d")
+ graph.addVertex(upstream)
+ graph.addVertex(vertex)
+ graph.addVertex(downstream)
+ graph.addEdge(upstream, vertex, RegionLink(upstream.id, vertex.id))
+ graph.addEdge(vertex, downstream, RegionLink(vertex.id, downstream.id))
+
+ SchedulingUtils.replaceVertex(graph, vertex, vertex)
+
+ assert(graph.vertexSet().asScala.toSet == Set(upstream, vertex, downstream))
+ assert(graph.edgeSet().size == 2)
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinatorSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinatorSpec.scala
new file mode 100644
index 00000000000..f4372e8b575
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinatorSpec.scala
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.scheduling
+
+import org.apache.pekko.actor.ActorSystem
+import org.apache.pekko.testkit.TestKit
+import org.apache.texera.amber.engine.architecture.controller.ControllerConfig
+import org.apache.texera.amber.engine.architecture.controller.execution.WorkflowExecution
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.scheduling.RegionCoordinatorTestSupport._
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpecLike
+
+import scala.collection.mutable
+
+class WorkflowExecutionCoordinatorSpec
+ extends TestKit(ActorSystem("WorkflowExecutionCoordinatorSpec", AmberRuntime.akkaConfig))
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll
+ with RegionCoordinatorTestSupport {
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ "WorkflowExecutionCoordinator" should "start the next region only after previous region termination succeeds" in {
+ val firstOp = createSourceOp("first-op")
+ val firstWorkerId = createWorkerId(firstOp)
+ val firstRegion = createSingleWorkerRegion(1, firstOp, firstWorkerId)
+
+ val secondOp = createSourceOp("second-op")
+ val secondWorkerId = createWorkerId(secondOp)
+ val secondRegion = createSingleWorkerRegion(2, secondOp, secondWorkerId)
+
+ val workflowExecution = WorkflowExecution()
+ seedReusableWorkerExecution(workflowExecution, seedRegionId = 101, firstOp, firstWorkerId)
+ seedReusableWorkerExecution(workflowExecution, seedRegionId = 102, secondOp, secondWorkerId)
+
+ // First region's worker holds endWorker pending until we explicitly fulfill it; the second
+ // region's worker terminates immediately. This lets us assert the second region cannot start
+ // until termination of the first finishes.
+ val rpcProbe = new ControllerRpcProbe(
+ endWorkerResponse = call => if (call.receiver == firstWorkerId) None else Some(EmptyReturn())
+ )
+ val controller = createControllerHarness()
+ registerLiveWorker(controller.actorRefService, firstWorkerId)
+ registerLiveWorker(controller.actorRefService, secondWorkerId)
+
+ val nextRegionLevels = mutable.Queue(Set(firstRegion), Set(secondRegion))
+ val workflowCoordinator = new WorkflowExecutionCoordinator(
+ () => if (nextRegionLevels.nonEmpty) nextRegionLevels.dequeue() else Set.empty,
+ workflowExecution,
+ ControllerConfig(None, None, None, None),
+ rpcProbe.asyncRPCClient
+ )
+ workflowCoordinator.setupActorRefService(controller.actorRefService)
+
+ await(workflowCoordinator.coordinateRegionExecutors(controller.actorService))
+ assert(rpcProbe.startedWorkers == Seq(firstWorkerId))
+
+ val coordination = workflowCoordinator.coordinateRegionExecutors(controller.actorService)
+
+ waitUntil(rpcProbe.endWorkerCalls.size == 1)
+ assert(coordination.poll.isEmpty)
+ assert(!rpcProbe.initializedWorkers.contains(secondWorkerId))
+ assert(controller.actorRefService.hasActorRef(firstWorkerId))
+
+ rpcProbe.fulfill(rpcProbe.onlyEndWorkerCall, EmptyReturn())
+ await(coordination)
+
+ assert(!controller.actorRefService.hasActorRef(firstWorkerId))
+ assert(rpcProbe.initializedWorkers.contains(secondWorkerId))
+ assert(rpcProbe.startedWorkers.contains(secondWorkerId))
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/DPThreadSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/DPThreadSpec.scala
new file mode 100644
index 00000000000..d8b5d57d638
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/DPThreadSpec.scala
@@ -0,0 +1,243 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import org.apache.texera.amber.core.executor.OperatorExecutor
+import org.apache.texera.amber.core.tuple.{AttributeType, Schema, Tuple, TupleLike}
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.engine.architecture.logreplay.{ReplayLogManager, ReplayLogRecord}
+import org.apache.texera.amber.engine.architecture.messaginglayer.WorkerTimerService
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.{
+ METHOD_PAUSE_WORKER,
+ METHOD_RESUME_WORKER
+}
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ DPInputQueueElement,
+ FIFOMessageElement,
+ TimerBasedControlElement
+}
+import org.apache.texera.amber.engine.common.ambermessage.{DataFrame, WorkflowFIFOMessage}
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.engine.common.virtualidentity.util.SELF
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.flatspec.AnyFlatSpec
+
+import java.net.URI
+import java.util.concurrent.LinkedBlockingQueue
+
+class DPThreadSpec extends AnyFlatSpec with MockFactory {
+
+ private val workerId: ActorVirtualIdentity = ActorVirtualIdentity("DP mock")
+ private val senderWorkerId: ActorVirtualIdentity = ActorVirtualIdentity("mock sender")
+ private val dataChannelId = ChannelIdentity(senderWorkerId, workerId, isControl = false)
+ private val controlChannelId = ChannelIdentity(senderWorkerId, workerId, isControl = true)
+ private val executor = mock[OperatorExecutor]
+ private val mockInputPortId = PortIdentity()
+
+ private val schema: Schema = Schema().add("field1", AttributeType.INTEGER)
+ private val tuples: Array[Tuple] = (0 until 5000)
+ .map(i => TupleLike(i).enforceSchema(schema))
+ .toArray
+ private val logStorage = SequentialRecordStorage.getStorage[ReplayLogRecord](None)
+ private val logManager: ReplayLogManager =
+ ReplayLogManager.createLogManager(logStorage, "none", x => {})
+
+ "DP Thread" should "handle pause/resume during processing" in {
+ val inputQueue = new LinkedBlockingQueue[DPInputQueueElement]()
+ val dp = new DataProcessor(workerId, x => {}, inputMessageQueue = inputQueue)
+ dp.executor = executor
+ dp.inputManager.addPort(mockInputPortId, schema, List.empty, List.empty)
+ dp.inputGateway.getChannel(dataChannelId).setPortId(mockInputPortId)
+ dp.adaptiveBatchingMonitor = mock[WorkerTimerService]
+ (dp.adaptiveBatchingMonitor.resumeAdaptiveBatching _).expects().anyNumberOfTimes()
+ val dpThread = new DPThread(workerId, dp, logManager, inputQueue)
+ dpThread.start()
+ tuples.foreach { x =>
+ (
+ (
+ tuple: Tuple,
+ input: Int
+ ) => executor.processTupleMultiPort(tuple, input)
+ )
+ .expects(x, 0)
+ }
+ val message = WorkflowFIFOMessage(dataChannelId, 0, DataFrame(tuples))
+ inputQueue.put(FIFOMessageElement(message))
+ inputQueue.put(
+ TimerBasedControlElement(
+ ControlInvocation(METHOD_PAUSE_WORKER, EmptyRequest(), AsyncRPCContext(SELF, SELF), 0)
+ )
+ )
+ Thread.sleep(1000)
+ assert(dp.pauseManager.isPaused)
+ inputQueue.put(
+ TimerBasedControlElement(
+ ControlInvocation(METHOD_RESUME_WORKER, EmptyRequest(), AsyncRPCContext(SELF, SELF), 1)
+ )
+ )
+ Thread.sleep(1000)
+ while (dp.inputManager.hasUnfinishedInput) {
+ Thread.sleep(100)
+ }
+ }
+
+ "DP Thread" should "handle pause/resume using fifo messages" in {
+ val inputQueue = new LinkedBlockingQueue[DPInputQueueElement]()
+ val dp = new DataProcessor(workerId, x => {}, inputMessageQueue = inputQueue)
+ dp.inputManager.addPort(mockInputPortId, schema, List.empty, List.empty)
+ dp.inputGateway.getChannel(dataChannelId).setPortId(mockInputPortId)
+ dp.adaptiveBatchingMonitor = mock[WorkerTimerService]
+ (dp.adaptiveBatchingMonitor.resumeAdaptiveBatching _).expects().anyNumberOfTimes()
+ val dpThread = new DPThread(workerId, dp, logManager, inputQueue)
+ dp.executor = executor
+ dpThread.start()
+ tuples.foreach { x =>
+ (
+ (
+ tuple: Tuple,
+ input: Int
+ ) => executor.processTupleMultiPort(tuple, input)
+ )
+ .expects(x, 0)
+ }
+ val message = WorkflowFIFOMessage(dataChannelId, 0, DataFrame(tuples))
+ val pauseControl = WorkflowFIFOMessage(
+ controlChannelId,
+ 0,
+ ControlInvocation(METHOD_PAUSE_WORKER, EmptyRequest(), AsyncRPCContext(SELF, SELF), 0)
+ )
+ val resumeControl =
+ WorkflowFIFOMessage(
+ controlChannelId,
+ 1,
+ ControlInvocation(METHOD_RESUME_WORKER, EmptyRequest(), AsyncRPCContext(SELF, SELF), 1)
+ )
+ inputQueue.put(FIFOMessageElement(message))
+ inputQueue.put(
+ FIFOMessageElement(pauseControl)
+ )
+ Thread.sleep(1000)
+ assert(dp.pauseManager.isPaused)
+ inputQueue.put(FIFOMessageElement(resumeControl))
+ Thread.sleep(1000)
+ while (dp.inputManager.hasUnfinishedInput) {
+ Thread.sleep(100)
+ }
+ }
+
+ "DP Thread" should "handle multiple batches from multiple sources" in {
+ val inputQueue = new LinkedBlockingQueue[DPInputQueueElement]()
+ val dp = new DataProcessor(workerId, x => {}, inputMessageQueue = inputQueue)
+ dp.executor = executor
+ val anotherSenderWorkerId = ActorVirtualIdentity("another")
+ dp.inputManager.addPort(mockInputPortId, schema, List.empty, List.empty)
+ dp.inputGateway.getChannel(dataChannelId).setPortId(mockInputPortId)
+ dp.inputGateway
+ .getChannel(ChannelIdentity(anotherSenderWorkerId, workerId, isControl = false))
+ .setPortId(mockInputPortId)
+ dp.adaptiveBatchingMonitor = mock[WorkerTimerService]
+ (dp.adaptiveBatchingMonitor.resumeAdaptiveBatching _).expects().anyNumberOfTimes()
+ val dpThread = new DPThread(workerId, dp, logManager, inputQueue)
+ dpThread.start()
+ tuples.foreach { x =>
+ (
+ (
+ tuple: Tuple,
+ input: Int
+ ) => executor.processTupleMultiPort(tuple, input)
+ )
+ .expects(x, 0)
+ }
+ val dataChannelID2 = ChannelIdentity(anotherSenderWorkerId, workerId, isControl = false)
+ val message1 = WorkflowFIFOMessage(dataChannelId, 0, DataFrame(tuples.slice(0, 100)))
+ val message2 = WorkflowFIFOMessage(dataChannelId, 1, DataFrame(tuples.slice(100, 200)))
+ val message3 = WorkflowFIFOMessage(dataChannelID2, 0, DataFrame(tuples.slice(300, 1000)))
+ val message4 = WorkflowFIFOMessage(dataChannelId, 2, DataFrame(tuples.slice(200, 300)))
+ val message5 = WorkflowFIFOMessage(dataChannelID2, 1, DataFrame(tuples.slice(1000, 5000)))
+ inputQueue.put(FIFOMessageElement(message1))
+ inputQueue.put(FIFOMessageElement(message2))
+ inputQueue.put(FIFOMessageElement(message3))
+ inputQueue.put(FIFOMessageElement(message4))
+ inputQueue.put(FIFOMessageElement(message5))
+ Thread.sleep(1000)
+ while (dp.inputManager.hasUnfinishedInput) {
+ Thread.sleep(100)
+ }
+ }
+
+ "DP Thread" should "write determinant logs to local storage while processing" in {
+ val inputQueue = new LinkedBlockingQueue[DPInputQueueElement]()
+ val dp = new DataProcessor(workerId, _ => {}, inputMessageQueue = inputQueue)
+ dp.executor = executor
+ val anotherSenderWorkerId = ActorVirtualIdentity("another")
+ dp.inputManager.addPort(mockInputPortId, schema, List.empty, List.empty)
+ dp.inputGateway.getChannel(dataChannelId).setPortId(mockInputPortId)
+ dp.inputGateway
+ .getChannel(ChannelIdentity(anotherSenderWorkerId, workerId, isControl = false))
+ .setPortId(mockInputPortId)
+ dp.adaptiveBatchingMonitor = mock[WorkerTimerService]
+ (dp.adaptiveBatchingMonitor.resumeAdaptiveBatching _).expects().anyNumberOfTimes()
+ val logStorage = SequentialRecordStorage.getStorage[ReplayLogRecord](
+ Some(new URI("ram:///recovery-logs/tmp"))
+ )
+ logStorage.deleteStorage()
+ val logManager: ReplayLogManager =
+ ReplayLogManager.createLogManager(logStorage, "tmpLog", _ => {})
+ val dpThread = new DPThread(workerId, dp, logManager, inputQueue)
+ dpThread.start()
+ tuples.foreach { x =>
+ (
+ (
+ tuple: Tuple,
+ input: Int
+ ) => executor.processTupleMultiPort(tuple, input)
+ )
+ .expects(x, 0)
+ }
+ val dataChannelId2 = ChannelIdentity(anotherSenderWorkerId, workerId, isControl = false)
+ val message1 = WorkflowFIFOMessage(dataChannelId, 0, DataFrame(tuples.slice(0, 100)))
+ val message2 = WorkflowFIFOMessage(dataChannelId, 1, DataFrame(tuples.slice(100, 200)))
+ val message3 = WorkflowFIFOMessage(dataChannelId2, 0, DataFrame(tuples.slice(300, 1000)))
+ val message4 = WorkflowFIFOMessage(dataChannelId, 2, DataFrame(tuples.slice(200, 300)))
+ val message5 = WorkflowFIFOMessage(dataChannelId2, 1, DataFrame(tuples.slice(1000, 5000)))
+ inputQueue.put(FIFOMessageElement(message1))
+ inputQueue.put(FIFOMessageElement(message2))
+ inputQueue.put(FIFOMessageElement(message3))
+ Thread.sleep(1000)
+ inputQueue.put(FIFOMessageElement(message4))
+ inputQueue.put(FIFOMessageElement(message5))
+ Thread.sleep(1000)
+ while (logManager.getStep < 4999) {
+ Thread.sleep(100)
+ }
+ logManager.sendCommitted(null) // drain in-mem records to flush
+ logManager.terminate()
+ val logs = logStorage.getReader("tmpLog").mkRecordIterator().toArray
+ logStorage.deleteStorage()
+ assert(logs.length > 1)
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessorSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessorSpec.scala
new file mode 100644
index 00000000000..0f4840f9653
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessorSpec.scala
@@ -0,0 +1,244 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import org.apache.texera.amber.core.executor.OperatorExecutor
+import org.apache.texera.amber.core.tuple.{AttributeType, Schema, Tuple, TupleLike}
+import org.apache.texera.amber.core.virtualidentity._
+import org.apache.texera.amber.core.workflow.PortIdentity
+import org.apache.texera.amber.core.workflow.WorkflowContext.DEFAULT_WORKFLOW_ID
+import org.apache.texera.amber.engine.architecture.logreplay.{ReplayLogManager, ReplayLogRecord}
+import org.apache.texera.amber.engine.architecture.messaginglayer.WorkerTimerService
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmbeddedControlMessage,
+ EmbeddedControlMessageType,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.{
+ METHOD_END_CHANNEL,
+ METHOD_FLUSH_NETWORK_BUFFER,
+ METHOD_OPEN_EXECUTOR
+}
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ DPInputQueueElement,
+ MainThreadDelegateMessage
+}
+import org.apache.texera.amber.engine.architecture.worker.statistics.WorkerState.READY
+import org.apache.texera.amber.engine.common.ambermessage.{DataFrame, WorkflowFIFOMessage}
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.util.VirtualIdentityUtils
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.flatspec.AnyFlatSpec
+
+import java.util.concurrent.LinkedBlockingQueue
+
+class DataProcessorSpec extends AnyFlatSpec with MockFactory with BeforeAndAfterEach {
+ private val testOpId = PhysicalOpIdentity(OperatorIdentity("testop"), "main")
+ private val upstreamOpId = PhysicalOpIdentity(OperatorIdentity("sender"), "main")
+ private val testWorkerId: ActorVirtualIdentity = VirtualIdentityUtils.createWorkerIdentity(
+ DEFAULT_WORKFLOW_ID,
+ testOpId,
+ 0
+ )
+ private val senderWorkerId: ActorVirtualIdentity = VirtualIdentityUtils.createWorkerIdentity(
+ DEFAULT_WORKFLOW_ID,
+ upstreamOpId,
+ 0
+ )
+
+ private val executor = mock[OperatorExecutor]
+ private val inputPortId = PortIdentity()
+ private val outputPortId = PortIdentity()
+ private val outputHandler = mock[Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit]
+ private val adaptiveBatchingMonitor = mock[WorkerTimerService]
+ private val schema: Schema = Schema().add("field1", AttributeType.INTEGER)
+ private val tuples: Array[Tuple] = (0 until 400)
+ .map(i => TupleLike(i).enforceSchema(schema))
+ .toArray
+ private val logStorage = SequentialRecordStorage.getStorage[ReplayLogRecord](None)
+ private val logManager: ReplayLogManager =
+ ReplayLogManager.createLogManager(logStorage, "none", x => {})
+ private val endChannelPayload = EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity("EndChannel"),
+ EmbeddedControlMessageType.PORT_ALIGNMENT,
+ Seq(),
+ Map(
+ testWorkerId.name ->
+ ControlInvocation(
+ METHOD_END_CHANNEL.getBareMethodName,
+ EmptyRequest(),
+ AsyncRPCContext(ActorVirtualIdentity(""), ActorVirtualIdentity("")),
+ -1
+ )
+ )
+ )
+
+ def mkDataProcessor: DataProcessor = {
+ val dp: DataProcessor = new DataProcessor(
+ testWorkerId,
+ outputHandler,
+ inputMessageQueue = new LinkedBlockingQueue[DPInputQueueElement]()
+ )
+ dp.initTimerService(adaptiveBatchingMonitor)
+ dp
+ }
+
+ "data processor" should "process data messages" in {
+ val dp = mkDataProcessor
+ dp.executor = executor
+ dp.stateManager.transitTo(READY)
+ (outputHandler.apply _).expects(*).once()
+ (executor.open _).expects().once()
+ tuples.foreach { x =>
+ (
+ (
+ tuple: Tuple,
+ input: Int
+ ) => executor.processTupleMultiPort(tuple, input)
+ )
+ .expects(x, 0)
+ }
+ (
+ (
+ input: Int
+ ) => executor.produceStateOnFinish(input)
+ )
+ .expects(0)
+ .returning(None)
+ (
+ (
+ input: Int
+ ) => executor.onFinishMultiPort(input)
+ )
+ .expects(
+ 0
+ )
+ (adaptiveBatchingMonitor.startAdaptiveBatching _).expects().anyNumberOfTimes()
+ (adaptiveBatchingMonitor.stopAdaptiveBatching _).expects().once()
+ (executor.close _).expects().once()
+ (outputHandler.apply _).expects(*).anyNumberOfTimes()
+ dp.inputManager.addPort(inputPortId, schema, List.empty, List.empty)
+ dp.inputGateway
+ .getChannel(ChannelIdentity(senderWorkerId, testWorkerId, isControl = false))
+ .setPortId(inputPortId)
+ dp.outputManager.addPort(outputPortId, schema, None)
+ dp.processDCM(
+ ChannelIdentity(CONTROLLER, testWorkerId, isControl = true),
+ ControlInvocation(
+ METHOD_OPEN_EXECUTOR,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, testWorkerId),
+ 0
+ )
+ )
+ dp.processDataPayload(
+ ChannelIdentity(senderWorkerId, testWorkerId, isControl = false),
+ DataFrame(tuples)
+ )
+ while (dp.inputManager.hasUnfinishedInput || dp.outputManager.hasUnfinishedOutput) {
+ dp.continueDataProcessing()
+ }
+ dp.processECM(
+ ChannelIdentity(senderWorkerId, testWorkerId, isControl = false),
+ endChannelPayload,
+ logManager
+ )
+
+ while (dp.inputManager.hasUnfinishedInput || dp.outputManager.hasUnfinishedOutput) {
+ dp.continueDataProcessing()
+ }
+ }
+
+ "data processor" should "process control messages during data processing" in {
+ val dp = mkDataProcessor
+ dp.executor = executor
+ dp.stateManager.transitTo(READY)
+ (outputHandler.apply _).expects(*).anyNumberOfTimes()
+ (executor.open _).expects().once()
+ tuples.foreach { x =>
+ (
+ (
+ tuple: Tuple,
+ input: Int
+ ) => executor.processTupleMultiPort(tuple, input)
+ )
+ .expects(x, 0)
+ }
+ (
+ (
+ input: Int
+ ) => executor.produceStateOnFinish(input)
+ )
+ .expects(0)
+ .returning(None)
+ (
+ (
+ input: Int
+ ) => executor.onFinishMultiPort(input)
+ )
+ .expects(0)
+ (adaptiveBatchingMonitor.startAdaptiveBatching _).expects().anyNumberOfTimes()
+ dp.inputManager.addPort(inputPortId, schema, List.empty, List.empty)
+ dp.inputGateway
+ .getChannel(ChannelIdentity(senderWorkerId, testWorkerId, isControl = false))
+ .setPortId(inputPortId)
+ dp.outputManager.addPort(outputPortId, schema, None)
+ dp.processDCM(
+ ChannelIdentity(CONTROLLER, testWorkerId, isControl = true),
+ ControlInvocation(
+ METHOD_OPEN_EXECUTOR,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, testWorkerId),
+ 0
+ )
+ )
+ dp.processDataPayload(
+ ChannelIdentity(senderWorkerId, testWorkerId, isControl = false),
+ DataFrame(tuples)
+ )
+ while (dp.inputManager.hasUnfinishedInput || dp.outputManager.hasUnfinishedOutput) {
+ dp.processDCM(
+ ChannelIdentity(CONTROLLER, testWorkerId, isControl = true),
+ ControlInvocation(
+ METHOD_FLUSH_NETWORK_BUFFER,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, testWorkerId),
+ 1
+ )
+ )
+ dp.continueDataProcessing()
+ }
+ (adaptiveBatchingMonitor.stopAdaptiveBatching _).expects().once()
+ (executor.close _).expects().once()
+ dp.processECM(
+ ChannelIdentity(senderWorkerId, testWorkerId, isControl = false),
+ endChannelPayload,
+ logManager
+ )
+ while (dp.inputManager.hasUnfinishedInput || dp.outputManager.hasUnfinishedOutput) {
+ dp.continueDataProcessing()
+ }
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/WorkerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/WorkerSpec.scala
new file mode 100644
index 00000000000..890fe97b852
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/WorkerSpec.scala
@@ -0,0 +1,317 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker
+
+import org.apache.pekko.actor.{ActorRef, ActorSystem, Props}
+import org.apache.pekko.testkit.{ImplicitSender, TestActorRef, TestKit}
+import org.apache.texera.amber.clustering.SingleNodeListener
+import org.apache.texera.amber.core.executor.{OpExecWithClassName, OperatorExecutor}
+import org.apache.texera.amber.core.tuple._
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity
+}
+import org.apache.texera.amber.core.workflow.{PhysicalLink, PortIdentity}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkMessage
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc._
+import org.apache.texera.amber.engine.architecture.scheduling.config.WorkerConfig
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.OneToOnePartitioning
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ DPInputQueueElement,
+ MainThreadDelegateMessage,
+ WorkerReplayInitialization
+}
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataFrame,
+ DataPayload,
+ WorkflowFIFOMessage
+}
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpecLike
+
+import java.util.concurrent.{CompletableFuture, LinkedBlockingQueue}
+import scala.collection.mutable
+import scala.concurrent.duration.MILLISECONDS
+import scala.util.Random
+class DummyOperatorExecutor extends OperatorExecutor {
+ override def processTuple(tuple: Tuple, port: Int): Iterator[TupleLike] = {
+ Iterator(tuple)
+ }
+}
+
+class WorkerSpec
+ extends TestKit(ActorSystem("WorkerSpec", AmberRuntime.akkaConfig))
+ with ImplicitSender
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll
+ with MockFactory {
+
+ def mkSchema(fields: Any*): Schema = {
+ var schema = Schema()
+ fields.indices.foreach { i =>
+ schema = schema.add(new Attribute("field" + i, AttributeType.ANY))
+ }
+ schema
+ }
+
+ def mkTuple(fields: Any*): Tuple = {
+ Tuple.builder(mkSchema(fields: _*)).addSequentially(fields.toArray).build()
+ }
+
+ override def beforeAll(): Unit = {
+ system.actorOf(Props[SingleNodeListener](), "cluster-info")
+ }
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ private val identifier1 = ActorVirtualIdentity("Worker:WF1-E1-op-layer-1")
+ private val identifier2 = ActorVirtualIdentity("Worker:WF1-E1-op-layer-2")
+
+ private val operatorIdentity = OperatorIdentity("testOperator")
+
+ private val mockPortId = PortIdentity()
+ private val mockLink =
+ PhysicalLink(
+ PhysicalOpIdentity(operatorIdentity, "1st-physical-op"),
+ mockPortId,
+ PhysicalOpIdentity(operatorIdentity, "2nd-physical-op"),
+ mockPortId
+ )
+
+ private val mockPolicy =
+ OneToOnePartitioning(10, Seq(ChannelIdentity(identifier1, identifier2, isControl = false)))
+
+ def sendControlToWorker(
+ worker: ActorRef,
+ controls: Array[ControlInvocation],
+ beginSeqNum: Long = 0
+ ): Unit = {
+ var seq = beginSeqNum
+ controls.foreach { ctrl =>
+ worker ! NetworkMessage(
+ seq,
+ WorkflowFIFOMessage(ChannelIdentity(CONTROLLER, identifier1, isControl = true), seq, ctrl)
+ )
+ seq += 1
+ }
+ }
+
+ def mkWorker(expectedOutput: Iterable[TupleLike]): (ActorRef, CompletableFuture[Boolean]) = {
+ val expected = mutable.Queue.from(expectedOutput)
+ val completeStatus = new CompletableFuture[Boolean]()
+ val mockHandler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit = {
+ case Left(value) => ???
+ case Right(value) =>
+ value match {
+ case WorkflowFIFOMessage(_, _, payload) =>
+ payload match {
+ case payload: DataPayload =>
+ payload.asInstanceOf[DataFrame].frame.foreach { item =>
+ val expectedOutput = expected.dequeue()
+ if (expectedOutput != item) {
+ completeStatus.complete(false)
+ } else {
+ if (expected.isEmpty) {
+ completeStatus.complete(true)
+ }
+ }
+ }
+ case _ => //skip
+ }
+ }
+ }
+ val worker = TestActorRef(
+ new WorkflowWorker(
+ WorkerConfig(identifier1),
+ WorkerReplayInitialization(restoreConfOpt = None, faultToleranceConfOpt = None)
+ ) {
+ this.dp = new DataProcessor(
+ identifier1,
+ mockHandler,
+ inputMessageQueue = new LinkedBlockingQueue[DPInputQueueElement]()
+ )
+ this.dp.initTimerService(timerService)
+ dpThread = new DPThread(
+ actorId,
+ dp,
+ logManager,
+ inputQueue
+ )
+ }
+ )
+ val invocation = AsyncRPCClient.ControlInvocation(
+ METHOD_ADD_PARTITIONING,
+ AddPartitioningRequest(mockLink, mockPolicy),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 0
+ )
+ val addPort1 = AsyncRPCClient.ControlInvocation(
+ METHOD_ASSIGN_PORT,
+ AssignPortRequest(mockPortId, input = true, mkSchema(1).toRawSchema, List(""), List()),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 1
+ )
+ val addPort2 = AsyncRPCClient.ControlInvocation(
+ METHOD_ASSIGN_PORT,
+ AssignPortRequest(mockPortId, input = false, mkSchema(1).toRawSchema, List(""), List()),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 2
+ )
+ val addInputChannel = AsyncRPCClient.ControlInvocation(
+ METHOD_ADD_INPUT_CHANNEL,
+ AddInputChannelRequest(
+ ChannelIdentity(identifier2, identifier1, isControl = false),
+ mockLink.toPortId
+ ),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 3
+ )
+
+ val initializeOperatorLogic = AsyncRPCClient.ControlInvocation(
+ METHOD_INITIALIZE_EXECUTOR,
+ InitializeExecutorRequest(
+ 1,
+ OpExecWithClassName(
+ "org.apache.texera.amber.engine.architecture.worker.DummyOperatorExecutor"
+ ),
+ isSource = false
+ ),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 4
+ )
+ sendControlToWorker(
+ worker,
+ Array(invocation, addPort1, addPort2, addInputChannel, initializeOperatorLogic)
+ )
+ (worker, completeStatus)
+ }
+
+ "Worker" should "process data messages correctly" in {
+ val (worker, future) = mkWorker(Array(mkTuple(1)))
+ worker ! NetworkMessage(
+ 0,
+ WorkflowFIFOMessage(
+ ChannelIdentity(identifier2, identifier1, isControl = false),
+ 0,
+ DataFrame(Array(mkTuple(1)))
+ )
+ )
+ worker ! AsyncRPCClient.ControlInvocation(
+ METHOD_FLUSH_NETWORK_BUFFER,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 1
+ )
+ //wait test to finish
+ assert(future.get(3000, MILLISECONDS))
+ }
+
+ "Worker" should "process batches correctly" in {
+ ignoreMsg {
+ case a => println(a); true
+ }
+
+ def mkBatch(start: Int, end: Int): Array[Tuple] = {
+ (start until end).map { x =>
+ mkTuple(x)
+ }.toArray
+ }
+
+ val batch1 = mkBatch(0, 400)
+ val batch2 = mkBatch(400, 500)
+ val batch3 = mkBatch(500, 800)
+ val (worker, future) = mkWorker(mkBatch(0, 800))
+ worker ! NetworkMessage(
+ 3,
+ WorkflowFIFOMessage(
+ ChannelIdentity(identifier2, identifier1, isControl = false),
+ 0,
+ DataFrame(batch1)
+ )
+ )
+ worker ! NetworkMessage(
+ 2,
+ WorkflowFIFOMessage(
+ ChannelIdentity(identifier2, identifier1, isControl = false),
+ 1,
+ DataFrame(batch2)
+ )
+ )
+ Thread.sleep(1000)
+ worker ! NetworkMessage(
+ 4,
+ WorkflowFIFOMessage(
+ ChannelIdentity(identifier2, identifier1, isControl = false),
+ 2,
+ DataFrame(batch3)
+ )
+ )
+ //wait test to finish
+ assert(future.get(3000, MILLISECONDS))
+ }
+
+ "Worker" should "accept messages in fifo order" in {
+ ignoreMsg {
+ case a => println(a); true
+ }
+ val (worker, future) = mkWorker((0 until 100).map(mkTuple(_)))
+ Random
+ .shuffle((0 until 50).map { i =>
+ NetworkMessage(
+ i + 2,
+ WorkflowFIFOMessage(
+ ChannelIdentity(identifier2, identifier1, isControl = false),
+ i,
+ DataFrame(Array(mkTuple(i)))
+ )
+ )
+ })
+ .foreach { x =>
+ worker ! x
+ }
+ Thread.sleep(1000)
+ Random
+ .shuffle((50 until 100).map { i =>
+ NetworkMessage(
+ i + 2,
+ WorkflowFIFOMessage(
+ ChannelIdentity(identifier2, identifier1, isControl = false),
+ i,
+ DataFrame(Array(mkTuple(i)))
+ )
+ )
+ })
+ .foreach { x =>
+ worker ! x
+ }
+ //wait test to finish
+ assert(future.get(3000, MILLISECONDS))
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala
new file mode 100644
index 00000000000..90e8b817bea
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala
@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.{Await, Duration, Future}
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_QUERY_STATISTICS
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{
+ ActorCommandElement,
+ DPInputQueueElement,
+ FIFOMessageElement,
+ MainThreadDelegateMessage
+}
+import org.apache.texera.amber.engine.architecture.worker.{
+ DataProcessor,
+ DataProcessorRPCHandlerInitializer
+}
+import org.apache.texera.amber.engine.common.actormessage.Backpressure
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.scalatest.flatspec.AnyFlatSpec
+
+import java.util.concurrent.LinkedBlockingQueue
+
+/**
+ * `endWorker` is the controller's acknowledgement point before it sends actor-level `gracefulStop`.
+ *
+ * A successful reply means the worker has drained every queued workflow message. If the queue still contains work,
+ * the handler must fail so the region coordinator can retry the kill instead of stopping the actor too early.
+ */
+class EndHandlerSpec extends AnyFlatSpec {
+ private val workerId = ActorVirtualIdentity("Worker:WF1-test-op-main-0")
+ private val rpcContext = AsyncRPCContext(CONTROLLER, workerId)
+ private val awaitTimeout = Duration.fromSeconds(1)
+
+ private def createEndHandlerForQueue(
+ queue: LinkedBlockingQueue[DPInputQueueElement]
+ ): DataProcessorRPCHandlerInitializer = {
+ val outputHandler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit = _ => ()
+ val dp = new DataProcessor(workerId, outputHandler, queue)
+ new DataProcessorRPCHandlerInitializer(dp)
+ }
+
+ private def await[T](future: Future[T]): T = Await.result(future, awaitTimeout)
+
+ private def assertEndWorkerFails(handler: DataProcessorRPCHandlerInitializer): Unit = {
+ val exception = intercept[IllegalStateException] {
+ await(handler.endWorker(EmptyRequest(), rpcContext))
+ }
+ assert(exception.getMessage == "worker still has unprocessed messages")
+ }
+
+ private def queueWithFifoControlMessage(): LinkedBlockingQueue[DPInputQueueElement] = {
+ val queue = new LinkedBlockingQueue[DPInputQueueElement]()
+ queue.put(
+ FIFOMessageElement(
+ WorkflowFIFOMessage(
+ ChannelIdentity(CONTROLLER, workerId, isControl = true),
+ 0,
+ ControlInvocation(METHOD_QUERY_STATISTICS, EmptyRequest(), rpcContext, 1)
+ )
+ )
+ )
+ queue
+ }
+
+ private def queueWithActorCommand(): LinkedBlockingQueue[DPInputQueueElement] = {
+ val queue = new LinkedBlockingQueue[DPInputQueueElement]()
+ queue.put(ActorCommandElement(Backpressure(enableBackpressure = true)))
+ queue
+ }
+
+ "EndHandler" should "reply successfully when there are no unprocessed messages" in {
+ val handler = createEndHandlerForQueue(new LinkedBlockingQueue[DPInputQueueElement]())
+
+ assert(await(handler.endWorker(EmptyRequest(), rpcContext)) == EmptyReturn())
+ }
+
+ it should "fail when a FIFO control message is still queued" in {
+ val handler = createEndHandlerForQueue(queueWithFifoControlMessage())
+
+ assertEndWorkerFails(handler)
+ }
+
+ it should "fail when an actor command is still queued" in {
+ val handler = createEndHandlerForQueue(queueWithActorCommand())
+
+ assertEndWorkerFails(handler)
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/BatchSizePropagationSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/BatchSizePropagationSpec.scala
new file mode 100644
index 00000000000..e9b830bdfdc
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/BatchSizePropagationSpec.scala
@@ -0,0 +1,291 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.e2e
+
+import org.apache.pekko.actor.{ActorSystem, Props}
+import org.apache.pekko.testkit.{ImplicitSender, TestKit}
+import org.apache.pekko.util.Timeout
+import org.apache.texera.amber.clustering.SingleNodeListener
+import org.apache.texera.amber.core.workflow.{PortIdentity, WorkflowContext, WorkflowSettings}
+import org.apache.texera.amber.engine.architecture.controller._
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings._
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.apache.texera.amber.engine.e2e.TestUtils.buildWorkflow
+import org.apache.texera.amber.operator.TestOperators
+import org.apache.texera.amber.operator.aggregate.AggregationFunction
+import org.apache.texera.workflow.LogicalLink
+import org.scalatest.flatspec.AnyFlatSpecLike
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
+
+import scala.concurrent.duration.DurationInt
+
+class BatchSizePropagationSpec
+ extends TestKit(ActorSystem("BatchSizePropagationSpec"))
+ with ImplicitSender
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach {
+
+ implicit val timeout: Timeout = Timeout(5.seconds)
+
+ override def beforeAll(): Unit = {
+ system.actorOf(Props[SingleNodeListener](), "cluster-info")
+ }
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ def verifyBatchSizeInPartitioning(
+ workflowScheduler: WorkflowScheduler,
+ expectedBatchSize: Int
+ ): Unit = {
+ var nextRegions = workflowScheduler.getNextRegions
+ while (nextRegions.nonEmpty) {
+ nextRegions.foreach { region =>
+ region.resourceConfig.foreach { resourceConfig =>
+ resourceConfig.linkConfigs.foreach {
+ case (_, linkConfig) =>
+ val partitioning = linkConfig.partitioning
+ partitioning match {
+ case oneToOne: OneToOnePartitioning =>
+ println(s"Testing OneToOnePartitioning with batch size: ${oneToOne.batchSize}")
+ assert(
+ oneToOne.batchSize == expectedBatchSize,
+ s"Batch size mismatch: ${oneToOne.batchSize} != $expectedBatchSize"
+ )
+
+ case roundRobin: RoundRobinPartitioning =>
+ println(
+ s"Testing RoundRobinPartitioning with batch size: ${roundRobin.batchSize}"
+ )
+ assert(
+ roundRobin.batchSize == expectedBatchSize,
+ s"Batch size mismatch: ${roundRobin.batchSize} != $expectedBatchSize"
+ )
+
+ case hashBased: HashBasedShufflePartitioning =>
+ println(
+ s"Testing HashBasedShufflePartitioning with batch size: ${hashBased.batchSize}"
+ )
+ assert(
+ hashBased.batchSize == expectedBatchSize,
+ s"Batch size mismatch: ${hashBased.batchSize} != $expectedBatchSize"
+ )
+
+ case rangeBased: RangeBasedShufflePartitioning =>
+ println(
+ s"Testing RangeBasedShufflePartitioning with batch size: ${rangeBased.batchSize}"
+ )
+ assert(
+ rangeBased.batchSize == expectedBatchSize,
+ s"Batch size mismatch: ${rangeBased.batchSize} != $expectedBatchSize"
+ )
+
+ case broadcast: BroadcastPartitioning =>
+ println(s"Testing BroadcastPartitioning with batch size: ${broadcast.batchSize}")
+ assert(
+ broadcast.batchSize == expectedBatchSize,
+ s"Batch size mismatch: ${broadcast.batchSize} != $expectedBatchSize"
+ )
+
+ case _ =>
+ throw new IllegalArgumentException("Unknown partitioning type encountered")
+ }
+ }
+ }
+ }
+ nextRegions = workflowScheduler.getNextRegions
+ }
+ }
+
+ "Engine" should "propagate the correct batch size for headerlessCsv workflow" in {
+ val expectedBatchSize = 1
+
+ val customWorkflowSettings = WorkflowSettings(dataTransferBatchSize = expectedBatchSize)
+
+ val context =
+ new WorkflowContext(workflowSettings = customWorkflowSettings)
+
+ val headerlessCsvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc()
+
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc),
+ List(),
+ context
+ )
+
+ val workflowScheduler = new WorkflowScheduler(context, CONTROLLER)
+ workflowScheduler.updateSchedule(workflow.physicalPlan)
+
+ verifyBatchSizeInPartitioning(workflowScheduler, 1)
+ }
+
+ "Engine" should "propagate the correct batch size for headerlessCsv->keyword workflow" in {
+ val expectedBatchSize = 500
+
+ val customWorkflowSettings = WorkflowSettings(dataTransferBatchSize = expectedBatchSize)
+
+ val context =
+ new WorkflowContext(workflowSettings = customWorkflowSettings)
+
+ val headerlessCsvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ context
+ )
+
+ val workflowScheduler = new WorkflowScheduler(context, CONTROLLER)
+ workflowScheduler.updateSchedule(workflow.physicalPlan)
+
+ verifyBatchSizeInPartitioning(workflowScheduler, 500)
+ }
+
+ "Engine" should "propagate the correct batch size for csv->keyword->count workflow" in {
+ val expectedBatchSize = 100
+
+ val customWorkflowSettings = WorkflowSettings(dataTransferBatchSize = expectedBatchSize)
+
+ val context =
+ new WorkflowContext(workflowSettings = customWorkflowSettings)
+
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val countOpDesc =
+ TestOperators.aggregateAndGroupByDesc("Region", AggregationFunction.COUNT, List[String]())
+
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc, countOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ countOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ context
+ )
+
+ val workflowScheduler = new WorkflowScheduler(context, CONTROLLER)
+ workflowScheduler.updateSchedule(workflow.physicalPlan)
+
+ verifyBatchSizeInPartitioning(workflowScheduler, 100)
+ }
+
+ "Engine" should "propagate the correct batch size for csv->keyword->averageAndGroupBy workflow" in {
+ val expectedBatchSize = 300
+
+ val customWorkflowSettings = WorkflowSettings(dataTransferBatchSize = expectedBatchSize)
+
+ val context =
+ new WorkflowContext(workflowSettings = customWorkflowSettings)
+
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val averageAndGroupByOpDesc =
+ TestOperators.aggregateAndGroupByDesc(
+ "Units Sold",
+ AggregationFunction.AVERAGE,
+ List[String]("Country")
+ )
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc, averageAndGroupByOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ averageAndGroupByOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ context
+ )
+
+ val workflowScheduler = new WorkflowScheduler(context, CONTROLLER)
+ workflowScheduler.updateSchedule(workflow.physicalPlan)
+
+ verifyBatchSizeInPartitioning(workflowScheduler, 300)
+ }
+
+ "Engine" should "propagate the correct batch size for csv->(csv->)->join workflow" in {
+ val expectedBatchSize = 1
+
+ val customWorkflowSettings = WorkflowSettings(dataTransferBatchSize = expectedBatchSize)
+
+ val context =
+ new WorkflowContext(workflowSettings = customWorkflowSettings)
+
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val headerlessCsvOpDesc2 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ headerlessCsvOpDesc2,
+ joinOpDesc
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc2.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ )
+ ),
+ context
+ )
+
+ val workflowScheduler = new WorkflowScheduler(context, CONTROLLER)
+ workflowScheduler.updateSchedule(workflow.physicalPlan)
+
+ verifyBatchSizeInPartitioning(workflowScheduler, 1)
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/DataProcessingSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/DataProcessingSpec.scala
new file mode 100644
index 00000000000..69ee9c6a5fb
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/DataProcessingSpec.scala
@@ -0,0 +1,486 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.e2e
+
+import org.apache.pekko.actor.{ActorSystem, Props}
+import org.apache.pekko.testkit.{ImplicitSender, TestKit}
+import org.apache.pekko.util.Timeout
+import com.twitter.util.{Await, Duration, Promise}
+import org.apache.texera.amber.clustering.SingleNodeListener
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.storage.model.VirtualDocument
+import org.apache.texera.amber.core.tuple.{AttributeType, Tuple}
+import org.apache.texera.amber.core.virtualidentity.OperatorIdentity
+import org.apache.texera.amber.core.workflow.{
+ ExecutionMode,
+ PortIdentity,
+ WorkflowContext,
+ WorkflowSettings
+}
+import org.apache.texera.amber.engine.architecture.controller._
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmptyRequest
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.COMPLETED
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.engine.e2e.TestUtils.{
+ buildWorkflow,
+ cleanupWorkflowExecutionData,
+ initiateTexeraDBForTestCases,
+ setUpWorkflowExecutionData
+}
+import org.apache.texera.amber.operator.TestOperators
+import org.apache.texera.amber.operator.aggregate.AggregationFunction
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource.getResultUriByLogicalPortId
+import org.apache.texera.workflow.LogicalLink
+import org.scalatest.flatspec.AnyFlatSpecLike
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach, Outcome, Retries}
+
+import scala.concurrent.duration.DurationInt
+
+class DataProcessingSpec
+ extends TestKit(ActorSystem("DataProcessingSpec", AmberRuntime.akkaConfig))
+ with ImplicitSender
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with Retries {
+
+ /**
+ * This block retries each test once if it fails.
+ * In the CI environment, there is a chance that executeWorkflow does not receive "COMPLETED" status.
+ * Until we find the root cause of this issue, we use a retry mechanism here to stablize CI runs.
+ */
+ override def withFixture(test: NoArgTest): Outcome =
+ withRetry { super.withFixture(test) }
+
+ implicit val timeout: Timeout = Timeout(5.seconds)
+
+ val workflowContext: WorkflowContext = new WorkflowContext()
+
+ val materializedWorkflowContext: WorkflowContext = new WorkflowContext(
+ workflowSettings = WorkflowSettings(
+ dataTransferBatchSize = 400,
+ executionMode = ExecutionMode.MATERIALIZED
+ )
+ )
+
+ override protected def beforeEach(): Unit = {
+ setUpWorkflowExecutionData()
+ }
+
+ override protected def afterEach(): Unit = {
+ cleanupWorkflowExecutionData()
+ }
+
+ override def beforeAll(): Unit = {
+ system.actorOf(Props[SingleNodeListener](), "cluster-info")
+ // These test cases access postgres in CI, but occasionally the jdbc driver cannot be found during CI run.
+ // Explicitly load the JDBC driver to avoid flaky CI failures.
+ Class.forName("org.postgresql.Driver")
+ initiateTexeraDBForTestCases()
+ }
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ def executeWorkflow(workflow: Workflow): Map[OperatorIdentity, List[Tuple]] = {
+ var results: Map[OperatorIdentity, List[Tuple]] = null
+ val client = new AmberClient(
+ system,
+ workflow.context,
+ workflow.physicalPlan,
+ ControllerConfig.default,
+ error => {}
+ )
+ val completion = Promise[Unit]()
+ client.registerCallback[FatalError](evt => {
+ completion.setException(evt.e)
+ client.shutdown()
+ })
+
+ client
+ .registerCallback[ExecutionStateUpdate](evt => {
+ if (evt.state == COMPLETED) {
+ results = workflow.logicalPlan.getTerminalOperatorIds
+ .filter(terminalOpId => {
+ val uri = getResultUriByLogicalPortId(
+ workflowContext.executionId,
+ terminalOpId,
+ PortIdentity()
+ )
+ uri.nonEmpty
+ })
+ .map(terminalOpId => {
+ //TODO: remove the delay after fixing the issue of reporting "completed" status too early.
+ Thread.sleep(1000)
+ val uri = getResultUriByLogicalPortId(
+ workflowContext.executionId,
+ terminalOpId,
+ PortIdentity()
+ ).get
+ terminalOpId -> DocumentFactory
+ .openDocument(uri)
+ ._1
+ .asInstanceOf[VirtualDocument[Tuple]]
+ .get()
+ .toList
+ })
+ .toMap
+ completion.setDone()
+ }
+ })
+ Await.result(client.controllerInterface.startWorkflow(EmptyRequest(), ()))
+ Await.result(completion, Duration.fromMinutes(1))
+ results
+ }
+
+ "Engine" should "execute headerlessCsv workflow normally" in {
+ val headerlessCsvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc()
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc),
+ List(),
+ workflowContext
+ )
+ val results = executeWorkflow(workflow)(headerlessCsvOpDesc.operatorIdentifier)
+
+ assert(results.size == 100)
+ }
+
+ "Engine" should "execute headerlessMultiLineDataCsv workflow normally" in {
+ val headerlessCsvOpDesc = TestOperators.headerlessSmallMultiLineDataCsvScanOpDesc()
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc),
+ List(),
+ workflowContext
+ )
+ val results = executeWorkflow(workflow)(headerlessCsvOpDesc.operatorIdentifier)
+
+ assert(results.size == 100)
+ }
+
+ "Engine" should "execute jsonl workflow normally" in {
+ val jsonlOp = TestOperators.smallJSONLScanOpDesc()
+ val workflow = buildWorkflow(
+ List(jsonlOp),
+ List(),
+ workflowContext
+ )
+ val results = executeWorkflow(workflow)(jsonlOp.operatorIdentifier)
+
+ assert(results.size == 100)
+
+ for (result <- results) {
+ val schema = result.getSchema
+ assert(schema.getAttribute("id").getType == AttributeType.LONG)
+ assert(schema.getAttribute("first_name").getType == AttributeType.STRING)
+ assert(schema.getAttribute("flagged").getType == AttributeType.BOOLEAN)
+ assert(schema.getAttribute("year").getType == AttributeType.INTEGER)
+ assert(schema.getAttribute("created_at").getType == AttributeType.TIMESTAMP)
+ assert(schema.getAttributes.length == 9)
+ }
+
+ }
+
+ "Engine" should "execute mediumFlattenJsonl workflow normally" in {
+ val jsonlOp = TestOperators.mediumFlattenJSONLScanOpDesc()
+ val workflow = buildWorkflow(
+ List(jsonlOp),
+ List(),
+ workflowContext
+ )
+ val results = executeWorkflow(workflow)(jsonlOp.operatorIdentifier)
+
+ assert(results.size == 1000)
+
+ for (result <- results) {
+ val schema = result.getSchema
+ assert(schema.getAttribute("id").getType == AttributeType.LONG)
+ assert(schema.getAttribute("first_name").getType == AttributeType.STRING)
+ assert(schema.getAttribute("flagged").getType == AttributeType.BOOLEAN)
+ assert(schema.getAttribute("year").getType == AttributeType.INTEGER)
+ assert(schema.getAttribute("created_at").getType == AttributeType.TIMESTAMP)
+ assert(schema.getAttribute("test_object.array2.another").getType == AttributeType.INTEGER)
+ assert(schema.getAttributes.length == 13)
+ }
+ }
+
+ "Engine" should "execute headerlessCsv->keyword workflow normally" in {
+ val headerlessCsvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ workflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv workflow normally" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val workflow = buildWorkflow(
+ List(csvOpDesc),
+ List(),
+ workflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv->keyword workflow normally" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ workflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv->keyword->count workflow normally" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val countOpDesc =
+ TestOperators.aggregateAndGroupByDesc("Region", AggregationFunction.COUNT, List[String]())
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc, countOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ countOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ workflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv->keyword->averageAndGroupBy workflow normally" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val averageAndGroupByOpDesc =
+ TestOperators.aggregateAndGroupByDesc(
+ "Units Sold",
+ AggregationFunction.AVERAGE,
+ List[String]("Country")
+ )
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc, averageAndGroupByOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ averageAndGroupByOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ workflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv->(csv->)->join workflow normally" in {
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val headerlessCsvOpDesc2 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ headerlessCsvOpDesc2,
+ joinOpDesc
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc2.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ )
+ ),
+ workflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute headerlessCsv->keyword workflow with MATERIALIZED mode" in {
+ val headerlessCsvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("column-1", "Asia")
+ val workflow = buildWorkflow(
+ List(headerlessCsvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ materializedWorkflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv workflow with MATERIALIZED mode" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val workflow = buildWorkflow(
+ List(csvOpDesc),
+ List(),
+ materializedWorkflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv->keyword workflow with MATERIALIZED mode" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ materializedWorkflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv->keyword->count workflow with MATERIALIZED mode" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val countOpDesc =
+ TestOperators.aggregateAndGroupByDesc("Region", AggregationFunction.COUNT, List[String]())
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc, countOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ countOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ materializedWorkflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv->keyword->averageAndGroupBy workflow with MATERIALIZED mode" in {
+ val csvOpDesc = TestOperators.smallCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val averageAndGroupByOpDesc =
+ TestOperators.aggregateAndGroupByDesc(
+ "Units Sold",
+ AggregationFunction.AVERAGE,
+ List[String]("Country")
+ )
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc, averageAndGroupByOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity(),
+ averageAndGroupByOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ materializedWorkflowContext
+ )
+ executeWorkflow(workflow)
+ }
+
+ "Engine" should "execute csv->(csv->)->join workflow with MATERIALIZED mode" in {
+ val headerlessCsvOpDesc1 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val headerlessCsvOpDesc2 = TestOperators.headerlessSmallCsvScanOpDesc()
+ val joinOpDesc = TestOperators.joinOpDesc("column-1", "column-1")
+ val workflow = buildWorkflow(
+ List(
+ headerlessCsvOpDesc1,
+ headerlessCsvOpDesc2,
+ joinOpDesc
+ ),
+ List(
+ LogicalLink(
+ headerlessCsvOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ headerlessCsvOpDesc2.operatorIdentifier,
+ PortIdentity(),
+ joinOpDesc.operatorIdentifier,
+ PortIdentity(1)
+ )
+ ),
+ materializedWorkflowContext
+ )
+ executeWorkflow(workflow)
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/PauseSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/PauseSpec.scala
new file mode 100644
index 00000000000..b459533c573
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/PauseSpec.scala
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.e2e
+
+import org.apache.pekko.actor.{ActorSystem, Props}
+import org.apache.pekko.testkit.{ImplicitSender, TestKit}
+import org.apache.pekko.util.Timeout
+import com.twitter.util.{Await, Duration, Promise}
+import com.typesafe.scalalogging.Logger
+import org.apache.texera.amber.clustering.SingleNodeListener
+import org.apache.texera.amber.core.workflow.{PortIdentity, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerConfig,
+ ExecutionStateUpdate
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmptyRequest
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{
+ COMPLETED,
+ PAUSED
+}
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.engine.e2e.TestUtils.{
+ cleanupWorkflowExecutionData,
+ initiateTexeraDBForTestCases,
+ setUpWorkflowExecutionData,
+ stateReached
+}
+import org.apache.texera.amber.operator.{LogicalOp, TestOperators}
+import org.apache.texera.workflow.LogicalLink
+import org.scalatest.flatspec.AnyFlatSpecLike
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach, Outcome, Retries}
+
+import scala.concurrent.duration._
+
+class PauseSpec
+ extends TestKit(ActorSystem("PauseSpec", AmberRuntime.akkaConfig))
+ with ImplicitSender
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with Retries {
+
+ /**
+ * This block retries each test once if it fails.
+ * In the CI environment, there is a chance that shouldPause does not receive "COMPLETED" status.
+ * Until we find the root cause of this issue, we use a retry mechanism here to stablize CI runs.
+ */
+ override def withFixture(test: NoArgTest): Outcome =
+ withRetry { super.withFixture(test) }
+
+ implicit val timeout: Timeout = Timeout(5.seconds)
+
+ val logger = Logger("PauseSpecLogger")
+
+ override protected def beforeEach(): Unit = {
+ setUpWorkflowExecutionData()
+ }
+
+ override protected def afterEach(): Unit = {
+ cleanupWorkflowExecutionData()
+ }
+
+ override def beforeAll(): Unit = {
+ system.actorOf(Props[SingleNodeListener](), "cluster-info")
+ // These test cases access postgres in CI, but occasionally the jdbc driver cannot be found during CI run.
+ // Explicitly load the JDBC driver to avoid flaky CI failures.
+ Class.forName("org.postgresql.Driver")
+ initiateTexeraDBForTestCases()
+ }
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ def shouldPause(
+ operators: List[LogicalOp],
+ links: List[LogicalLink]
+ ): Unit = {
+ val workflow =
+ TestUtils.buildWorkflow(operators, links, new WorkflowContext())
+ val client =
+ new AmberClient(
+ system,
+ workflow.context,
+ workflow.physicalPlan,
+ ControllerConfig.default,
+ error => {}
+ )
+ val completion = Promise[Unit]()
+ client
+ .registerCallback[ExecutionStateUpdate](evt => {
+ if (evt.state == COMPLETED) {
+ completion.setDone()
+ }
+ })
+ val stateWaitTimeout = Duration.fromSeconds(10)
+ Await.result(client.controllerInterface.startWorkflow(EmptyRequest(), ()))
+ val firstPaused = stateReached(client, PAUSED)
+ Await.result(client.controllerInterface.pauseWorkflow(EmptyRequest(), ()))
+ Await.result(firstPaused, stateWaitTimeout)
+ Await.result(client.controllerInterface.resumeWorkflow(EmptyRequest(), ()))
+ val secondPaused = stateReached(client, PAUSED)
+ Await.result(client.controllerInterface.pauseWorkflow(EmptyRequest(), ()))
+ Await.result(secondPaused, stateWaitTimeout)
+ Await.result(client.controllerInterface.resumeWorkflow(EmptyRequest(), ()))
+ Await.result(completion, Duration.fromMinutes(1))
+ }
+
+ "Engine" should "be able to pause csv workflow" in {
+ val csvOpDesc = TestOperators.mediumCsvScanOpDesc()
+ logger.info(s"csv-id ${csvOpDesc.operatorIdentifier}")
+ shouldPause(
+ List(csvOpDesc),
+ List()
+ )
+ }
+
+ "Engine" should "be able to pause csv->keyword workflow" in {
+ val csvOpDesc = TestOperators.mediumCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ logger.info(
+ s"csv-id ${csvOpDesc.operatorIdentifier}, keyword-id ${keywordOpDesc.operatorIdentifier}"
+ )
+ shouldPause(
+ List(csvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ )
+ )
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/ReconfigurationSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/ReconfigurationSpec.scala
new file mode 100644
index 00000000000..92dfba19de1
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/ReconfigurationSpec.scala
@@ -0,0 +1,327 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.e2e
+
+import com.twitter.util.{Await, Duration, Promise}
+import com.typesafe.scalalogging.Logger
+import org.apache.pekko.actor.{ActorSystem, Props}
+import org.apache.pekko.testkit.{ImplicitSender, TestKit}
+import org.apache.pekko.util.Timeout
+import org.apache.texera.amber.clustering.SingleNodeListener
+import org.apache.texera.amber.core.executor.{OpExecInitInfo, OpExecWithCode}
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.storage.model.VirtualDocument
+import org.apache.texera.amber.core.tuple.Tuple
+import org.apache.texera.amber.core.virtualidentity.OperatorIdentity
+import org.apache.texera.amber.core.workflow.{PortIdentity, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerConfig,
+ ExecutionStateUpdate
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ EmptyRequest,
+ UpdateExecutorRequest,
+ WorkflowReconfigureRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{
+ COMPLETED,
+ PAUSED
+}
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.engine.e2e.TestUtils.{
+ cleanupWorkflowExecutionData,
+ initiateTexeraDBForTestCases,
+ setUpWorkflowExecutionData,
+ stateReached
+}
+import org.apache.texera.amber.operator.{LogicalOp, TestOperators}
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowExecutionsResource.getResultUriByLogicalPortId
+import org.apache.texera.workflow.LogicalLink
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach, Outcome, Retries}
+import org.scalatest.flatspec.AnyFlatSpecLike
+
+import scala.concurrent.duration._
+
+class ReconfigurationSpec
+ extends TestKit(ActorSystem("ReconfigurationSpec", AmberRuntime.akkaConfig))
+ with ImplicitSender
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with Retries {
+
+ /**
+ * This block retries each test once if it fails.
+ * In the CI environment, there is a chance that executeWorkflow does not receive "COMPLETED" status.
+ * Until we find the root cause of this issue, we use a retry mechanism here to stabilize CI runs.
+ */
+ override def withFixture(test: NoArgTest): Outcome =
+ withRetry { super.withFixture(test) }
+
+ implicit val timeout: Timeout = Timeout(5.seconds)
+
+ val logger = Logger("ReconfigurationSpecLogger")
+ val ctx = new WorkflowContext()
+
+ override protected def beforeEach(): Unit = {
+ setUpWorkflowExecutionData()
+ }
+
+ override protected def afterEach(): Unit = {
+ cleanupWorkflowExecutionData()
+ }
+
+ override def beforeAll(): Unit = {
+ system.actorOf(Props[SingleNodeListener](), "cluster-info")
+ // These test cases access postgres in CI, but occasionally the jdbc driver cannot be found during CI run.
+ // Explicitly load the JDBC driver to avoid flaky CI failures.
+ Class.forName("org.postgresql.Driver")
+ initiateTexeraDBForTestCases()
+ }
+
+ override def afterAll(): Unit = {
+ TestKit.shutdownActorSystem(system)
+ }
+
+ def shouldReconfigure(
+ operators: List[LogicalOp],
+ links: List[LogicalLink],
+ targetOps: Seq[LogicalOp],
+ newOpExecInitInfo: OpExecInitInfo
+ ): Map[OperatorIdentity, List[Tuple]] = {
+ val workflow =
+ TestUtils.buildWorkflow(operators, links, ctx)
+ val client =
+ new AmberClient(
+ system,
+ workflow.context,
+ workflow.physicalPlan,
+ ControllerConfig.default,
+ error => {}
+ )
+ val completion = Promise[Unit]()
+ var result: Map[OperatorIdentity, List[Tuple]] = null
+ client
+ .registerCallback[ExecutionStateUpdate](evt => {
+ if (evt.state == COMPLETED) {
+ result = workflow.logicalPlan.getTerminalOperatorIds
+ .filter(terminalOpId => {
+ val uri = getResultUriByLogicalPortId(
+ workflow.context.executionId,
+ terminalOpId,
+ PortIdentity()
+ )
+ uri.nonEmpty
+ })
+ .map(terminalOpId => {
+ //TODO: remove the delay after fixing the issue of reporting "completed" status too early.
+ Thread.sleep(1000)
+ val uri = getResultUriByLogicalPortId(
+ workflow.context.executionId,
+ terminalOpId,
+ PortIdentity()
+ ).get
+ terminalOpId -> DocumentFactory
+ .openDocument(uri)
+ ._1
+ .asInstanceOf[VirtualDocument[Tuple]]
+ .get()
+ .toList
+ })
+ .toMap
+ completion.setDone()
+ }
+ })
+ Await.result(
+ client.controllerInterface.startWorkflow(EmptyRequest(), ()),
+ Duration.fromSeconds(5)
+ )
+ val pausedReached = stateReached(client, PAUSED)
+ Await.result(
+ client.controllerInterface.pauseWorkflow(EmptyRequest(), ()),
+ Duration.fromSeconds(5)
+ )
+ Await.result(pausedReached, Duration.fromSeconds(10))
+ val physicalOps = targetOps.flatMap(op =>
+ workflow.physicalPlan.getPhysicalOpsOfLogicalOp(op.operatorIdentifier)
+ )
+ Await.result(
+ client.controllerInterface.reconfigureWorkflow(
+ WorkflowReconfigureRequest(
+ reconfiguration = physicalOps.map(op => UpdateExecutorRequest(op.id, newOpExecInitInfo)),
+ reconfigurationId = "test-reconfigure-1"
+ ),
+ ()
+ ),
+ Duration.fromSeconds(5)
+ )
+ Await.result(
+ client.controllerInterface.resumeWorkflow(EmptyRequest(), ()),
+ Duration.fromSeconds(5)
+ )
+ Await.result(completion, Duration.fromMinutes(1))
+ result
+ }
+
+ "Engine" should "be able to modify a python UDF worker in workflow" in {
+ val sourceOpDesc = TestOperators.smallCsvScanOpDesc()
+ val udfOpDesc = TestOperators.pythonOpDesc()
+ val code = """
+ |from pytexera import *
+ |
+ |class ProcessTupleOperator(UDFOperatorV2):
+ | @overrides
+ | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ | tuple_['Region'] = tuple_['Region'] + '_reconfigured'
+ | yield tuple_
+ |""".stripMargin
+
+ val result = shouldReconfigure(
+ List(sourceOpDesc, udfOpDesc),
+ List(
+ LogicalLink(
+ sourceOpDesc.operatorIdentifier,
+ PortIdentity(),
+ udfOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ Seq(udfOpDesc),
+ OpExecWithCode(code, "python")
+ )
+ assert(result(udfOpDesc.operatorIdentifier).exists { t =>
+ t.getField("Region").asInstanceOf[String].contains("_reconfigured")
+ })
+ }
+
+ "Engine" should "be able to modify a java operator in workflow" in {
+ val sourceOpDesc = TestOperators.mediumCsvScanOpDesc()
+ val keywordMatchNoneOpDesc = TestOperators.keywordSearchOpDesc("Region", "ShouldMatchNone")
+ val keywordMatchManyOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val result = shouldReconfigure(
+ List(sourceOpDesc, keywordMatchNoneOpDesc),
+ List(
+ LogicalLink(
+ sourceOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordMatchNoneOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ Seq(keywordMatchNoneOpDesc),
+ keywordMatchManyOpDesc.getPhysicalOp(ctx.workflowId, ctx.executionId).opExecInitInfo
+ )
+ assert(result(keywordMatchNoneOpDesc.operatorIdentifier).nonEmpty)
+ }
+
+ "Engine" should "not be able to modify a source operator in workflow" in {
+ val sourceOpDesc = TestOperators.mediumCsvScanOpDesc()
+ val sourceOpDesc2 = TestOperators.mediumCsvScanOpDesc()
+ val keywordMatchNoneOpDesc = TestOperators.keywordSearchOpDesc("Region", "ShouldMatchNone")
+ val ex = intercept[Throwable] {
+ shouldReconfigure(
+ List(sourceOpDesc, keywordMatchNoneOpDesc),
+ List(
+ LogicalLink(
+ sourceOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordMatchNoneOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ Seq(sourceOpDesc),
+ sourceOpDesc2.getPhysicalOp(ctx.workflowId, ctx.executionId).opExecInitInfo
+ )
+ }
+ assert(
+ ex.getMessage == "java.lang.IllegalStateException: Reconfiguration cannot be applied to source operators"
+ )
+ }
+
+ "Engine" should "propagate reconfiguration through a source operator in workflow" in {
+ val sourceOpDesc = TestOperators.pythonSourceOpDesc(10000)
+ val udfOpDesc = TestOperators.pythonOpDesc()
+ val code = """
+ |from pytexera import *
+ |
+ |class ProcessTupleOperator(UDFOperatorV2):
+ | @overrides
+ | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ | tuple_['field_1'] = tuple_['field_1'] + '_reconfigured'
+ | yield tuple_
+ |""".stripMargin
+ val result = shouldReconfigure(
+ List(sourceOpDesc, udfOpDesc),
+ List(
+ LogicalLink(
+ sourceOpDesc.operatorIdentifier,
+ PortIdentity(),
+ udfOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ Seq(udfOpDesc),
+ OpExecWithCode(code, "python")
+ )
+ assert(result(udfOpDesc.operatorIdentifier).exists { t =>
+ t.getField("field_1").asInstanceOf[String].contains("_reconfigured")
+ })
+ }
+
+ "Engine" should "be able to modify two python UDFs in workflow" in {
+ val sourceOpDesc = TestOperators.smallCsvScanOpDesc()
+ val udfOpDesc1 = TestOperators.pythonOpDesc()
+ val udfOpDesc2 = TestOperators.pythonOpDesc()
+ val code = """
+ |from pytexera import *
+ |
+ |class ProcessTupleOperator(UDFOperatorV2):
+ | @overrides
+ | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+ | tuple_['Region'] = tuple_['Region'] + '_reconfigured'
+ | yield tuple_
+ |""".stripMargin
+
+ val result = shouldReconfigure(
+ List(sourceOpDesc, udfOpDesc1, udfOpDesc2),
+ List(
+ LogicalLink(
+ sourceOpDesc.operatorIdentifier,
+ PortIdentity(),
+ udfOpDesc1.operatorIdentifier,
+ PortIdentity()
+ ),
+ LogicalLink(
+ udfOpDesc1.operatorIdentifier,
+ PortIdentity(),
+ udfOpDesc2.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ Seq(udfOpDesc1, udfOpDesc2),
+ OpExecWithCode(code, "python")
+ )
+ assert(result(udfOpDesc2.operatorIdentifier).exists { t =>
+ t.getField("Region").asInstanceOf[String].contains("_reconfigured_reconfigured")
+ })
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala
new file mode 100644
index 00000000000..fab5d5a16c1
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.e2e
+
+import com.twitter.util.{Promise, Return}
+import org.apache.texera.amber.config.StorageConfig
+import org.apache.texera.amber.core.workflow.WorkflowContext
+import org.apache.texera.amber.engine.architecture.controller.{ExecutionStateUpdate, Workflow}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
+import org.apache.texera.amber.engine.common.client.AmberClient
+import org.apache.texera.amber.operator.LogicalOp
+import org.apache.texera.dao.SqlServer
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.{
+ UserDao,
+ WorkflowDao,
+ WorkflowExecutionsDao,
+ WorkflowVersionDao
+}
+import org.apache.texera.dao.jooq.generated.tables.pojos.{
+ User,
+ WorkflowExecutions,
+ WorkflowVersion,
+ Workflow => WorkflowPojo
+}
+import org.apache.texera.web.model.websocket.request.LogicalPlanPojo
+import org.apache.texera.workflow.{LogicalLink, WorkflowCompiler}
+
+object TestUtils {
+
+ def buildWorkflow(
+ operators: List[LogicalOp],
+ links: List[LogicalLink],
+ context: WorkflowContext
+ ): Workflow = {
+ val workflowCompiler = new WorkflowCompiler(
+ context
+ )
+ workflowCompiler.compile(
+ LogicalPlanPojo(operators, links, List(), List())
+ )
+ }
+
+ /**
+ * If a test case accesses the user system through singleton resources that cache the DSLContext (e.g., executes a
+ * workflow, which accesses WorkflowExecutionsResource), we use a separate texera_db specifically for such test cases.
+ * Note such test cases need to clean up the database at the end of running each test case.
+ */
+ def initiateTexeraDBForTestCases(): Unit = {
+ SqlServer.initConnection(
+ StorageConfig.jdbcUrlForTestCases,
+ StorageConfig.jdbcUsername,
+ StorageConfig.jdbcPassword
+ )
+ }
+
+ val testUser: User = {
+ val user = new User
+ user.setUid(Integer.valueOf(1))
+ user.setName("test_user")
+ user.setRole(UserRoleEnum.ADMIN)
+ user.setPassword("123")
+ user.setEmail("test_user@test.com")
+ user
+ }
+
+ val testWorkflowEntry: WorkflowPojo = {
+ val workflow = new WorkflowPojo
+ workflow.setName("test workflow")
+ workflow.setWid(Integer.valueOf(1))
+ workflow.setContent("test workflow content")
+ workflow.setDescription("test description")
+ workflow
+ }
+
+ val testWorkflowVersionEntry: WorkflowVersion = {
+ val workflowVersion = new WorkflowVersion
+ workflowVersion.setWid(Integer.valueOf(1))
+ workflowVersion.setVid(Integer.valueOf(1))
+ workflowVersion.setContent("test version content")
+ workflowVersion
+ }
+
+ val testWorkflowExecutionEntry: WorkflowExecutions = {
+ val workflowExecution = new WorkflowExecutions
+ workflowExecution.setEid(Integer.valueOf(1))
+ workflowExecution.setVid(Integer.valueOf(1))
+ workflowExecution.setUid(Integer.valueOf(1))
+ workflowExecution.setStatus(3.toByte)
+ workflowExecution.setEnvironmentVersion("test engine")
+ workflowExecution
+ }
+
+ def setUpWorkflowExecutionData(): Unit = {
+ val dslConfig = SqlServer.getInstance().context.configuration()
+ val userDao = new UserDao(dslConfig)
+ val workflowDao = new WorkflowDao(dslConfig)
+ val workflowExecutionsDao = new WorkflowExecutionsDao(dslConfig)
+ val workflowVersionDao = new WorkflowVersionDao(dslConfig)
+ userDao.insert(testUser)
+ workflowDao.insert(testWorkflowEntry)
+ workflowVersionDao.insert(testWorkflowVersionEntry)
+ workflowExecutionsDao.insert(testWorkflowExecutionEntry)
+ }
+
+ /**
+ * Returns a Promise that completes the next time the client emits an
+ * ExecutionStateUpdate with the given target state. Must be called BEFORE
+ * the action that triggers the state change, since AmberClient observables
+ * do not replay past events.
+ */
+ def stateReached(
+ client: AmberClient,
+ target: WorkflowAggregatedState
+ ): Promise[Unit] = {
+ val p = Promise[Unit]()
+ client.registerCallback[ExecutionStateUpdate](evt => {
+ if (evt.state == target) {
+ p.updateIfEmpty(Return(()))
+ }
+ })
+ p
+ }
+
+ def cleanupWorkflowExecutionData(): Unit = {
+ val dslConfig = SqlServer.getInstance().context.configuration()
+ val userDao = new UserDao(dslConfig)
+ val workflowDao = new WorkflowDao(dslConfig)
+ val workflowExecutionsDao = new WorkflowExecutionsDao(dslConfig)
+ val workflowVersionDao = new WorkflowVersionDao(dslConfig)
+ workflowExecutionsDao.deleteById(1)
+ workflowVersionDao.deleteById(1)
+ workflowDao.deleteById(1)
+ userDao.deleteById(1)
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/CheckpointSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/CheckpointSpec.scala
new file mode 100644
index 00000000000..ee5a0b9a609
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/CheckpointSpec.scala
@@ -0,0 +1,162 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.faulttolerance
+
+import org.apache.pekko.actor.{ActorSystem, Props}
+import org.apache.texera.amber.clustering.SingleNodeListener
+import org.apache.texera.amber.core.workflow.{PortIdentity, WorkflowContext}
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerConfig,
+ ControllerProcessor
+}
+import org.apache.texera.amber.engine.architecture.worker.DataProcessor
+import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.DPInputQueueElement
+import org.apache.texera.amber.engine.common.SerializedState.{CP_STATE_KEY, DP_STATE_KEY}
+import org.apache.texera.amber.engine.common.virtualidentity.util.{CONTROLLER, SELF}
+import org.apache.texera.amber.engine.common.{AmberRuntime, CheckpointState}
+import org.apache.texera.amber.engine.e2e.TestUtils.buildWorkflow
+import org.apache.texera.amber.operator.TestOperators
+import org.apache.texera.workflow.LogicalLink
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpecLike
+
+import java.util.concurrent.LinkedBlockingQueue
+
+class CheckpointSpec extends AnyFlatSpecLike with BeforeAndAfterAll {
+
+ var system: ActorSystem = _
+
+ val csvOpDesc = TestOperators.mediumCsvScanOpDesc()
+ val keywordOpDesc = TestOperators.keywordSearchOpDesc("Region", "Asia")
+ val workflow = buildWorkflow(
+ List(csvOpDesc, keywordOpDesc),
+ List(
+ LogicalLink(
+ csvOpDesc.operatorIdentifier,
+ PortIdentity(),
+ keywordOpDesc.operatorIdentifier,
+ PortIdentity()
+ )
+ ),
+ new WorkflowContext()
+ )
+
+ override def beforeAll(): Unit = {
+ system = ActorSystem("CheckpointSpec", AmberRuntime.akkaConfig)
+ system.actorOf(Props[SingleNodeListener](), "cluster-info")
+ }
+
+ "Default controller state" should "be serializable" in {
+ val cp =
+ new ControllerProcessor(
+ workflow.context,
+ ControllerConfig.default,
+ CONTROLLER,
+ msg => {}
+ )
+ val chkpt = new CheckpointState()
+ chkpt.save(CP_STATE_KEY, cp)
+ }
+
+ "Default worker state" should "be serializable" in {
+ val dp = new DataProcessor(
+ SELF,
+ msg => {},
+ inputMessageQueue = new LinkedBlockingQueue[DPInputQueueElement]()
+ )
+ val chkpt = new CheckpointState()
+ chkpt.save(DP_STATE_KEY, dp)
+ }
+
+// "CSVScanOperator" should "be serializable" in {
+// val chkpt = new CheckpointState()
+// val headerlessCsvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc()
+// val context = new WorkflowContext()
+// headerlessCsvOpDesc.setContext(context)
+// val phyOp = headerlessCsvOpDesc.getPhysicalOp(WorkflowIdentity(1), ExecutionIdentity(1))
+// phyOp.opExecInitInfo match {
+// case OpExecInitInfoWithCode(codeGen) => ???
+// case OpExecInitInfoWithFunc(opGen) =>
+// val operator = opGen(1, 1)
+// operator.open()
+// val outputIter =
+// operator.asInstanceOf[SourceOperatorExecutor].produceTuple().map(t => (t, None))
+// outputIter.next()
+// outputIter.next()
+// operator.asInstanceOf[CheckpointSupport].serializeState(outputIter, chkpt)
+// chkpt.save("deserialization", opGen)
+// val opGen2 = chkpt.load("deserialization").asInstanceOf[(Int, Int) => OperatorExecutor]
+// val op = opGen2.apply(1, 1)
+// op.asInstanceOf[CheckpointSupport].deserializeState(chkpt)
+// }
+// }
+//
+// "Workflow " should "take global checkpoint, reload and continue" in {
+// val client1 = new AmberClient(
+// system,
+// workflow.context,
+// workflow.physicalPlan,
+// resultStorage,
+// ControllerConfig.default,
+// error => {}
+// )
+// Await.result(client1.controllerInterface.startWorkflow(EmptyRequest(), ()))
+// Thread.sleep(100)
+// Await.result(client1.controllerInterface.pauseWorkflow(EmptyRequest(), ()))
+// val checkpointId = EmbeddedControlMessageIdentity(s"Checkpoint_test_1")
+// val uri = new URI("ram:///recovery-logs/tmp/")
+// Await.result(
+// client1.controllerInterface.takeGlobalCheckpoint(
+// TakeGlobalCheckpointRequest(estimationOnly = false, checkpointId, uri.toString),
+// ()
+// ),
+// Duration.fromSeconds(30)
+// )
+// client1.shutdown()
+// Thread.sleep(100)
+// var controllerConfig = ControllerConfig.default
+// controllerConfig =
+// controllerConfig.copy(stateRestoreConfOpt = Some(StateRestoreConfig(uri, checkpointId)))
+// val completableFuture = new CompletableFuture[Unit]()
+// val client2 = new AmberClient(
+// system,
+// workflow.context,
+// workflow.physicalPlan,
+// resultStorage,
+// controllerConfig,
+// error => {}
+// )
+// client2.registerCallback[ExecutionStateUpdate] { evt =>
+// if (evt.state == COMPLETED) {
+// completableFuture.complete(())
+// }
+// }
+// Thread.sleep(1000)
+// assert(
+// Await
+// .result(client2.controllerInterface.startWorkflow(EmptyRequest(), ()))
+// .workflowState == PAUSED
+// )
+// Thread.sleep(5000)
+// Await.result(client2.controllerInterface.resumeWorkflow(EmptyRequest(), ()))
+// completableFuture.get(30000, TimeUnit.MILLISECONDS)
+// }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/LoggingSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/LoggingSpec.scala
new file mode 100644
index 00000000000..87e3ca148ee
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/LoggingSpec.scala
@@ -0,0 +1,152 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.faulttolerance
+
+import org.apache.pekko.actor.ActorSystem
+import org.apache.pekko.testkit.{ImplicitSender, TestKit}
+import org.apache.texera.amber.core.tuple.{AttributeType, Schema, TupleLike}
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity
+}
+import org.apache.texera.amber.core.workflow.{PhysicalLink, PortIdentity}
+import org.apache.texera.amber.engine.architecture.logreplay.{ReplayLogManager, ReplayLogRecord}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AddPartitioningRequest,
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controllerservice.ControllerServiceGrpc.METHOD_WORKER_EXECUTION_COMPLETED
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.{
+ METHOD_ADD_PARTITIONING,
+ METHOD_PAUSE_WORKER,
+ METHOD_RESUME_WORKER,
+ METHOD_START_WORKER
+}
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.OneToOnePartitioning
+import org.apache.texera.amber.engine.common.AmberRuntime
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataFrame,
+ WorkflowFIFOMessage,
+ WorkflowFIFOMessagePayload
+}
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.engine.common.virtualidentity.util.{CONTROLLER, SELF}
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.concurrent.TimeLimitedTests
+import org.scalatest.flatspec.AnyFlatSpecLike
+import org.scalatest.time.Span
+import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
+
+import java.net.URI
+
+class LoggingSpec
+ extends TestKit(ActorSystem("LoggingSpec", AmberRuntime.akkaConfig))
+ with ImplicitSender
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll
+ with TimeLimitedTests {
+
+ private val identifier1 = ActorVirtualIdentity("Worker:WF1-E1-op-layer-1")
+ private val identifier2 = ActorVirtualIdentity("Worker:WF1-E1-op-layer-2")
+ private val operatorIdentity = OperatorIdentity("testOperator")
+ private val physicalOpId1 = PhysicalOpIdentity(operatorIdentity, "1st-layer")
+ private val physicalOpId2 = PhysicalOpIdentity(operatorIdentity, "2nd-layer")
+ private val mockLink = PhysicalLink(physicalOpId1, PortIdentity(), physicalOpId2, PortIdentity())
+
+ private val mockPolicy =
+ OneToOnePartitioning(10, Seq(ChannelIdentity(identifier1, identifier2, isControl = false)))
+ val payloadToLog: Array[WorkflowFIFOMessagePayload] = Array(
+ ControlInvocation(
+ METHOD_START_WORKER,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 0
+ ),
+ ControlInvocation(
+ METHOD_ADD_PARTITIONING,
+ AddPartitioningRequest(mockLink, mockPolicy),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 0
+ ),
+ ControlInvocation(
+ METHOD_PAUSE_WORKER,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 0
+ ),
+ ControlInvocation(
+ METHOD_RESUME_WORKER,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 0
+ ),
+ DataFrame(
+ (0 to 400)
+ .map(i =>
+ TupleLike(i, i.toString, i.toDouble).enforceSchema(
+ Schema()
+ .add("field1", AttributeType.INTEGER)
+ .add("field2", AttributeType.STRING)
+ .add("field3", AttributeType.DOUBLE)
+ )
+ )
+ .toArray
+ ),
+ ControlInvocation(
+ METHOD_START_WORKER,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, identifier1),
+ 0
+ ),
+ ControlInvocation(
+ METHOD_WORKER_EXECUTION_COMPLETED,
+ EmptyRequest(),
+ AsyncRPCContext(identifier1, CONTROLLER),
+ 0
+ )
+ )
+
+ "determinant logger" should "log processing steps in local storage" in {
+ Thread.sleep(1000) // wait for serializer to be registered
+ val logStorage = SequentialRecordStorage.getStorage[ReplayLogRecord](
+ Some(new URI("ram:///recovery-logs/tmp"))
+ )
+ logStorage.deleteStorage()
+ val logManager = ReplayLogManager.createLogManager(logStorage, "tmpLog", x => {})
+ payloadToLog.foreach { payload =>
+ val channel = ChannelIdentity(CONTROLLER, SELF, isControl = true)
+ val msgOpt = Some(WorkflowFIFOMessage(channel, 0, payload))
+ logManager.withFaultTolerant(channel, msgOpt) {
+ // do nothing
+ }
+ }
+ logManager.sendCommitted(null)
+ logManager.terminate()
+ val logRecords = logStorage.getReader("tmpLog").mkRecordIterator().toArray
+ logStorage.deleteStorage()
+ assert(logRecords.length == 15)
+ }
+
+ override def timeLimit: Span = 30.seconds
+}
diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/ReplaySpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/ReplaySpec.scala
new file mode 100644
index 00000000000..57c97d4d444
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/ReplaySpec.scala
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.faulttolerance
+
+import org.apache.pekko.actor.ActorSystem
+import org.apache.pekko.testkit.{ImplicitSender, TestKit}
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.logreplay.{
+ ProcessingStep,
+ ReplayLogManagerImpl,
+ ReplayLogRecord,
+ ReplayOrderEnforcer
+}
+import org.apache.texera.amber.engine.architecture.messaginglayer.NetworkInputGateway
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_START_WORKER
+import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
+import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage
+import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage.SequentialRecordReader
+import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpecLike
+
+import scala.collection.mutable
+
+class ReplaySpec
+ extends TestKit(ActorSystem("ReplaySpec"))
+ with ImplicitSender
+ with AnyFlatSpecLike
+ with BeforeAndAfterAll {
+
+ class IterableReadOnlyLogStore(iter: Iterable[ReplayLogRecord])
+ extends SequentialRecordStorage[ReplayLogRecord] {
+ override def getWriter(
+ fileName: String
+ ): SequentialRecordStorage.SequentialRecordWriter[ReplayLogRecord] = ???
+
+ override def getReader(
+ fileName: String
+ ): SequentialRecordStorage.SequentialRecordReader[ReplayLogRecord] =
+ new SequentialRecordReader[ReplayLogRecord](null) {
+ override def mkRecordIterator(): Iterator[ReplayLogRecord] = iter.iterator
+ }
+
+ override def deleteStorage(): Unit = ???
+
+ override def containsFolder(folderName: String): Boolean = ???
+ }
+
+ private val actorId = ActorVirtualIdentity("test")
+ private val actorId2 = ActorVirtualIdentity("upstream1")
+ private val actorId3 = ActorVirtualIdentity("upstream2")
+ private val channelId1 = ChannelIdentity(CONTROLLER, actorId, isControl = true)
+ private val channelId2 = ChannelIdentity(actorId2, actorId, isControl = false)
+ private val channelId3 = ChannelIdentity(actorId3, actorId, isControl = false)
+ private val channelId4 = ChannelIdentity(actorId2, actorId, isControl = true)
+ private val logManager = new ReplayLogManagerImpl(x => {})
+
+ "replay input gate" should "replay the message payload in log order" in {
+ val logRecords = mutable.Queue[ProcessingStep](
+ ProcessingStep(channelId1, -1),
+ ProcessingStep(channelId4, 1),
+ ProcessingStep(channelId3, 2),
+ ProcessingStep(channelId1, 3),
+ ProcessingStep(channelId2, 4)
+ )
+ val inputGateway = new NetworkInputGateway(actorId)
+
+ def inputMessage(channelId: ChannelIdentity, seq: Long): Unit = {
+ inputGateway
+ .getChannel(channelId)
+ .acceptMessage(
+ WorkflowFIFOMessage(
+ channelId,
+ seq,
+ ControlInvocation(
+ METHOD_START_WORKER,
+ EmptyRequest(),
+ AsyncRPCContext(CONTROLLER, actorId),
+ 0
+ )
+ )
+ )
+ }
+
+ val orderEnforcer = new ReplayOrderEnforcer(logManager, logRecords, -1, () => {})
+ inputGateway.addEnforcer(orderEnforcer)
+
+ def processMessage(channelId: ChannelIdentity, seq: Long): Unit = {
+ val msg = inputGateway.tryPickChannel.get.take
+ logManager.withFaultTolerant(msg.channelId, Some(msg)) {
+ assert(msg.channelId == channelId && msg.sequenceNumber == seq)
+ }
+ }
+
+ assert(inputGateway.tryPickChannel.isEmpty)
+ inputMessage(channelId2, 0)
+ assert(inputGateway.tryPickChannel.isEmpty)
+ inputMessage(channelId4, 0)
+ assert(inputGateway.tryPickChannel.isEmpty)
+ inputMessage(channelId1, 0)
+ inputMessage(channelId1, 1)
+ inputMessage(channelId1, 2)
+ assert(
+ inputGateway.tryPickChannel.nonEmpty && inputGateway.tryPickChannel.get.channelId == channelId1
+ )
+ processMessage(channelId1, 0)
+ assert(inputGateway.tryPickChannel.nonEmpty)
+ processMessage(channelId1, 1)
+ assert(inputGateway.tryPickChannel.nonEmpty)
+ processMessage(channelId4, 0)
+ assert(inputGateway.tryPickChannel.isEmpty)
+ inputMessage(channelId3, 0)
+ processMessage(channelId3, 0)
+ assert(inputGateway.tryPickChannel.nonEmpty)
+ processMessage(channelId1, 2)
+ assert(inputGateway.tryPickChannel.nonEmpty)
+ processMessage(channelId2, 0)
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/web/resource/dashboard/file/WorkflowResourceSpec.scala b/amber/src/test/scala/org/apache/texera/web/resource/dashboard/file/WorkflowResourceSpec.scala
new file mode 100644
index 00000000000..74a68ee65ea
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/web/resource/dashboard/file/WorkflowResourceSpec.scala
@@ -0,0 +1,783 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.file
+
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.MockTexeraDB
+import org.apache.texera.dao.jooq.generated.Tables.{USER, WORKFLOW, WORKFLOW_OF_PROJECT}
+import org.apache.texera.dao.jooq.generated.enums.UserRoleEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.UserDao
+import org.apache.texera.dao.jooq.generated.tables.pojos.{Project, User, Workflow}
+import org.apache.texera.web.resource.dashboard.DashboardResource.SearchQueryParams
+import org.apache.texera.web.resource.dashboard.user.project.ProjectResource
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource
+import org.apache.texera.web.resource.dashboard.user.workflow.WorkflowResource.{
+ DashboardWorkflow,
+ WorkflowIDs
+}
+import org.apache.texera.web.resource.dashboard.{DashboardResource, FulltextSearchQueryUtils}
+import org.jooq.Condition
+import org.jooq.impl.DSL.noCondition
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
+
+import java.sql.Timestamp
+import java.text.{ParseException, SimpleDateFormat}
+import java.time.{Duration, OffsetDateTime, ZoneOffset}
+import java.util
+import java.util.Collections
+import java.util.concurrent.TimeUnit
+
+class WorkflowResourceSpec
+ extends AnyFlatSpec
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with MockTexeraDB {
+
+ // An example creation time to test Account Creation Time attribute
+ private val exampleCreationTime: OffsetDateTime =
+ OffsetDateTime.parse("2025-01-01T00:00:00Z")
+
+ private val testUser: User = {
+ val user = new User
+ user.setUid(Integer.valueOf(1))
+ user.setName("test_user")
+ user.setRole(UserRoleEnum.ADMIN)
+ user.setPassword("123")
+ user.setComment("test_comment")
+ user.setAccountCreationTime(exampleCreationTime)
+ user
+ }
+
+ private val testUser2: User = {
+ val user = new User
+ user.setUid(Integer.valueOf(2))
+ user.setName("test_user2")
+ user.setRole(UserRoleEnum.ADMIN)
+ user.setPassword("123")
+ user.setComment("test_comment2")
+ user.setAccountCreationTime(exampleCreationTime)
+ user
+ }
+
+ private val keywordInWorkflow1Content = "keyword_in_workflow1_content"
+ private val textPhrase = "text phrases"
+ private val exampleContent =
+ "{\"x\":5,\"y\":\"" + keywordInWorkflow1Content + "\",\"z\":\"" + textPhrase + "\"}"
+
+ private val testWorkflow1: Workflow = {
+ val workflow = new Workflow()
+ workflow.setName("test_workflow1")
+ workflow.setDescription("keyword_in_workflow_description")
+ workflow.setContent(exampleContent)
+
+ workflow
+ }
+
+ private val testWorkflow2: Workflow = {
+ val workflow = new Workflow()
+ workflow.setName("test_workflow2")
+ workflow.setDescription("another_text")
+ workflow.setContent("{\"x\":5,\"y\":\"example2\",\"z\":\"\"}")
+
+ workflow
+ }
+
+ private val testWorkflow3: Workflow = {
+ val workflow = new Workflow()
+ workflow.setName("test_workflow3")
+ workflow.setDescription("")
+ workflow.setContent("{\"x\":5,\"y\":\"example3\",\"z\":\"\"}")
+
+ workflow
+ }
+
+ private val testProject1: Project = {
+ val project = new Project()
+ project.setName("test_project1")
+ project.setDescription("this is project description")
+ project
+ }
+
+ private val exampleEmailAddress = "name@example.com"
+ private val exampleWord1 = "Lorem"
+ private val exampleWord2 = "Ipsum"
+
+ private val testWorkflowWithSpecialCharacters: Workflow = {
+ val workflow = new Workflow()
+ workflow.setName("workflow_with_special_characters")
+ workflow.setDescription(exampleWord1 + " " + exampleWord2 + " " + exampleEmailAddress)
+ workflow.setContent(exampleContent)
+
+ workflow
+ }
+
+ private val sessionUser1: SessionUser = {
+ new SessionUser(testUser)
+ }
+
+ private val sessionUser2: SessionUser = {
+ new SessionUser(testUser2)
+ }
+
+ private val workflowResource: WorkflowResource = {
+ new WorkflowResource()
+ }
+
+ private val projectResource: ProjectResource = {
+ new ProjectResource()
+ }
+
+ private val dashboardResource: DashboardResource = {
+ new DashboardResource()
+ }
+
+ override protected def beforeAll(): Unit = {
+ initializeDBAndReplaceDSLContext()
+ FulltextSearchQueryUtils.usePgroonga = false // disable pgroonga
+ // add test user directly
+ val userDao = new UserDao(getDSLContext.configuration())
+ userDao.insert(testUser)
+ userDao.insert(testUser2)
+ }
+
+ override protected def beforeEach(): Unit = {
+ // Clean up environment before each test case
+ // Delete all workflows, or reset the state of the `workflowResource` object
+ }
+
+ override protected def afterEach(): Unit = {
+ // Clean up environment after each test case if necessary
+ // delete all workflows in the database
+ var workflows = workflowResource.retrieveWorkflowsBySessionUser(sessionUser1)
+ workflows.foreach(workflow =>
+ workflowResource.deleteWorkflow(
+ WorkflowIDs(List(workflow.workflow.getWid), None),
+ sessionUser1
+ )
+ )
+
+ workflows = workflowResource.retrieveWorkflowsBySessionUser(sessionUser2)
+ workflows.foreach(workflow =>
+ workflowResource.deleteWorkflow(
+ WorkflowIDs(List(workflow.workflow.getWid), None),
+ sessionUser2
+ )
+ )
+
+ // delete all projects in the database
+ var projects = projectResource.getProjectList(sessionUser1)
+ projects.forEach(project => projectResource.deleteProject(project.pid))
+
+ projects = projectResource.getProjectList(sessionUser2)
+ projects.forEach(project => projectResource.deleteProject(project.pid))
+
+ }
+
+ override protected def afterAll(): Unit = {
+ shutdownDB()
+ }
+
+ private def getKeywordsArray(keywords: String*): util.ArrayList[String] = {
+ val keywordsList = new util.ArrayList[String]()
+ for (keyword <- keywords) {
+ keywordsList.add(keyword)
+ }
+ keywordsList
+ }
+
+ private def insertAndAssertAccountCreation(uid: Int, ts: OffsetDateTime): Unit = {
+ val userDao = new UserDao(getDSLContext.configuration())
+ val u = new User
+ u.setUid(Integer.valueOf(uid))
+ u.setName(s"tmp_user_$uid")
+ u.setRole(UserRoleEnum.REGULAR)
+ u.setPassword("pw")
+ u.setComment("tmp")
+ u.setAccountCreationTime(ts)
+ userDao.insert(u)
+
+ try {
+ val fetched = userDao.fetchOneByUid(Integer.valueOf(uid))
+ assert(fetched.getAccountCreationTime != null)
+ assert(fetched.getAccountCreationTime.isEqual(ts))
+ } finally {
+ userDao.deleteById(Integer.valueOf(uid))
+ }
+ }
+
+ private def assertSameWorkflow(a: Workflow, b: DashboardWorkflow): Unit = {
+ assert(a.getName == b.workflow.getName)
+ }
+
+ "User.accountCreationTime" should "be persisted and retrievable via UserDao" in {
+ val userDao = new UserDao(getDSLContext.configuration())
+ val u1 = userDao.fetchOneByUid(Integer.valueOf(1))
+ val u2 = userDao.fetchOneByUid(Integer.valueOf(2))
+
+ assert(u1.getAccountCreationTime != null)
+ assert(u2.getAccountCreationTime != null)
+
+ assert(u1.getAccountCreationTime.isEqual(exampleCreationTime))
+ assert(u2.getAccountCreationTime.isEqual(exampleCreationTime))
+ }
+
+ it should "remain unchanged when updating unrelated fields" in {
+ val userDao = new UserDao(getDSLContext.configuration())
+ val u1 = userDao.fetchOneByUid(Integer.valueOf(1))
+ val originalTime = u1.getAccountCreationTime
+
+ u1.setComment("updated_comment")
+ userDao.update(u1)
+
+ val test_u1 = userDao.fetchOneByUid(Integer.valueOf(1))
+ assert(test_u1.getAccountCreationTime.isEqual(originalTime))
+ }
+
+ it should "fallback to DB default when not explicitly set on insert" in {
+ // account_creation_time TIMESTAMPTZ NOT NULL DEFAULT now()
+ val userDao = new UserDao(getDSLContext.configuration())
+ // Test user 3 on top of test user 1 and 2
+ val userId = 3
+ val tmp = new User
+ tmp.setUid(Integer.valueOf(userId))
+ tmp.setName("tmp_user")
+ tmp.setRole(UserRoleEnum.REGULAR)
+ tmp.setPassword("pw")
+ tmp.setComment("tmp")
+ // Account creation time not set
+ userDao.insert(tmp)
+
+ val fetched = userDao.fetchOneByUid(Integer.valueOf(3))
+ assert(fetched.getAccountCreationTime != null)
+
+ val now = OffsetDateTime.now(ZoneOffset.UTC)
+ val diff = Duration.between(fetched.getAccountCreationTime, now).abs()
+ assert(diff.toMinutes <= 2)
+ }
+
+ // Testing with user id 4
+ it should "persist and retrieve a non-UTC offset time (ex: +09:00 JST)" in {
+ val userId = 4
+ insertAndAssertAccountCreation(
+ uid = userId,
+ ts = OffsetDateTime.parse("2020-06-15T12:34:56+09:00")
+ )
+ }
+
+ // Testing with user id 5
+ it should "persist and retrieve a leap day timestamp" in {
+ val userId = 5
+ insertAndAssertAccountCreation(
+ uid = userId,
+ ts = OffsetDateTime.parse("2024-02-29T23:59:59Z")
+ )
+ }
+
+ // Testing with user id 6
+ it should "persist and retrieve a future timestamp" in {
+ val userId = 6
+ insertAndAssertAccountCreation(
+ uid = userId,
+ ts = OffsetDateTime.parse("2100-12-31T23:59:59Z")
+ )
+ }
+
+ "WorkflowResource /owner_name" should "return owner name as plain text" in {
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+
+ val workflows = workflowResource.retrieveWorkflowsBySessionUser(sessionUser1)
+ assert(workflows.nonEmpty)
+
+ val wid =
+ workflows
+ .find(_.workflow.getName == testWorkflow1.getName)
+ .map(_.workflow.getWid)
+ .getOrElse(workflows.head.workflow.getWid)
+
+ val ownerName = workflowResource.getOwnerName(wid)
+
+ assert(ownerName == testUser.getName)
+ }
+
+ "/search API " should "be able to search for workflows in different columns in Workflow table" in {
+ // testWorkflow1: {name: test_name, descrption: test_description, content: test_content}
+ // search "test_name" or "test_description" or "test_content" should return testWorkflow1
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow3, sessionUser1)
+ // search
+ val DashboardWorkflowEntryList =
+ dashboardResource
+ .searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(keywords = getKeywordsArray(keywordInWorkflow1Content))
+ )
+ .results
+ assert(DashboardWorkflowEntryList.head.workflow.get.ownerName.equals(testUser.getName))
+ assert(DashboardWorkflowEntryList.length == 1)
+ assertSameWorkflow(testWorkflow1, DashboardWorkflowEntryList.head.workflow.get)
+ }
+
+ it should "be able to search text phrases" in {
+ // testWorkflow1: {name: "test_name", descrption: "test_description", content: "text phrase"}
+ // search "text phrase" should return testWorkflow1
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow3, sessionUser1)
+ val DashboardWorkflowEntryList =
+ dashboardResource
+ .searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(keywords = getKeywordsArray(keywordInWorkflow1Content))
+ )
+ .results
+ assert(DashboardWorkflowEntryList.length == 1)
+ assertSameWorkflow(testWorkflow1, DashboardWorkflowEntryList.head.workflow.get)
+ val DashboardWorkflowEntryList1 =
+ dashboardResource
+ .searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(keywords = getKeywordsArray("text sear"))
+ )
+ .results
+ assert(DashboardWorkflowEntryList1.isEmpty)
+ }
+
+ it should "return an all workflows when given an empty list of keywords" in {
+ // search "" should return all workflows
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow3, sessionUser1)
+ val DashboardWorkflowEntryList =
+ dashboardResource.searchAllResourcesCall(sessionUser1, SearchQueryParams())
+ assert(DashboardWorkflowEntryList.results.length == 2)
+ }
+
+ it should "be able to search with arbitrary number of keywords in different combinations" in {
+ // testWorkflow1: {name: test_name, description: test_description, content: "key pair"}
+ // search ["key"] or ["pair", "key"] should return the testWorkflow1
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow3, sessionUser1)
+ // search with multiple keywords
+ val keywords = new util.ArrayList[String]()
+ keywords.add(keywordInWorkflow1Content)
+ keywords.add(testWorkflow1.getDescription)
+ val DashboardWorkflowEntryList = dashboardResource
+ .searchAllResourcesCall(sessionUser1, SearchQueryParams(keywords = keywords))
+ .results
+ assert(DashboardWorkflowEntryList.size == 1)
+ assert(DashboardWorkflowEntryList.head.workflow.get.ownerName.equals(testUser.getName))
+ assertSameWorkflow(testWorkflow1, DashboardWorkflowEntryList.head.workflow.get)
+
+ keywords.add("nonexistent")
+ val DashboardWorkflowEntryList2 = dashboardResource
+ .searchAllResourcesCall(sessionUser1, SearchQueryParams(keywords = keywords))
+ .results
+ assert(DashboardWorkflowEntryList2.isEmpty)
+
+ val keywordsReverseOrder = new util.ArrayList[String]()
+ keywordsReverseOrder.add(testWorkflow1.getDescription)
+ keywordsReverseOrder.add(keywordInWorkflow1Content)
+ val DashboardWorkflowEntryList1 =
+ dashboardResource
+ .searchAllResourcesCall(sessionUser1, SearchQueryParams(keywords = keywordsReverseOrder))
+ .results
+ assert(DashboardWorkflowEntryList1.size == 1)
+ assert(DashboardWorkflowEntryList1.head.workflow.get.ownerName.equals(testUser.getName))
+ assertSameWorkflow(testWorkflow1, DashboardWorkflowEntryList1.head.workflow.get)
+
+ }
+
+ it should "handle reserved characters in the keywords" in {
+ // testWorkflow1: {name: test_name, description: test_description, content: "key pair"}
+ // search "key+-pair" or "key@pair" or "key+" or "+key" should return testWorkflow1
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow3, sessionUser1)
+
+ def testInner(keywords: String): Unit = {
+ val DashboardWorkflowEntryList = dashboardResource
+ .searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(keywords = getKeywordsArray(keywords))
+ )
+ .results
+ assert(DashboardWorkflowEntryList.size == 1)
+ assert(DashboardWorkflowEntryList.head.workflow.get.ownerName.equals(testUser.getName))
+ assertSameWorkflow(testWorkflow1, DashboardWorkflowEntryList.head.workflow.get)
+ }
+
+ testInner(keywordInWorkflow1Content + "+-@()<>~*\"" + keywordInWorkflow1Content)
+ testInner(keywordInWorkflow1Content + "@" + keywordInWorkflow1Content)
+ testInner(keywordInWorkflow1Content + "+-@()<>~*\"")
+ testInner("+-@()<>~*\"" + keywordInWorkflow1Content)
+
+ }
+
+ it should "return all workflows when keywords only contains reserved keywords +-@()<>~*\"" in {
+ // search "+-@()<>~*"" should return all workflows
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow3, sessionUser1)
+
+ val DashboardWorkflowEntryList =
+ dashboardResource
+ .searchAllResourcesCall(sessionUser1, SearchQueryParams(getKeywordsArray("+-@()<>~*\"")))
+ .results
+ assert(DashboardWorkflowEntryList.size == 2)
+
+ }
+
+ it should "not be able to search workflows from different user accounts" in {
+ // user1 has workflow1
+ // user2 has workflow2
+ // users should only be able to search for workflows they have access to
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow2, sessionUser2)
+ workflowResource.persistWorkflow(testWorkflow3, sessionUser1)
+
+ def test(user: SessionUser, workflow: Workflow): Unit = {
+ // search with reserved characters in keywords
+ val DashboardWorkflowEntryList =
+ dashboardResource
+ .searchAllResourcesCall(
+ user,
+ SearchQueryParams(getKeywordsArray(workflow.getDescription))
+ )
+ .results
+ assert(DashboardWorkflowEntryList.size == 1)
+ assert(DashboardWorkflowEntryList.head.workflow.get.ownerName.equals(user.getName()))
+ assertSameWorkflow(workflow, DashboardWorkflowEntryList.head.workflow.get)
+ }
+
+ test(sessionUser1, testWorkflow1)
+ test(sessionUser2, testWorkflow2)
+ }
+
+ it should "return a proper condition for a single owner" in {
+ val ownerList = new java.util.ArrayList[String](util.Arrays.asList("owner1"))
+ val ownerFilter: Condition =
+ FulltextSearchQueryUtils.getContainsFilter(ownerList, USER.EMAIL)
+ assert(ownerFilter.toString == USER.EMAIL.eq("owner1").toString)
+ }
+
+ it should "return a proper condition for multiple owners" in {
+ val ownerList = new java.util.ArrayList[String](util.Arrays.asList("owner1", "owner2"))
+ val ownerFilter: Condition =
+ FulltextSearchQueryUtils.getContainsFilter(ownerList, USER.EMAIL)
+ assert(ownerFilter.toString == USER.EMAIL.eq("owner1").or(USER.EMAIL.eq("owner2")).toString)
+ }
+
+ it should "return a proper condition for a single projectId" in {
+ val projectIdList = new java.util.ArrayList[Integer](util.Arrays.asList(Integer.valueOf(1)))
+ val projectFilter: Condition =
+ FulltextSearchQueryUtils.getContainsFilter(projectIdList, WORKFLOW_OF_PROJECT.PID)
+ assert(projectFilter.toString == WORKFLOW_OF_PROJECT.PID.eq(Integer.valueOf(1)).toString)
+ }
+
+ it should "return a proper condition for multiple projectIds" in {
+ val projectIdList = new java.util.ArrayList[Integer](
+ util.Arrays.asList(Integer.valueOf(1), Integer.valueOf(2))
+ )
+ val projectFilter: Condition =
+ FulltextSearchQueryUtils.getContainsFilter(projectIdList, WORKFLOW_OF_PROJECT.PID)
+ assert(
+ projectFilter.toString == WORKFLOW_OF_PROJECT.PID
+ .eq(Integer.valueOf(1))
+ .or(WORKFLOW_OF_PROJECT.PID.eq(Integer.valueOf(2)))
+ .toString
+ )
+ }
+
+ it should "return a proper condition for a single workflowID" in {
+ val workflowIdList = new java.util.ArrayList[Integer](util.Arrays.asList(Integer.valueOf(1)))
+ val workflowIdFilter: Condition =
+ FulltextSearchQueryUtils.getContainsFilter(workflowIdList, WORKFLOW.WID)
+ assert(workflowIdFilter.toString == WORKFLOW.WID.eq(Integer.valueOf(1)).toString)
+ }
+
+ it should "return a proper condition for multiple workflowIDs" in {
+ val workflowIdList = new java.util.ArrayList[Integer](
+ util.Arrays.asList(Integer.valueOf(1), Integer.valueOf(2))
+ )
+ val workflowIdFilter: Condition =
+ FulltextSearchQueryUtils.getContainsFilter(workflowIdList, WORKFLOW.WID)
+ assert(
+ workflowIdFilter.toString == WORKFLOW.WID
+ .eq(Integer.valueOf(1))
+ .or(WORKFLOW.WID.eq(Integer.valueOf(2)))
+ .toString
+ )
+ }
+
+ it should "return a proper condition for creation date type with specific start and end date" in {
+ val dateFilter: Condition =
+ FulltextSearchQueryUtils.getDateFilter(
+ "2023-01-01",
+ "2023-12-31",
+ WORKFLOW.CREATION_TIME
+ )
+ val dateFormat = new SimpleDateFormat("yyyy-MM-dd")
+ val startTimestamp = new Timestamp(dateFormat.parse("2023-01-01").getTime)
+ val endTimestamp =
+ new Timestamp(
+ dateFormat.parse("2023-12-31").getTime + TimeUnit.DAYS.toMillis(1) - 1
+ )
+ assert(
+ dateFilter.toString == WORKFLOW.CREATION_TIME.between(startTimestamp, endTimestamp).toString
+ )
+ }
+
+ it should "return a proper condition for modification date type with specific start and end date" in {
+ val dateFilter: Condition =
+ FulltextSearchQueryUtils.getDateFilter(
+ "2023-01-01",
+ "2023-12-31",
+ WORKFLOW.LAST_MODIFIED_TIME
+ )
+ val dateFormat = new SimpleDateFormat("yyyy-MM-dd")
+ val startTimestamp = new Timestamp(dateFormat.parse("2023-01-01").getTime)
+ val endTimestamp =
+ new Timestamp(
+ dateFormat.parse("2023-12-31").getTime + TimeUnit.DAYS.toMillis(1) - 1
+ )
+ assert(
+ dateFilter.toString == WORKFLOW.LAST_MODIFIED_TIME
+ .between(startTimestamp, endTimestamp)
+ .toString
+ )
+ }
+
+ it should "throw a ParseException when endDate is invalid" in {
+ assertThrows[ParseException] {
+ FulltextSearchQueryUtils.getDateFilter(
+ "2023-01-01",
+ "invalidDate",
+ WORKFLOW.CREATION_TIME
+ )
+ }
+ }
+
+ "getOperatorsFilter" should "return a noCondition when the input operators list is empty" in {
+ val operatorsFilter: Condition =
+ FulltextSearchQueryUtils.getOperatorsFilter(
+ Collections.emptyList[String](),
+ WORKFLOW.CONTENT
+ )
+ assert(operatorsFilter.toString == noCondition().toString)
+ }
+
+ it should "return a proper condition for a single operator" in {
+ val operatorsList = new java.util.ArrayList[String](util.Arrays.asList("operator1"))
+ val operatorsFilter: Condition =
+ FulltextSearchQueryUtils.getOperatorsFilter(operatorsList, WORKFLOW.CONTENT)
+ val searchKey = "%\"operatorType\":\"operator1\"%"
+ assert(operatorsFilter.toString == WORKFLOW.CONTENT.likeIgnoreCase(searchKey).toString)
+ }
+
+ it should "return a proper condition for multiple operators" in {
+ val operatorsList =
+ new java.util.ArrayList[String](util.Arrays.asList("operator1", "operator2"))
+ val operatorsFilter: Condition =
+ FulltextSearchQueryUtils.getOperatorsFilter(operatorsList, WORKFLOW.CONTENT)
+ val searchKey1 = "%\"operatorType\":\"operator1\"%"
+ val searchKey2 = "%\"operatorType\":\"operator2\"%"
+ assert(
+ operatorsFilter.toString == WORKFLOW.CONTENT
+ .likeIgnoreCase(searchKey1)
+ .or(WORKFLOW.CONTENT.likeIgnoreCase(searchKey2))
+ .toString
+ )
+ }
+
+ "/search API" should "be able to search for resources in different tables" in {
+
+ // create different types of resources, project, workflow, and file
+ projectResource.createProject(sessionUser1, "test project1")
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ // search
+ val DashboardClickableFileEntryList =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(getKeywordsArray("test"))
+ )
+ assert(DashboardClickableFileEntryList.results.length == 2)
+
+ }
+
+ it should "return all resources when no keyword provided" in {
+ projectResource.createProject(sessionUser1, "test project1")
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ val DashboardClickableFileEntryList =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(getKeywordsArray(""))
+ )
+ assert(DashboardClickableFileEntryList.results.length == 2)
+ }
+
+ it should "return multiple matching resources from a single resource type" in {
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ projectResource.createProject(sessionUser1, "common project1")
+ projectResource.createProject(sessionUser1, "common project2")
+ val DashboardClickableFileEntryList =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(getKeywordsArray("common"))
+ )
+ assert(DashboardClickableFileEntryList.results.length == 2)
+ }
+
+ it should "handle multiple keywords correctly" in {
+ projectResource.createProject(sessionUser1, "test project1")
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ val DashboardClickableFileEntryList =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(getKeywordsArray("test", "project1"))
+ )
+ assert(
+ DashboardClickableFileEntryList.results.length == 1
+ ) // should only return the project
+ }
+
+ it should "filter results by different resourceType" in {
+ // create different types of resources
+ // 3 projects, 2 file, and 1 workflow,
+ projectResource.createProject(sessionUser1, "test project1")
+ projectResource.createProject(sessionUser1, "test project2")
+ projectResource.createProject(sessionUser1, "test project3")
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ // search resources with all resourceType
+ var DashboardClickableFileEntryList =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(getKeywordsArray("test"))
+ )
+ assert(DashboardClickableFileEntryList.results.length == 4)
+
+ // filter resources by workflow
+ DashboardClickableFileEntryList = dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(resourceType = "workflow", keywords = getKeywordsArray("test"))
+ )
+ assert(DashboardClickableFileEntryList.results.length == 1)
+
+ // filter resources by project
+ DashboardClickableFileEntryList = dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(resourceType = "project", keywords = getKeywordsArray("test"))
+ )
+ assert(DashboardClickableFileEntryList.results.length == 3)
+ }
+
+ it should "return resources that match any of all provided keywords" in {
+ // This test is designed to verify that the searchAllResources function correctly
+ // returns resources that match all of the provided keywords
+
+ // Create different types of resources, a project, a workflow, and a file
+ projectResource.createProject(sessionUser1, "test project")
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ // Perform search with multiple keywords
+ val DashboardClickableFileEntryList =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(keywords = getKeywordsArray("test", "project"))
+ )
+
+ // Assert that the search results include resources that match any of the provided keywords
+ assert(DashboardClickableFileEntryList.results.length == 1)
+ }
+
+ it should "not return resources that belong to a different user" in {
+ // This test is designed to verify that the searchAllResources function does not return resources that belong to a different user
+
+ // Create a project for a different user (sessionUser2)
+ projectResource.createProject(sessionUser2, "test project2")
+
+ // Perform search for resources using sessionUser1
+ val DashboardClickableFileEntryList =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(keywords = getKeywordsArray("test"))
+ )
+
+ // Assert that the search results do not include the project that belongs to the different user
+ // Assuming that DashboardClickableFileEntryList is a list of resources where each resource has a `user` property
+ assert(DashboardClickableFileEntryList.results.isEmpty)
+ }
+
+ it should "paginate results correctly" in {
+ // This test is designed to verify that the pagination works correctly
+
+ // Create 1 workflow, 10 projects
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ for (i <- 1 to 10) {
+ projectResource.createProject(sessionUser1, s"test project $i")
+ }
+
+ // Request the first page of results (page size is 10)
+ val firstPage =
+ dashboardResource.searchAllResourcesCall(sessionUser1, SearchQueryParams(count = 10))
+
+ // Assert that the first page has 10 results
+ assert(firstPage.results.length == 10)
+ assert(firstPage.more) // Assert that there are more results to be fetched
+
+ // Request the second page of results
+ val secondPage =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(count = 10, offset = 10)
+ )
+
+ // Assert that the second page has 1 results
+ assert(secondPage.results.length == 1)
+
+ // Assert that the results are unique across all pages
+ val allResults = firstPage.results ++ secondPage.results
+ assert(allResults.distinct.length == allResults.length)
+ }
+
+ it should "order workflow by name correctly" in {
+ // Create several resources with different names
+ workflowResource.persistWorkflow(testWorkflow1, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow3, sessionUser1)
+ workflowResource.persistWorkflow(testWorkflow2, sessionUser1)
+
+ // Retrieve resources ordered by name in ascending order
+ var resources =
+ dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(resourceType = "workflow", orderBy = "NameAsc")
+ )
+
+ // Check the order of the results
+ assert(resources.results(0).workflow.get.workflow.getName == "test_workflow1")
+ assert(resources.results(1).workflow.get.workflow.getName == "test_workflow2")
+ assert(resources.results(2).workflow.get.workflow.getName == "test_workflow3")
+
+ resources = dashboardResource.searchAllResourcesCall(
+ sessionUser1,
+ SearchQueryParams(resourceType = "workflow", orderBy = "NameDesc")
+ )
+ // Check the order of the results
+ assert(resources.results(0).workflow.get.workflow.getName == "test_workflow3")
+ assert(resources.results(1).workflow.get.workflow.getName == "test_workflow2")
+ assert(resources.results(2).workflow.get.workflow.getName == "test_workflow1")
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowAccessResourceSpec.scala b/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowAccessResourceSpec.scala
new file mode 100644
index 00000000000..163f7b2683b
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowAccessResourceSpec.scala
@@ -0,0 +1,405 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.workflow
+
+import org.apache.texera.auth.SessionUser
+import org.apache.texera.dao.MockTexeraDB
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
+import org.apache.texera.dao.jooq.generated.tables.daos.{
+ UserDao,
+ WorkflowDao,
+ WorkflowOfUserDao,
+ WorkflowUserAccessDao
+}
+import org.apache.texera.dao.jooq.generated.tables.pojos.{
+ User,
+ Workflow,
+ WorkflowOfUser,
+ WorkflowUserAccess
+}
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
+
+import java.sql.Timestamp
+import javax.ws.rs.{BadRequestException, ForbiddenException}
+
+class WorkflowAccessResourceSpec
+ extends AnyFlatSpec
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with MockTexeraDB {
+
+ private val ownerUid = 1000 + scala.util.Random.nextInt(1000)
+ private val userWithWriteUid = 2000 + scala.util.Random.nextInt(1000)
+ private val userWithReadUid = 3000 + scala.util.Random.nextInt(1000)
+ private val targetUserUid = 4000 + scala.util.Random.nextInt(1000)
+ private val testWorkflowWid = 5000 + scala.util.Random.nextInt(1000)
+
+ private var owner: User = _
+ private var userWithWrite: User = _
+ private var userWithRead: User = _
+ private var targetUser: User = _
+ private var testWorkflow: Workflow = _
+
+ private var userDao: UserDao = _
+ private var workflowDao: WorkflowDao = _
+ private var workflowOfUserDao: WorkflowOfUserDao = _
+ private var workflowUserAccessDao: WorkflowUserAccessDao = _
+ private var workflowAccessResource: WorkflowAccessResource = _
+
+ override protected def beforeAll(): Unit = {
+ initializeDBAndReplaceDSLContext()
+ }
+
+ override protected def beforeEach(): Unit = {
+ // Initialize DAOs
+ userDao = new UserDao(getDSLContext.configuration())
+ workflowDao = new WorkflowDao(getDSLContext.configuration())
+ workflowOfUserDao = new WorkflowOfUserDao(getDSLContext.configuration())
+ workflowUserAccessDao = new WorkflowUserAccessDao(getDSLContext.configuration())
+ workflowAccessResource = new WorkflowAccessResource()
+
+ // Create test users
+ owner = new User
+ owner.setUid(ownerUid)
+ owner.setName("owner")
+ owner.setEmail("owner@test.com")
+ owner.setPassword("password")
+
+ userWithWrite = new User
+ userWithWrite.setUid(userWithWriteUid)
+ userWithWrite.setName("user_with_write")
+ userWithWrite.setEmail("write@test.com")
+ userWithWrite.setPassword("password")
+
+ userWithRead = new User
+ userWithRead.setUid(userWithReadUid)
+ userWithRead.setName("user_with_read")
+ userWithRead.setEmail("read@test.com")
+ userWithRead.setPassword("password")
+
+ targetUser = new User
+ targetUser.setUid(targetUserUid)
+ targetUser.setName("target_user")
+ targetUser.setEmail("target@test.com")
+ targetUser.setPassword("password")
+
+ // Create test workflow
+ testWorkflow = new Workflow
+ testWorkflow.setWid(testWorkflowWid)
+ testWorkflow.setName("test_workflow")
+ testWorkflow.setContent("{}")
+ testWorkflow.setDescription("test description")
+ testWorkflow.setCreationTime(new Timestamp(System.currentTimeMillis()))
+ testWorkflow.setLastModifiedTime(new Timestamp(System.currentTimeMillis()))
+
+ // Clean up before each test
+ cleanupTestData()
+
+ // Insert test data
+ userDao.insert(owner)
+ userDao.insert(userWithWrite)
+ userDao.insert(userWithRead)
+ userDao.insert(targetUser)
+ workflowDao.insert(testWorkflow)
+
+ // Set up workflow ownership
+ val workflowOfUser = new WorkflowOfUser
+ workflowOfUser.setUid(ownerUid)
+ workflowOfUser.setWid(testWorkflowWid)
+ workflowOfUserDao.insert(workflowOfUser)
+
+ // Grant write access to userWithWrite
+ val writeAccess = new WorkflowUserAccess
+ writeAccess.setUid(userWithWriteUid)
+ writeAccess.setWid(testWorkflowWid)
+ writeAccess.setPrivilege(PrivilegeEnum.WRITE)
+ workflowUserAccessDao.insert(writeAccess)
+
+ // Grant read access to userWithRead
+ val readAccess = new WorkflowUserAccess
+ readAccess.setUid(userWithReadUid)
+ readAccess.setWid(testWorkflowWid)
+ readAccess.setPrivilege(PrivilegeEnum.READ)
+ workflowUserAccessDao.insert(readAccess)
+
+ // Grant write access to targetUser
+ val targetAccess = new WorkflowUserAccess
+ targetAccess.setUid(targetUserUid)
+ targetAccess.setWid(testWorkflowWid)
+ targetAccess.setPrivilege(PrivilegeEnum.WRITE)
+ workflowUserAccessDao.insert(targetAccess)
+ }
+
+ override protected def afterEach(): Unit = {
+ cleanupTestData()
+ }
+
+ private def cleanupTestData(): Unit = {
+ getDSLContext
+ .deleteFrom(WORKFLOW_USER_ACCESS)
+ .where(WORKFLOW_USER_ACCESS.WID.eq(testWorkflowWid))
+ .execute()
+
+ getDSLContext
+ .deleteFrom(WORKFLOW_OF_USER)
+ .where(WORKFLOW_OF_USER.WID.eq(testWorkflowWid))
+ .execute()
+
+ getDSLContext
+ .deleteFrom(WORKFLOW)
+ .where(WORKFLOW.WID.eq(testWorkflowWid))
+ .execute()
+
+ getDSLContext
+ .deleteFrom(USER)
+ .where(
+ USER.UID.in(ownerUid, userWithWriteUid, userWithReadUid, targetUserUid)
+ )
+ .execute()
+ }
+
+ override protected def afterAll(): Unit = {
+ shutdownDB()
+ }
+
+ "WorkflowAccessResource.revokeAccess" should "successfully revoke access when user has WRITE permission" in {
+ val sessionUser = new SessionUser(userWithWrite)
+
+ // Verify target user has access before revocation
+ val accessBefore = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(targetUserUid)
+ )
+ )
+ .fetchOne()
+ assert(accessBefore != null, "Target user should have access before revocation")
+
+ // Revoke access
+ workflowAccessResource.revokeAccess(testWorkflowWid, "target@test.com", sessionUser)
+
+ // Verify access has been revoked
+ val accessAfter = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(targetUserUid)
+ )
+ )
+ .fetchOne()
+
+ assert(accessAfter == null, "Target user's access should be revoked")
+ }
+
+ it should "successfully allow user to revoke their own access" in {
+ val sessionUser = new SessionUser(userWithRead)
+
+ // Verify user has access before revocation
+ val accessBefore = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(userWithReadUid)
+ )
+ )
+ .fetchOne()
+ assert(accessBefore != null, "User should have access before revocation")
+
+ // User revokes their own access
+ workflowAccessResource.revokeAccess(testWorkflowWid, "read@test.com", sessionUser)
+
+ // Verify access has been revoked
+ val accessAfter = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(userWithReadUid)
+ )
+ )
+ .fetchOne()
+
+ assert(accessAfter == null, "User's own access should be revoked")
+ }
+
+ it should "throw ForbiddenException when user without WRITE permission tries to revoke others' access" in {
+ val sessionUser = new SessionUser(userWithRead)
+
+ assertThrows[ForbiddenException] {
+ workflowAccessResource.revokeAccess(testWorkflowWid, "target@test.com", sessionUser)
+ }
+
+ // Verify target user's access is still intact
+ val access = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(targetUserUid)
+ )
+ )
+ .fetchOne()
+
+ assert(access != null, "Target user's access should remain intact")
+ }
+
+ it should "throw ForbiddenException when trying to revoke owner's access" in {
+ val sessionUser = new SessionUser(userWithWrite)
+
+ val exception = intercept[ForbiddenException] {
+ workflowAccessResource.revokeAccess(testWorkflowWid, "owner@test.com", sessionUser)
+ }
+
+ assert(
+ exception.getMessage.contains("owner cannot revoke their own access"),
+ "Exception message should indicate owner cannot revoke their own access"
+ )
+ }
+
+ it should "throw ForbiddenException when owner tries to revoke their own access" in {
+ val sessionUser = new SessionUser(owner)
+
+ val exception = intercept[ForbiddenException] {
+ workflowAccessResource.revokeAccess(testWorkflowWid, "owner@test.com", sessionUser)
+ }
+
+ assert(
+ exception.getMessage.contains("owner cannot revoke their own access"),
+ "Exception message should indicate owner cannot revoke their own access"
+ )
+ }
+
+ it should "throw BadRequestException when email does not exist" in {
+ val sessionUser = new SessionUser(userWithWrite)
+
+ assertThrows[BadRequestException] {
+ workflowAccessResource.revokeAccess(
+ testWorkflowWid,
+ "nonexistent@test.com",
+ sessionUser
+ )
+ }
+ }
+
+ it should "not affect other users' access when revoking one user's access" in {
+ val sessionUser = new SessionUser(userWithWrite)
+
+ // Verify both users have access before revocation
+ val readAccessBefore = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(userWithReadUid)
+ )
+ )
+ .fetchOne()
+ assert(readAccessBefore != null, "Read user should have access before revocation")
+
+ val targetAccessBefore = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(targetUserUid)
+ )
+ )
+ .fetchOne()
+ assert(targetAccessBefore != null, "Target user should have access before revocation")
+
+ // Revoke only target user's access
+ workflowAccessResource.revokeAccess(testWorkflowWid, "target@test.com", sessionUser)
+
+ // Verify read user's access is still intact
+ val readAccessAfter = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(userWithReadUid)
+ )
+ )
+ .fetchOne()
+ assert(readAccessAfter != null, "Read user's access should remain intact")
+
+ // Verify target user's access has been revoked
+ val targetAccessAfter = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(targetUserUid)
+ )
+ )
+ .fetchOne()
+ assert(targetAccessAfter == null, "Target user's access should be revoked")
+ }
+
+ it should "handle revoking access for a user who already has no access gracefully" in {
+ val sessionUser = new SessionUser(userWithWrite)
+
+ // First revocation
+ workflowAccessResource.revokeAccess(testWorkflowWid, "target@test.com", sessionUser)
+
+ // Verify access has been revoked
+ val accessAfterFirst = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(targetUserUid)
+ )
+ )
+ .fetchOne()
+ assert(accessAfterFirst == null, "Target user's access should be revoked")
+
+ // Second revocation attempt (should not throw an error, just do nothing)
+ workflowAccessResource.revokeAccess(testWorkflowWid, "target@test.com", sessionUser)
+
+ // Verify access is still revoked
+ val accessAfterSecond = getDSLContext
+ .selectFrom(WORKFLOW_USER_ACCESS)
+ .where(
+ WORKFLOW_USER_ACCESS.WID
+ .eq(testWorkflowWid)
+ .and(
+ WORKFLOW_USER_ACCESS.UID.eq(targetUserUid)
+ )
+ )
+ .fetchOne()
+ assert(accessAfterSecond == null, "Target user's access should still be revoked")
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowExecutionsResourceSpec.scala b/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowExecutionsResourceSpec.scala
new file mode 100644
index 00000000000..bd55124a729
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowExecutionsResourceSpec.scala
@@ -0,0 +1,219 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.workflow
+
+import org.apache.texera.amber.core.virtualidentity.{
+ ExecutionIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity
+}
+import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PortIdentity}
+import org.apache.texera.amber.util.serde.GlobalPortIdentitySerde.SerdeOps
+import org.apache.texera.dao.MockTexeraDB
+import org.apache.texera.dao.jooq.generated.Tables._
+import org.apache.texera.dao.jooq.generated.tables.daos.{
+ UserDao,
+ WorkflowDao,
+ WorkflowExecutionsDao,
+ WorkflowVersionDao
+}
+import org.apache.texera.dao.jooq.generated.tables.pojos.{
+ User,
+ Workflow,
+ WorkflowExecutions,
+ WorkflowVersion
+}
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach, PrivateMethodTester}
+
+import java.net.URI
+import java.sql.Timestamp
+import java.util.UUID
+import java.util.concurrent.TimeUnit
+import scala.collection.mutable.ArrayBuffer
+
+class WorkflowExecutionsResourceSpec
+ extends AnyFlatSpec
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with MockTexeraDB
+ with PrivateMethodTester {
+
+ private val testWorkflowWid = 3000 + scala.util.Random.nextInt(1000)
+ private val testUserId = 1000 + scala.util.Random.nextInt(1000)
+
+ private var testWorkflow: Workflow = _
+ private var testVersion: WorkflowVersion = _
+ private var testUser: User = _
+ private var userDao: UserDao = _
+ private var workflowDao: WorkflowDao = _
+ private var workflowVersionDao: WorkflowVersionDao = _
+ private var workflowExecutionsDao: WorkflowExecutionsDao = _
+
+ override protected def beforeAll(): Unit = {
+ initializeDBAndReplaceDSLContext()
+ }
+
+ override protected def beforeEach(): Unit = {
+ testUser = new User
+ testUser.setUid(testUserId)
+ testUser.setName("test_user")
+ testUser.setEmail("test@example.com")
+ testUser.setPassword("password")
+ testUser.setGoogleAvatar("avatar_url")
+
+ testWorkflow = new Workflow
+ testWorkflow.setWid(testWorkflowWid)
+ testWorkflow.setName("test_workflow_" + UUID.randomUUID().toString.substring(0, 8))
+ testWorkflow.setContent("{}")
+ testWorkflow.setDescription("test description")
+ testWorkflow.setCreationTime(new Timestamp(System.currentTimeMillis()))
+ testWorkflow.setLastModifiedTime(new Timestamp(System.currentTimeMillis()))
+
+ testVersion = new WorkflowVersion
+ testVersion.setWid(testWorkflowWid)
+ testVersion.setContent("{}")
+ testVersion.setCreationTime(new Timestamp(System.currentTimeMillis()))
+
+ workflowDao = new WorkflowDao(getDSLContext.configuration())
+ workflowVersionDao = new WorkflowVersionDao(getDSLContext.configuration())
+ userDao = new UserDao(getDSLContext.configuration())
+ workflowExecutionsDao = new WorkflowExecutionsDao(getDSLContext.configuration())
+
+ cleanupTestData()
+
+ userDao.insert(testUser)
+ workflowDao.insert(testWorkflow)
+ workflowVersionDao.insert(testVersion)
+ }
+
+ override protected def afterEach(): Unit = {
+ cleanupTestData()
+ }
+
+ private def cleanupTestData(): Unit = {
+ getDSLContext
+ .deleteFrom(WORKFLOW_EXECUTIONS)
+ .where(
+ WORKFLOW_EXECUTIONS.VID.in(
+ getDSLContext
+ .select(WORKFLOW_VERSION.VID)
+ .from(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.WID.eq(testWorkflowWid))
+ )
+ )
+ .execute()
+
+ getDSLContext
+ .deleteFrom(WORKFLOW_VERSION)
+ .where(WORKFLOW_VERSION.WID.eq(testWorkflowWid))
+ .execute()
+
+ getDSLContext
+ .deleteFrom(WORKFLOW)
+ .where(WORKFLOW.WID.eq(testWorkflowWid))
+ .execute()
+
+ getDSLContext
+ .deleteFrom(USER)
+ .where(USER.UID.eq(testUserId))
+ .execute()
+ }
+
+ override protected def afterAll(): Unit = {
+ shutdownDB()
+ }
+
+ "WorkflowExecutionsResource.getWorkflowExecutions" should "return executions with EIDs in descending order" in {
+ val numExecutions = 10
+ val executionIds = ArrayBuffer.empty[Integer]
+
+ for (i <- 1 to numExecutions) {
+ val execution = new WorkflowExecutions
+ execution.setVid(testVersion.getVid)
+ execution.setUid(testUser.getUid)
+ execution.setStatus(0.toByte)
+ execution.setResult("")
+ execution.setStartingTime(
+ new Timestamp(System.currentTimeMillis() - TimeUnit.DAYS.toMillis(numExecutions - i))
+ )
+ execution.setBookmarked(false)
+ execution.setName(s"Execution ${i}")
+ execution.setEnvironmentVersion("test-env-1.0")
+
+ workflowExecutionsDao.insert(execution)
+ executionIds.append(execution.getEid)
+ }
+
+ val result = WorkflowExecutionsResource.getWorkflowExecutions(testWorkflowWid, getDSLContext)
+
+ assert(result.nonEmpty, "Result should not be empty")
+ assert(
+ result.size == numExecutions,
+ s"Expected $numExecutions executions, but got ${result.size}"
+ )
+
+ for (i <- 0 until result.size - 1) {
+ assert(
+ result(i).eId > result(i + 1).eId,
+ s"Executions are not in descending order: ${result(i).eId} should be > ${result(i + 1).eId}"
+ )
+ }
+
+ val returnedIds = result.map(_.eId).toSet
+ assert(
+ executionIds.toSet.subsetOf(returnedIds),
+ "All inserted execution IDs should be returned"
+ )
+ }
+
+ "WorkflowExecutionsResource.insertOperatorPortResultUri" should "insert a result URI row" in {
+ val execution = new WorkflowExecutions
+ execution.setVid(testVersion.getVid)
+ execution.setUid(testUser.getUid)
+ execution.setStatus(0.toByte)
+ execution.setResult("")
+ execution.setStartingTime(new Timestamp(System.currentTimeMillis()))
+ execution.setBookmarked(false)
+ execution.setName("Execution with duplicate result URI insert")
+ execution.setEnvironmentVersion("test-env-1.0")
+ workflowExecutionsDao.insert(execution)
+
+ val executionId = ExecutionIdentity(execution.getEid.longValue())
+ val globalPortId = GlobalPortIdentity(
+ PhysicalOpIdentity(OperatorIdentity("operator-1"), "main"),
+ PortIdentity(),
+ input = false
+ )
+ val uri = URI.create("vfs:///test-result")
+
+ WorkflowExecutionsResource.insertOperatorPortResultUri(executionId, globalPortId, uri)
+
+ val rows = getDSLContext
+ .selectFrom(OPERATOR_PORT_EXECUTIONS)
+ .where(OPERATOR_PORT_EXECUTIONS.WORKFLOW_EXECUTION_ID.eq(execution.getEid))
+ .and(OPERATOR_PORT_EXECUTIONS.GLOBAL_PORT_ID.eq(globalPortId.serializeAsString))
+ .fetch()
+
+ assert(rows.size() == 1)
+ assert(rows.get(0).getResultUri == uri.toString)
+ }
+
+}
diff --git a/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowVersionResourceSpec.scala b/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowVersionResourceSpec.scala
new file mode 100644
index 00000000000..ecf704f663b
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/web/resource/dashboard/user/workflow/WorkflowVersionResourceSpec.scala
@@ -0,0 +1,195 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.dashboard.user.workflow
+
+import org.apache.texera.amber.util.JSONUtils.objectMapper
+import org.apache.texera.dao.MockTexeraDB
+import org.apache.texera.dao.jooq.generated.Tables
+import org.apache.texera.dao.jooq.generated.tables.daos.{WorkflowDao, WorkflowVersionDao}
+import org.apache.texera.dao.jooq.generated.tables.pojos.{Workflow, WorkflowVersion}
+import org.jooq.impl.DSL
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
+
+import java.sql.Timestamp
+import java.util.UUID
+import java.util.concurrent.TimeUnit
+import scala.collection.mutable.ArrayBuffer
+
+class WorkflowVersionResourceSpec
+ extends AnyFlatSpec
+ with BeforeAndAfterAll
+ with BeforeAndAfterEach
+ with MockTexeraDB {
+
+ private val testWorkflowWid = 2000 + scala.util.Random.nextInt(1000)
+
+ private var testWorkflow: Workflow = _
+ private var workflowDao: WorkflowDao = _
+ private var workflowVersionDao: WorkflowVersionDao = _
+
+ private val capturedVersions = ArrayBuffer.empty[Integer]
+
+ override protected def beforeAll(): Unit = {
+ initializeDBAndReplaceDSLContext()
+ }
+
+ override protected def beforeEach(): Unit = {
+ testWorkflow = new Workflow
+ testWorkflow.setWid(Integer.valueOf(testWorkflowWid))
+ testWorkflow.setName("test_workflow_" + UUID.randomUUID().toString.substring(0, 8))
+ testWorkflow.setContent(createWorkflowContent("initial"))
+ testWorkflow.setDescription("test description")
+
+ workflowDao = new WorkflowDao(getDSLContext.configuration())
+ workflowVersionDao = new WorkflowVersionDao(getDSLContext.configuration())
+
+ cleanupTestData()
+ workflowDao.insert(testWorkflow)
+ capturedVersions.clear()
+ }
+
+ override protected def afterEach(): Unit = {
+ cleanupTestData()
+ }
+
+ private def cleanupTestData(): Unit = {
+ getDSLContext
+ .deleteFrom(Tables.WORKFLOW_VERSION)
+ .where(Tables.WORKFLOW_VERSION.WID.eq(testWorkflowWid))
+ .execute()
+
+ getDSLContext
+ .deleteFrom(Tables.WORKFLOW)
+ .where(Tables.WORKFLOW.WID.eq(testWorkflowWid))
+ .execute()
+ }
+
+ override protected def afterAll(): Unit = {
+ shutdownDB()
+ }
+
+ private def createWorkflowContent(value: String): String = {
+ val jsonNode = objectMapper.createObjectNode()
+ jsonNode.put("value", value)
+ jsonNode.toString
+ }
+
+ private def createVersionDiff(oldValue: String, newValue: String): String = {
+ val oldJson = objectMapper.createObjectNode()
+ oldJson.put("value", oldValue)
+
+ val newJson = objectMapper.createObjectNode()
+ newJson.put("value", newValue)
+
+ val patch = com.flipkart.zjsonpatch.JsonDiff.asJson(
+ oldJson,
+ newJson
+ )
+ patch.toString
+ }
+
+ "WorkflowVersionResource" should "return versions in descending order from fetchSubsequentVersions and apply patches correctly" in {
+ var currentContent = "initial"
+ for (i <- 1 to 10) {
+ val newContent = s"version_$i"
+ val diffContent = createVersionDiff(currentContent, newContent)
+
+ val version = new WorkflowVersion
+ version.setWid(testWorkflow.getWid)
+ version.setContent(diffContent)
+ version.setCreationTime(
+ new Timestamp(System.currentTimeMillis() - TimeUnit.MINUTES.toMillis(10 - i))
+ )
+ workflowVersionDao.insert(version)
+
+ currentContent = newContent
+ }
+
+ testWorkflow.setContent(createWorkflowContent(currentContent))
+ workflowDao.update(testWorkflow)
+
+ val midVersionId = 5
+ val versions = WorkflowVersionResource.fetchSubsequentVersions(
+ testWorkflow.getWid,
+ midVersionId,
+ getDSLContext
+ )
+
+ assert(versions.nonEmpty, "No versions were returned")
+
+ for (i <- 0 until versions.length - 1) {
+ assert(
+ versions(i).getVid > versions(i + 1).getVid,
+ s"Versions not in descending order: ${versions(i).getVid} should be > ${versions(i + 1).getVid}"
+ )
+ }
+
+ val highestVersionId = getDSLContext
+ .select(DSL.max(Tables.WORKFLOW_VERSION.VID))
+ .from(Tables.WORKFLOW_VERSION)
+ .where(Tables.WORKFLOW_VERSION.WID.eq(testWorkflowWid))
+ .fetchOneInto(classOf[Integer])
+
+ assert(versions.head.getVid === highestVersionId, "First version should have the highest VID")
+
+ capturedVersions.clear()
+ versions.foreach(v => capturedVersions.append(v.getVid))
+
+ val workflowFromDb = workflowDao.fetchOneByWid(testWorkflow.getWid)
+
+ val workflowVersionDirect = WorkflowVersionResource.applyPatch(versions, workflowFromDb)
+ val directVersionContent =
+ objectMapper.readTree(workflowVersionDirect.getContent).get("value").asText()
+
+ assert(
+ directVersionContent === s"version_$midVersionId",
+ s"Workflow content from direct applyPatch should be 'version_$midVersionId' but was '$directVersionContent'"
+ )
+
+ val combinedVersions = WorkflowVersionResource.fetchSubsequentVersions(
+ testWorkflow.getWid,
+ midVersionId,
+ getDSLContext
+ )
+ val currentWorkflowForCombined = workflowDao.fetchOneByWid(testWorkflow.getWid)
+ val workflowVersion =
+ WorkflowVersionResource.applyPatch(combinedVersions, currentWorkflowForCombined)
+
+ assert(capturedVersions.nonEmpty, "No versions were captured")
+ assert(
+ capturedVersions.length === versions.length,
+ "Captured versions length doesn't match fetched versions"
+ )
+
+ for (i <- versions.indices) {
+ assert(
+ capturedVersions(i) === versions(i).getVid,
+ s"Captured version ${capturedVersions(i)} doesn't match fetched version ${versions(i).getVid} at index $i"
+ )
+ }
+
+ val midVersionContent = objectMapper.readTree(workflowVersion.getContent).get("value").asText()
+ assert(
+ midVersionContent === s"version_$midVersionId",
+ s"Workflow content should be 'version_$midVersionId' but was '$midVersionContent'"
+ )
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveResourceSpec.scala b/amber/src/test/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveResourceSpec.scala
new file mode 100644
index 00000000000..a093cf1ad2f
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveResourceSpec.scala
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.resource.pythonvirtualenvironment
+
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+import java.nio.file.{Files, Path, Paths}
+import java.util.concurrent.LinkedBlockingQueue
+import scala.jdk.CollectionConverters._
+
+class PveResourceSpec extends AnyFlatSpec with Matchers with BeforeAndAfterEach {
+
+ private val testCuid = 256
+ private var testPveName: String = _
+ private var testRoot: Path = _
+ private var queue: LinkedBlockingQueue[String] = _
+
+ override protected def beforeEach(): Unit = {
+ testPveName = s"testenv${System.currentTimeMillis()}"
+ testRoot = Paths.get("/tmp/texera-pve/venvs").resolve(testCuid.toString)
+ queue = new LinkedBlockingQueue[String]()
+ }
+
+ override protected def afterEach(): Unit = {
+ PveManager.deleteEnvironments(testCuid)
+ }
+
+ private def queueText(): String = {
+ queue.iterator().asScala.toList.mkString("\n")
+ }
+
+ "PveManager" should "create a new PVE and list it" in {
+ PveManager.createNewPve(testCuid, queue, testPveName, isLocal = true)
+
+ val logs = queueText()
+
+ logs should not include "[PVE][ERR]"
+ logs should include(s"[PVE] Created new environment for cuid = $testCuid")
+
+ val pvePath = testRoot.resolve(testPveName).resolve("pve")
+ val pythonPath = pvePath.resolve("bin").resolve("python")
+ val pipPath = pvePath.resolve("bin").resolve("pip")
+
+ Files.exists(pvePath) shouldBe true
+ Files.exists(pythonPath) shouldBe true
+ Files.exists(pipPath) shouldBe true
+
+ PveManager.getEnvironments(testCuid) should contain(testPveName)
+ }
+
+ "PveManager" should "delete all PVEs for a computing unit" in {
+ PveManager.createNewPve(testCuid, queue, testPveName, isLocal = true)
+
+ Files.exists(testRoot.resolve(testPveName)) shouldBe true
+
+ PveManager.deleteEnvironments(testCuid)
+
+ Files.exists(testRoot) shouldBe false
+ PveManager.getEnvironments(testCuid) shouldBe empty
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/web/service/ExecutionConsoleServiceSpec.scala b/amber/src/test/scala/org/apache/texera/web/service/ExecutionConsoleServiceSpec.scala
new file mode 100644
index 00000000000..d4753984cf1
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/web/service/ExecutionConsoleServiceSpec.scala
@@ -0,0 +1,234 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import com.google.protobuf.timestamp.Timestamp
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ ConsoleMessage,
+ ConsoleMessageType
+}
+import org.apache.texera.amber.engine.common.executionruntimestate.ExecutionConsoleStore
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+import java.time.Instant
+
+class ExecutionConsoleServiceSpec extends AnyFlatSpec with Matchers {
+
+ // Constants for testing
+ val standardBufferSize: Int = 100
+ val smallBufferSize: Int = 2
+ val messageDisplayLength: Int = 100
+
+ "processConsoleMessage" should "truncate message title when it exceeds display length" in {
+ // Create a long message title that exceeds display length
+ val longTitle = "a" * (messageDisplayLength + 10)
+ val expectedTruncatedTitle = "a" * (messageDisplayLength - 3) + "..."
+
+ // Create a console message with a long title
+ val consoleMessage = new ConsoleMessage(
+ "worker1",
+ Timestamp(Instant.now),
+ ConsoleMessageType.PRINT,
+ "test",
+ longTitle,
+ "message content"
+ )
+
+ // Call the method under test
+ val processedMessage =
+ ConsoleMessageProcessor.processConsoleMessage(consoleMessage, messageDisplayLength)
+
+ // Verify the title was truncated
+ processedMessage.title shouldBe expectedTruncatedTitle
+ }
+
+ it should "not truncate message title when it does not exceed display length" in {
+ // Create a short message title that doesn't exceed display length
+ val shortTitle = "Short Title"
+
+ // Create a console message with a short title
+ val consoleMessage = new ConsoleMessage(
+ "worker1",
+ Timestamp(Instant.now),
+ ConsoleMessageType.PRINT,
+ "test",
+ shortTitle,
+ "message content"
+ )
+
+ // Call the method under test
+ val processedMessage =
+ ConsoleMessageProcessor.processConsoleMessage(consoleMessage, messageDisplayLength)
+
+ // Verify the title was not truncated
+ processedMessage.title shouldBe shortTitle
+ }
+
+ "addMessageToOperatorConsole" should "add message to buffer when buffer is not full" in {
+ // Create a test console store
+ val consoleStore = new ExecutionConsoleStore()
+ val opId = "op1"
+
+ // Create console messages
+ val message1 = new ConsoleMessage(
+ "worker1",
+ Timestamp(Instant.now),
+ ConsoleMessageType.PRINT,
+ "test",
+ "Message 1",
+ "content 1"
+ )
+
+ val message2 = new ConsoleMessage(
+ "worker1",
+ Timestamp(Instant.now),
+ ConsoleMessageType.PRINT,
+ "test",
+ "Message 2",
+ "content 2"
+ )
+
+ // Add first message
+ val storeWithMessage1 =
+ ConsoleMessageProcessor.addMessageToOperatorConsole(
+ consoleStore,
+ opId,
+ message1,
+ standardBufferSize
+ )
+
+ // Add second message
+ val storeWithMessage2 = ConsoleMessageProcessor.addMessageToOperatorConsole(
+ storeWithMessage1,
+ opId,
+ message2,
+ standardBufferSize
+ )
+
+ // Verify both messages are in the buffer
+ val opInfo = storeWithMessage2.operatorConsole(opId)
+ opInfo.consoleMessages.size shouldBe 2
+ opInfo.consoleMessages.head.title shouldBe "Message 1"
+ opInfo.consoleMessages(1).title shouldBe "Message 2"
+ }
+
+ it should "remove oldest message when buffer is full" in {
+ // Create a test console store
+ val consoleStore = new ExecutionConsoleStore()
+ val opId = "op1"
+
+ // Create console messages
+ val message1 = new ConsoleMessage(
+ "worker1",
+ Timestamp(Instant.now),
+ ConsoleMessageType.PRINT,
+ "test",
+ "Message 1",
+ "content 1"
+ )
+
+ val message2 = new ConsoleMessage(
+ "worker1",
+ Timestamp(Instant.now),
+ ConsoleMessageType.PRINT,
+ "test",
+ "Message 2",
+ "content 2"
+ )
+
+ val message3 = new ConsoleMessage(
+ "worker1",
+ Timestamp(Instant.now),
+ ConsoleMessageType.PRINT,
+ "test",
+ "Message 3",
+ "content 3"
+ )
+
+ // Fill the buffer
+ val storeWithMessage1 =
+ ConsoleMessageProcessor.addMessageToOperatorConsole(
+ consoleStore,
+ opId,
+ message1,
+ smallBufferSize
+ )
+ val storeWithMessage2 =
+ ConsoleMessageProcessor.addMessageToOperatorConsole(
+ storeWithMessage1,
+ opId,
+ message2,
+ smallBufferSize
+ )
+
+ // Add one more message which should remove the oldest
+ val storeWithMessage3 =
+ ConsoleMessageProcessor.addMessageToOperatorConsole(
+ storeWithMessage2,
+ opId,
+ message3,
+ smallBufferSize
+ )
+
+ // Verify the first message was removed and only the second and third remain
+ val opInfo = storeWithMessage3.operatorConsole(opId)
+ opInfo.consoleMessages.size shouldBe 2
+ opInfo.consoleMessages.head.title shouldBe "Message 2"
+ opInfo.consoleMessages(1).title shouldBe "Message 3"
+ }
+
+ "the complete message processing flow" should "handle messages correctly" in {
+ // Create a test console store
+ val consoleStore = new ExecutionConsoleStore()
+ val opId = "op1"
+
+ // Create a message with a title that needs truncation
+ val longTitle = "a" * (messageDisplayLength + 10)
+ val consoleMessage = new ConsoleMessage(
+ "worker1",
+ Timestamp(Instant.now),
+ ConsoleMessageType.PRINT,
+ "test",
+ longTitle,
+ "message content"
+ )
+
+ // Process the message first
+ val processedMessage =
+ ConsoleMessageProcessor.processConsoleMessage(consoleMessage, messageDisplayLength)
+
+ // Then update the store
+ val updatedStore = ConsoleMessageProcessor.addMessageToOperatorConsole(
+ consoleStore,
+ opId,
+ processedMessage,
+ standardBufferSize
+ )
+
+ // Verify correct processing
+ val opInfo = updatedStore.operatorConsole(opId)
+ opInfo.consoleMessages.size shouldBe 1
+
+ // Check that title was truncated
+ val expectedTruncatedTitle = "a" * (messageDisplayLength - 3) + "..."
+ opInfo.consoleMessages.head.title shouldBe expectedTruncatedTitle
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/web/service/ExecutionReconfigurationServiceSpec.scala b/amber/src/test/scala/org/apache/texera/web/service/ExecutionReconfigurationServiceSpec.scala
new file mode 100644
index 00000000000..974db13286d
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/web/service/ExecutionReconfigurationServiceSpec.scala
@@ -0,0 +1,149 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import org.apache.texera.amber.core.executor.OpExecWithClassName
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ExecutionIdentity,
+ OperatorIdentity,
+ PhysicalOpIdentity,
+ WorkflowIdentity
+}
+import org.apache.texera.amber.core.workflow.PhysicalOp
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.WorkflowReconfigureRequest
+import org.apache.texera.web.storage.{ExecutionReconfigurationStore, ExecutionStateStore}
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+import scala.collection.mutable.ArrayBuffer
+
+/**
+ * Web-service-layer tests for ExecutionReconfigurationService.
+ *
+ * The end-to-end engine path (reconfigureWorkflow → Fries algorithm →
+ * UpdateExecutor on workers) is covered by ReconfigurationSpec.
+ * This spec focuses on the wiring inside performReconfigurationOnResume:
+ * empty short-circuit, request construction, and store reset semantics.
+ */
+class ExecutionReconfigurationServiceSpec extends AnyFlatSpec with Matchers {
+
+ private def mkPhysicalOp(name: String): PhysicalOp =
+ PhysicalOp(
+ id = PhysicalOpIdentity(OperatorIdentity(name), "main"),
+ workflowId = WorkflowIdentity(0L),
+ executionId = ExecutionIdentity(0L),
+ opExecInitInfo = OpExecWithClassName(s"$name.Class", "")
+ )
+
+ /** Service variant that records dispatched requests and skips the AmberClient
+ * registration / workflow-dependent diff handler so it can be constructed
+ * without a live engine.
+ */
+ private class RecordingService(stateStore: ExecutionStateStore)
+ extends ExecutionReconfigurationService(client = null, stateStore, workflow = null) {
+ val captured: ArrayBuffer[WorkflowReconfigureRequest] = ArrayBuffer.empty
+ override protected def dispatch(request: WorkflowReconfigureRequest): Unit =
+ captured += request
+ override protected def registerWorkerCompletionCallback(): Unit = ()
+ override protected def registerCompletionDiffHandler(): Unit = ()
+ }
+
+ "performReconfigurationOnResume" should
+ "return without dispatching when no reconfigurations are pending" in {
+ val stateStore = new ExecutionStateStore()
+ val service = new RecordingService(stateStore)
+
+ noException should be thrownBy service.performReconfigurationOnResume()
+
+ service.captured shouldBe empty
+ val state = stateStore.reconfigurationStore.getState
+ state.unscheduledReconfigurations shouldBe empty
+ state.currentReconfigId shouldBe None
+ state.completedReconfigurations shouldBe empty
+ }
+
+ it should "dispatch one request carrying every pending reconfiguration and reset the store" in {
+ val stateStore = new ExecutionStateStore()
+ val service = new RecordingService(stateStore)
+
+ val op1 = mkPhysicalOp("op-1")
+ val op2 = mkPhysicalOp("op-2")
+ stateStore.reconfigurationStore.updateState(_ =>
+ ExecutionReconfigurationStore(unscheduledReconfigurations = List((op1, None), (op2, None)))
+ )
+
+ service.performReconfigurationOnResume()
+
+ service.captured should have size 1
+ val request = service.captured.head
+ request.reconfigurationId should not be empty
+ request.reconfiguration.map(_.targetOpId) should contain theSameElementsInOrderAs Seq(
+ op1.id,
+ op2.id
+ )
+ request.reconfiguration.map(_.newExecInitInfo) should contain theSameElementsInOrderAs Seq(
+ op1.opExecInitInfo,
+ op2.opExecInitInfo
+ )
+
+ val state = stateStore.reconfigurationStore.getState
+ state.unscheduledReconfigurations shouldBe empty
+ state.currentReconfigId shouldBe Some(request.reconfigurationId)
+ state.completedReconfigurations shouldBe empty
+ }
+
+ it should "use a fresh reconfigurationId on each dispatch" in {
+ val stateStore = new ExecutionStateStore()
+ val service = new RecordingService(stateStore)
+
+ def queueAndDispatch(opName: String): String = {
+ stateStore.reconfigurationStore.updateState(old =>
+ old.copy(unscheduledReconfigurations = List((mkPhysicalOp(opName), None)))
+ )
+ service.performReconfigurationOnResume()
+ service.captured.last.reconfigurationId
+ }
+
+ val firstId = queueAndDispatch("op-a")
+ val secondId = queueAndDispatch("op-b")
+
+ firstId should not be secondId
+ stateStore.reconfigurationStore.getState.currentReconfigId shouldBe Some(secondId)
+ }
+
+ "onWorkerReconfigured" should
+ "add the worker id to completedReconfigurations so the diff handler can fire" in {
+ val stateStore = new ExecutionStateStore()
+ val service = new RecordingService(stateStore)
+
+ val w1 = ActorVirtualIdentity("Worker:WF1-E1-op-main-0")
+ val w2 = ActorVirtualIdentity("Worker:WF1-E1-op-main-1")
+ service.onWorkerReconfigured(w1)
+ service.onWorkerReconfigured(w2)
+ // duplicate completion is idempotent (Set semantics).
+ service.onWorkerReconfigured(w1)
+
+ stateStore.reconfigurationStore.getState.completedReconfigurations should contain theSameElementsAs Set(
+ w1,
+ w2
+ )
+ }
+}
diff --git a/amber/src/test/scala/org/apache/texera/web/service/ExecutionResultServiceSpec.scala b/amber/src/test/scala/org/apache/texera/web/service/ExecutionResultServiceSpec.scala
new file mode 100644
index 00000000000..d9b3f60e6ff
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/web/service/ExecutionResultServiceSpec.scala
@@ -0,0 +1,478 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.web.service
+
+import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema, Tuple}
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+class ExecutionResultServiceSpec extends AnyFlatSpec with Matchers {
+
+ "convertTuplesToJson" should "convert tuples with various field types correctly" in {
+ // Create a schema with different attribute types
+ val attributes = List(
+ new Attribute("stringCol", AttributeType.STRING),
+ new Attribute("intCol", AttributeType.INTEGER),
+ new Attribute("boolCol", AttributeType.BOOLEAN),
+ new Attribute("nullCol", AttributeType.ANY),
+ new Attribute("longStringCol", AttributeType.STRING),
+ new Attribute("shortBinaryCol", AttributeType.BINARY),
+ new Attribute("longBinaryCol", AttributeType.BINARY)
+ )
+
+ val schema = new Schema(attributes)
+
+ // Create a string longer than maxStringLength (100)
+ val longString = "a" * 150
+
+ // Create binary data
+ val shortBinaryData = Array[Byte](1, 2, 3, 4, 5)
+ val longBinaryData = Array.tabulate[Byte](100)(_.toByte)
+
+ // Create a tuple with all the test data
+ val tuple = Tuple
+ .builder(schema)
+ .add("stringCol", AttributeType.STRING, "regular string")
+ .add("intCol", AttributeType.INTEGER, 42)
+ .add("boolCol", AttributeType.BOOLEAN, true)
+ .add("nullCol", AttributeType.ANY, null)
+ .add("longStringCol", AttributeType.STRING, longString)
+ .add("shortBinaryCol", AttributeType.BINARY, shortBinaryData)
+ .add("longBinaryCol", AttributeType.BINARY, longBinaryData)
+ .build()
+
+ // Convert to JSON
+ val result = ExecutionResultService.convertTuplesToJson(List(tuple))
+
+ // Verify the result
+ result should have size 1
+ val jsonNode = result.head
+
+ // Check regular values
+ jsonNode.get("stringCol").asText() shouldBe "regular string"
+ jsonNode.get("intCol").asInt() shouldBe 42
+ jsonNode.get("boolCol").asBoolean() shouldBe true
+
+ // Check NULL value
+ jsonNode.get("nullCol").asText() shouldBe "NULL"
+
+ // Check long string truncation
+ jsonNode.get("longStringCol").asText() should (
+ have length 103 and // 100 chars + "..."
+ startWith("a" * 100) and
+ endWith("...")
+ )
+
+ // Check short binary representation
+ val shortBinaryString = jsonNode.get("shortBinaryCol").asText()
+ shortBinaryString should (
+ startWith("bytes'") and
+ include("01 02 03 04 05") and
+ include("(length: 5)")
+ )
+
+ // Check long binary representation
+ val longBinaryString = jsonNode.get("longBinaryCol").asText()
+ longBinaryString should (
+ startWith("bytes'") and
+ include("...") and
+ include("(length: 100)")
+ )
+ }
+
+ it should "handle empty collections of tuples" in {
+ val result = ExecutionResultService.convertTuplesToJson(List())
+ result shouldBe empty
+ }
+
+ it should "handle collections with multiple tuples" in {
+ // Create a simple schema
+ val attributes = List(
+ new Attribute("id", AttributeType.INTEGER),
+ new Attribute("name", AttributeType.STRING)
+ )
+
+ val schema = new Schema(attributes)
+
+ // Create multiple tuples
+ val tuple1 = Tuple
+ .builder(schema)
+ .add("id", AttributeType.INTEGER, 1)
+ .add("name", AttributeType.STRING, "Alice")
+ .build()
+
+ val tuple2 = Tuple
+ .builder(schema)
+ .add("id", AttributeType.INTEGER, 2)
+ .add("name", AttributeType.STRING, "Bob")
+ .build()
+
+ // Convert to JSON
+ val results = ExecutionResultService.convertTuplesToJson(List(tuple1, tuple2))
+
+ // Verify the results
+ results should have size 2
+ results.head.get("id").asInt() shouldBe 1
+ results.head.get("name").asText() shouldBe "Alice"
+ results(1).get("id").asInt() shouldBe 2
+ results(1).get("name").asText() shouldBe "Bob"
+ }
+
+ it should "handle string exactly at the maximum length" in {
+ val attributes = List(
+ new Attribute("exactLengthString", AttributeType.STRING)
+ )
+ val schema = new Schema(attributes)
+
+ // Create string exactly at maxStringLength (100)
+ val exactLengthString = "x" * 100
+
+ val tuple = Tuple
+ .builder(schema)
+ .add("exactLengthString", AttributeType.STRING, exactLengthString)
+ .build()
+
+ val result = ExecutionResultService.convertTuplesToJson(List(tuple))
+
+ result should have size 1
+ val jsonNode = result.head
+
+ jsonNode.get("exactLengthString").asText() shouldBe exactLengthString
+ jsonNode.get("exactLengthString").asText() should have length 100
+ }
+
+ it should "handle empty binary data" in {
+ val attributes = List(
+ new Attribute("emptyBinary", AttributeType.BINARY)
+ )
+ val schema = new Schema(attributes)
+
+ // Empty binary data
+ val emptyBinaryData = Array[Byte]()
+
+ val tuple = Tuple
+ .builder(schema)
+ .add("emptyBinary", AttributeType.BINARY, emptyBinaryData)
+ .build()
+
+ val result = ExecutionResultService.convertTuplesToJson(List(tuple))
+
+ result should have size 1
+ val jsonNode = result.head
+
+ val emptyBinaryString = jsonNode.get("emptyBinary").asText()
+ emptyBinaryString should include("(length: 0)")
+ }
+
+ it should "handle binary data with single ByteBuffer" in {
+ val attributes = List(
+ new Attribute("singleBufferBinary", AttributeType.BINARY)
+ )
+ val schema = new Schema(attributes)
+
+ // Create binary data with a single ByteBuffer
+ val singleBufferData = "Hello, world!".getBytes()
+
+ val tuple = Tuple
+ .builder(schema)
+ .add("singleBufferBinary", AttributeType.BINARY, singleBufferData)
+ .build()
+
+ val result = ExecutionResultService.convertTuplesToJson(List(tuple))
+
+ result should have size 1
+ val jsonNode = result.head
+
+ val binaryString = jsonNode.get("singleBufferBinary").asText()
+ binaryString should (
+ startWith("bytes'") and
+ include("(length: 13)") // "Hello, world!" is 13 bytes
+ )
+ }
+
+ it should "handle various numeric types correctly" in {
+ val attributes = List(
+ new Attribute("intValue", AttributeType.INTEGER),
+ new Attribute("doubleValue", AttributeType.DOUBLE),
+ new Attribute("longValue", AttributeType.LONG)
+ )
+ val schema = new Schema(attributes)
+
+ val tuple = Tuple
+ .builder(schema)
+ .add("intValue", AttributeType.INTEGER, Int.MaxValue)
+ .add("doubleValue", AttributeType.DOUBLE, 3.14159)
+ .add("longValue", AttributeType.LONG, Long.MaxValue)
+ .build()
+
+ val result = ExecutionResultService.convertTuplesToJson(List(tuple))
+
+ result should have size 1
+ val jsonNode = result.head
+
+ jsonNode.get("intValue").asInt() shouldBe Int.MaxValue
+ jsonNode.get("doubleValue").asDouble() shouldBe 3.14159
+ jsonNode.get("longValue").asLong() shouldBe Long.MaxValue
+ }
+
+ it should "handle multiple binary fields within the same tuple" in {
+ val attributes = List(
+ new Attribute("binaryField1", AttributeType.BINARY),
+ new Attribute("binaryField2", AttributeType.BINARY)
+ )
+ val schema = new Schema(attributes)
+
+ val binaryData1 = Array[Byte](10, 20, 30)
+ val binaryData2 = Array[Byte](40, 50, 60)
+
+ val tuple = Tuple
+ .builder(schema)
+ .add("binaryField1", AttributeType.BINARY, binaryData1)
+ .add("binaryField2", AttributeType.BINARY, binaryData2)
+ .build()
+
+ val result = ExecutionResultService.convertTuplesToJson(List(tuple))
+
+ result should have size 1
+ val jsonNode = result.head
+
+ val binaryString1 = jsonNode.get("binaryField1").asText()
+ binaryString1 should (
+ include("0A 14 1E") and // Hex representation of 10, 20, 30
+ include("(length: 3)")
+ )
+
+ val binaryString2 = jsonNode.get("binaryField2").asText()
+ binaryString2 should (
+ include("28 32 3C") and // Hex representation of 40, 50, 60
+ include("(length: 3)")
+ )
+ }
+
+ it should "not truncate long strings when isVisualization is true" in {
+ val attributes = List(
+ new Attribute("longStringCol", AttributeType.STRING)
+ )
+ val schema = new Schema(attributes)
+
+ // Create a string longer than maxStringLength (100)
+ val longString = "a" * 150
+ val htmlVisualizationString = """
+
+
+
+
+
+
+