diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 000000000000..8ef26023dd50 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,2 @@ +docker/development/docker-volume +docker/docker-compose-quickstart/docker-volume diff --git a/.github/actions/setup-openmetadata-test-environment/action.yml b/.github/actions/setup-openmetadata-test-environment/action.yml index de417eab6b04..afe52f9938fe 100644 --- a/.github/actions/setup-openmetadata-test-environment/action.yml +++ b/.github/actions/setup-openmetadata-test-environment/action.yml @@ -15,6 +15,10 @@ inputs: description: Arguments to pass to run_local_docker.sh required: false default: "-m no-ui -d mysql" # Use "-d postgresql" for postgres and Opensearch + startup-script: + description: Startup script used to launch the local OpenMetadata test environment + required: false + default: "./docker/run_local_docker.sh" ingestion_dependency: description: Ingestion dependency to pass to run_local_docker.sh required: false @@ -97,4 +101,4 @@ runs: timeout_minutes: 60 max_attempts: 2 retry_on: error - command: ./docker/run_local_docker.sh ${{ inputs.args }} + command: ${{ inputs.startup-script }} ${{ inputs.args }} diff --git a/.github/workflows/playwright-knowledge-graph-postgresql-e2e.yml b/.github/workflows/playwright-knowledge-graph-postgresql-e2e.yml new file mode 100644 index 000000000000..f2642eb5761d --- /dev/null +++ b/.github/workflows/playwright-knowledge-graph-postgresql-e2e.yml @@ -0,0 +1,224 @@ +# Copyright 2021 Collate +# 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. + +name: Postgresql PR Knowledge Graph E2E Tests +on: + workflow_dispatch: + pull_request: + types: + - opened + - synchronize + - reopened + - ready_for_review + paths: + - ".github/actions/setup-openmetadata-test-environment/action.yml" + - ".github/workflows/playwright-knowledge-graph-postgresql-e2e.yml" + - "docker/run_local_docker.sh" + - "docker/run_local_docker_common.sh" + - "docker/run_local_docker_rdf.sh" + - "docker/validate_compose.py" + - "docker/development/docker-compose-fuseki.yml" + - "docker/development/docker-compose-postgres-fuseki.yml" + - "docs/rdf-local-development.md" + - "openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/**" + - "openmetadata-service/src/main/java/org/openmetadata/service/rdf/**" + - "openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/**" + - "openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/**" + - "openmetadata-service/src/test/java/org/openmetadata/service/rdf/**" + - "openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/**" + - "openmetadata-spec/src/main/resources/rdf/**" + - "openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/KnowledgeGraph.spec.ts" + - "openmetadata-ui/src/main/resources/ui/playwright.config.ts" + - "openmetadata-ui/src/main/resources/ui/src/components/KnowledgeGraph/**" + - "openmetadata-ui/src/main/resources/ui/src/components/OntologyExplorer/**" + - "openmetadata-ui/src/main/resources/ui/src/rest/rdfAPI.ts" + - "openmetadata-ui/src/main/resources/ui/src/types/knowledgeGraph.types.ts" + - "openmetadata-ui/src/main/resources/ui/src/utils/TableUtils.tsx" + +permissions: + contents: read + +concurrency: + group: playwright-knowledge-graph-pr-postgresql-${{ github.event.pull_request.number || github.run_id }} + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-latest + if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }} + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + + - name: Setup JDK 21 + uses: actions/setup-java@v4 + with: + java-version: '21' + distribution: 'temurin' + + - name: Cache Maven Dependencies + uses: actions/cache@v4 + with: + path: ~/.m2 + key: ${{ runner.os }}-maven-${{ hashFiles('**/pom.xml') }} + restore-keys: | + ${{ runner.os }}-maven- + + - name: Install antlr cli + run: sudo make install_antlr_cli + + - name: Build with Maven + run: mvn -DskipTests clean package + + - name: Upload Maven build artifact + uses: actions/upload-artifact@v4 + with: + name: openmetadata-build + path: openmetadata-dist/target/openmetadata-*.tar.gz + retention-days: 1 + + playwright-knowledge-graph-postgresql: + needs: [build] + runs-on: ubuntu-latest + if: ${{ !cancelled() && needs.build.result == 'success' && (github.event_name != 'pull_request' || !github.event.pull_request.head.repo.fork) }} + environment: test + steps: + - name: Free Disk Space (Ubuntu) + uses: jlumbroso/free-disk-space@main + with: + tool-cache: false + android: true + dotnet: true + haskell: true + large-packages: false + swap-storage: true + docker-images: false + + - name: Checkout + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + + - name: Prepare temporary directory for Maven build artifact + run: mkdir -p "${{ runner.temp }}/openmetadata-build-artifact" + + - name: Download Maven build artifact + uses: actions/download-artifact@v4 + with: + name: openmetadata-build + path: ${{ runner.temp }}/openmetadata-build-artifact + + - name: Copy Maven build artifact into workspace + run: | + mkdir -p openmetadata-dist/target + cp -a "${{ runner.temp }}/openmetadata-build-artifact/." openmetadata-dist/target/ + + - name: Setup Openmetadata Test Environment + uses: ./.github/actions/setup-openmetadata-test-environment + with: + python-version: "3.10" + args: "-d postgresql -s true" + startup-script: "./docker/run_local_docker_rdf.sh" + ingestion_dependency: "all" + + - name: Wait for Fuseki to be healthy + run: | + echo "Verifying Fuseki is healthy before running tests..." + for i in $(seq 1 30); do + if curl -sf "http://localhost:3030/\$/ping" > /dev/null 2>&1; then + echo "Fuseki is healthy" + exit 0 + fi + echo "Waiting for Fuseki ($i/30)..." + sleep 10 + done + echo "Fuseki failed health check. Container logs:" + docker logs openmetadata-fuseki --tail 100 + exit 1 + + - name: Setup Node.js + uses: actions/setup-node@v4 + with: + node-version-file: "openmetadata-ui/src/main/resources/ui/.nvmrc" + + - name: Install dependencies + working-directory: openmetadata-ui/src/main/resources/ui/ + run: yarn --ignore-scripts --frozen-lockfile + + - name: Install Playwright Browsers + run: npx playwright@1.57.0 install chromium --with-deps + + - name: Run Knowledge Graph Playwright tests + working-directory: openmetadata-ui/src/main/resources/ui/ + run: npx playwright test --project="Knowledge Graph" + env: + PLAYWRIGHT_IS_OSS: true + PLAYWRIGHT_SNOWFLAKE_USERNAME: ${{ secrets.TEST_SNOWFLAKE_USERNAME }} + PLAYWRIGHT_SNOWFLAKE_PASSWORD: ${{ secrets.TEST_SNOWFLAKE_PASSWORD }} + PLAYWRIGHT_SNOWFLAKE_ACCOUNT: ${{ secrets.TEST_SNOWFLAKE_ACCOUNT }} + PLAYWRIGHT_SNOWFLAKE_DATABASE: ${{ secrets.TEST_SNOWFLAKE_DATABASE }} + PLAYWRIGHT_SNOWFLAKE_WAREHOUSE: ${{ secrets.TEST_SNOWFLAKE_WAREHOUSE }} + PLAYWRIGHT_SNOWFLAKE_PASSPHRASE: ${{ secrets.TEST_SNOWFLAKE_PASSPHRASE }} + PLAYWRIGHT_BQ_PRIVATE_KEY: ${{ secrets.TEST_BQ_PRIVATE_KEY }} + PLAYWRIGHT_BQ_PROJECT_ID: ${{ secrets.PLAYWRIGHT_BQ_PROJECT_ID }} + PLAYWRIGHT_BQ_PRIVATE_KEY_ID: ${{ secrets.TEST_BQ_PRIVATE_KEY_ID }} + PLAYWRIGHT_BQ_PROJECT_ID_TAXONOMY: ${{ secrets.TEST_BQ_PROJECT_ID_TAXONOMY }} + PLAYWRIGHT_BQ_CLIENT_EMAIL: ${{ secrets.TEST_BQ_CLIENT_EMAIL }} + PLAYWRIGHT_BQ_CLIENT_ID: ${{ secrets.TEST_BQ_CLIENT_ID }} + PLAYWRIGHT_REDSHIFT_HOST: ${{ secrets.E2E_REDSHIFT_HOST_PORT }} + PLAYWRIGHT_REDSHIFT_USERNAME: ${{ secrets.E2E_REDSHIFT_USERNAME }} + PLAYWRIGHT_REDSHIFT_PASSWORD: ${{ secrets.E2E_REDSHIFT_PASSWORD }} + PLAYWRIGHT_REDSHIFT_DATABASE: ${{ secrets.TEST_REDSHIFT_DATABASE }} + PLAYWRIGHT_METABASE_USERNAME: ${{ secrets.TEST_METABASE_USERNAME }} + PLAYWRIGHT_METABASE_PASSWORD: ${{ secrets.TEST_METABASE_PASSWORD }} + PLAYWRIGHT_METABASE_DB_SERVICE_NAME: ${{ secrets.TEST_METABASE_DB_SERVICE_NAME }} + PLAYWRIGHT_METABASE_HOST_PORT: ${{ secrets.TEST_METABASE_HOST_PORT }} + PLAYWRIGHT_SUPERSET_USERNAME: ${{ secrets.TEST_SUPERSET_USERNAME }} + PLAYWRIGHT_SUPERSET_PASSWORD: ${{ secrets.TEST_SUPERSET_PASSWORD }} + PLAYWRIGHT_SUPERSET_HOST_PORT: ${{ secrets.TEST_SUPERSET_HOST_PORT }} + PLAYWRIGHT_KAFKA_BOOTSTRAP_SERVERS: ${{ secrets.TEST_KAFKA_BOOTSTRAP_SERVERS }} + PLAYWRIGHT_KAFKA_SCHEMA_REGISTRY_URL: ${{ secrets.TEST_KAFKA_SCHEMA_REGISTRY_URL }} + PLAYWRIGHT_GLUE_ACCESS_KEY: ${{ secrets.TEST_GLUE_ACCESS_KEY }} + PLAYWRIGHT_GLUE_SECRET_KEY: ${{ secrets.TEST_GLUE_SECRET_KEY }} + PLAYWRIGHT_GLUE_AWS_REGION: ${{ secrets.TEST_GLUE_AWS_REGION }} + PLAYWRIGHT_GLUE_ENDPOINT: ${{ secrets.TEST_GLUE_ENDPOINT }} + PLAYWRIGHT_GLUE_STORAGE_SERVICE: ${{ secrets.TEST_GLUE_STORAGE_SERVICE }} + PLAYWRIGHT_MYSQL_USERNAME: ${{ secrets.TEST_MYSQL_USERNAME }} + PLAYWRIGHT_MYSQL_PASSWORD: ${{ secrets.TEST_MYSQL_PASSWORD }} + PLAYWRIGHT_MYSQL_HOST_PORT: ${{ secrets.TEST_MYSQL_HOST_PORT }} + PLAYWRIGHT_MYSQL_DATABASE_SCHEMA: ${{ secrets.TEST_MYSQL_DATABASE_SCHEMA }} + PLAYWRIGHT_POSTGRES_USERNAME: ${{ secrets.TEST_POSTGRES_USERNAME }} + PLAYWRIGHT_POSTGRES_PASSWORD: ${{ secrets.TEST_POSTGRES_PASSWORD }} + PLAYWRIGHT_POSTGRES_HOST_PORT: ${{ secrets.TEST_POSTGRES_HOST_PORT }} + PLAYWRIGHT_POSTGRES_DATABASE: ${{ secrets.TEST_POSTGRES_DATABASE }} + PLAYWRIGHT_AIRFLOW_HOST_PORT: ${{ secrets.TEST_AIRFLOW_HOST_PORT }} + PLAYWRIGHT_ML_MODEL_TRACKING_URI: ${{ secrets.TEST_ML_MODEL_TRACKING_URI }} + PLAYWRIGHT_ML_MODEL_REGISTRY_URI: ${{ secrets.TEST_ML_MODEL_REGISTRY_URI }} + PLAYWRIGHT_S3_STORAGE_ACCESS_KEY_ID: ${{ secrets.TEST_S3_STORAGE_ACCESS_KEY_ID }} + PLAYWRIGHT_S3_STORAGE_SECRET_ACCESS_KEY: ${{ secrets.TEST_S3_STORAGE_SECRET_ACCESS_KEY }} + PLAYWRIGHT_S3_STORAGE_END_POINT_URL: ${{ secrets.TEST_S3_STORAGE_END_POINT_URL }} + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + - uses: actions/upload-artifact@v4 + if: ${{ !cancelled() }} + with: + name: playwright-knowledge-graph-report + path: openmetadata-ui/src/main/resources/ui/playwright/output/playwright-report + retention-days: 5 + + - name: Clean Up + if: always() + run: | + docker compose -f docker/development/docker-compose-postgres.yml -f docker/development/docker-compose-fuseki.yml down --remove-orphans || true + docker compose -f docker/development/docker-compose-postgres.yml down --remove-orphans || true + sudo rm -rf ${PWD}/docker/development/docker-volume diff --git a/bootstrap/sql/migrations/native/1.13.0/mysql/schemaChanges.sql b/bootstrap/sql/migrations/native/1.13.0/mysql/schemaChanges.sql index ee41e5655fab..e16a031b4254 100644 --- a/bootstrap/sql/migrations/native/1.13.0/mysql/schemaChanges.sql +++ b/bootstrap/sql/migrations/native/1.13.0/mysql/schemaChanges.sql @@ -129,3 +129,82 @@ SELECT ue.id, re.id, 'user', 'role', 10 FROM user_entity ue, role_entity re WHERE ue.name = 'mcpapplicationbot' AND re.name = 'ApplicationBotImpersonationRole'; + +-- RDF distributed indexing state tables +CREATE TABLE IF NOT EXISTS rdf_index_job ( + id VARCHAR(36) NOT NULL, + status VARCHAR(32) NOT NULL, + jobConfiguration JSON NOT NULL, + totalRecords BIGINT NOT NULL DEFAULT 0, + processedRecords BIGINT NOT NULL DEFAULT 0, + successRecords BIGINT NOT NULL DEFAULT 0, + failedRecords BIGINT NOT NULL DEFAULT 0, + stats JSON, + createdBy VARCHAR(256) NOT NULL, + createdAt BIGINT NOT NULL, + startedAt BIGINT, + completedAt BIGINT, + updatedAt BIGINT NOT NULL, + errorMessage TEXT, + PRIMARY KEY (id), + INDEX idx_rdf_index_job_status (status), + INDEX idx_rdf_index_job_created (createdAt DESC) +); + +CREATE TABLE IF NOT EXISTS rdf_index_partition ( + id VARCHAR(36) NOT NULL, + jobId VARCHAR(36) NOT NULL, + entityType VARCHAR(128) NOT NULL, + partitionIndex INT NOT NULL, + rangeStart BIGINT NOT NULL, + rangeEnd BIGINT NOT NULL, + estimatedCount BIGINT NOT NULL, + workUnits BIGINT NOT NULL, + priority INT NOT NULL DEFAULT 50, + status VARCHAR(32) NOT NULL DEFAULT 'PENDING', + processingCursor BIGINT NOT NULL DEFAULT 0, + processedCount BIGINT NOT NULL DEFAULT 0, + successCount BIGINT NOT NULL DEFAULT 0, + failedCount BIGINT NOT NULL DEFAULT 0, + assignedServer VARCHAR(255), + claimedAt BIGINT, + startedAt BIGINT, + completedAt BIGINT, + lastUpdateAt BIGINT, + lastError TEXT, + retryCount INT NOT NULL DEFAULT 0, + claimableAt BIGINT NOT NULL DEFAULT 0, + PRIMARY KEY (id), + UNIQUE KEY uk_rdf_partition_job_entity_idx (jobId, entityType, partitionIndex), + INDEX idx_rdf_partition_job (jobId), + INDEX idx_rdf_partition_status_priority (status, priority DESC), + INDEX idx_rdf_partition_claimable (jobId, status, claimableAt), + INDEX idx_rdf_partition_assigned_server (jobId, assignedServer), + CONSTRAINT fk_rdf_partition_job FOREIGN KEY (jobId) REFERENCES rdf_index_job(id) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS rdf_reindex_lock ( + lockKey VARCHAR(64) NOT NULL, + jobId VARCHAR(36) NOT NULL, + serverId VARCHAR(255) NOT NULL, + acquiredAt BIGINT NOT NULL, + lastHeartbeat BIGINT NOT NULL, + expiresAt BIGINT NOT NULL, + PRIMARY KEY (lockKey) +); + +CREATE TABLE IF NOT EXISTS rdf_index_server_stats ( + id VARCHAR(36) NOT NULL, + jobId VARCHAR(36) NOT NULL, + serverId VARCHAR(256) NOT NULL, + entityType VARCHAR(128) NOT NULL, + processedRecords BIGINT DEFAULT 0, + successRecords BIGINT DEFAULT 0, + failedRecords BIGINT DEFAULT 0, + partitionsCompleted INT DEFAULT 0, + partitionsFailed INT DEFAULT 0, + lastUpdatedAt BIGINT NOT NULL, + PRIMARY KEY (id), + UNIQUE INDEX idx_rdf_index_server_stats_job_server_entity (jobId, serverId, entityType), + INDEX idx_rdf_index_server_stats_job_id (jobId) +); diff --git a/bootstrap/sql/migrations/native/1.13.0/postgres/schemaChanges.sql b/bootstrap/sql/migrations/native/1.13.0/postgres/schemaChanges.sql index 87918fb2a7d1..df571d6ef053 100644 --- a/bootstrap/sql/migrations/native/1.13.0/postgres/schemaChanges.sql +++ b/bootstrap/sql/migrations/native/1.13.0/postgres/schemaChanges.sql @@ -150,3 +150,85 @@ FROM user_entity ue, role_entity re WHERE ue.name = 'mcpapplicationbot' AND re.name = 'ApplicationBotImpersonationRole' ON CONFLICT DO NOTHING; + +-- RDF distributed indexing state tables +CREATE TABLE IF NOT EXISTS rdf_index_job ( + id VARCHAR(36) NOT NULL, + status VARCHAR(32) NOT NULL, + jobConfiguration JSONB NOT NULL, + totalRecords BIGINT NOT NULL DEFAULT 0, + processedRecords BIGINT NOT NULL DEFAULT 0, + successRecords BIGINT NOT NULL DEFAULT 0, + failedRecords BIGINT NOT NULL DEFAULT 0, + stats JSONB, + createdBy VARCHAR(256) NOT NULL, + createdAt BIGINT NOT NULL, + startedAt BIGINT, + completedAt BIGINT, + updatedAt BIGINT NOT NULL, + errorMessage TEXT, + PRIMARY KEY (id) +); + +CREATE INDEX IF NOT EXISTS idx_rdf_index_job_status ON rdf_index_job(status); +CREATE INDEX IF NOT EXISTS idx_rdf_index_job_created ON rdf_index_job(createdAt DESC); + +CREATE TABLE IF NOT EXISTS rdf_index_partition ( + id VARCHAR(36) NOT NULL, + jobId VARCHAR(36) NOT NULL, + entityType VARCHAR(128) NOT NULL, + partitionIndex INT NOT NULL, + rangeStart BIGINT NOT NULL, + rangeEnd BIGINT NOT NULL, + estimatedCount BIGINT NOT NULL, + workUnits BIGINT NOT NULL, + priority INT NOT NULL DEFAULT 50, + status VARCHAR(32) NOT NULL DEFAULT 'PENDING', + processingCursor BIGINT NOT NULL DEFAULT 0, + processedCount BIGINT NOT NULL DEFAULT 0, + successCount BIGINT NOT NULL DEFAULT 0, + failedCount BIGINT NOT NULL DEFAULT 0, + assignedServer VARCHAR(255), + claimedAt BIGINT, + startedAt BIGINT, + completedAt BIGINT, + lastUpdateAt BIGINT, + lastError TEXT, + retryCount INT NOT NULL DEFAULT 0, + claimableAt BIGINT NOT NULL DEFAULT 0, + PRIMARY KEY (id), + UNIQUE (jobId, entityType, partitionIndex), + CONSTRAINT fk_rdf_partition_job FOREIGN KEY (jobId) REFERENCES rdf_index_job(id) ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS idx_rdf_partition_job ON rdf_index_partition(jobId); +CREATE INDEX IF NOT EXISTS idx_rdf_partition_status_priority ON rdf_index_partition(status, priority DESC); +CREATE INDEX IF NOT EXISTS idx_rdf_partition_claimable ON rdf_index_partition(jobId, status, claimableAt); +CREATE INDEX IF NOT EXISTS idx_rdf_partition_assigned_server ON rdf_index_partition(jobId, assignedServer); + +CREATE TABLE IF NOT EXISTS rdf_reindex_lock ( + lockKey VARCHAR(64) NOT NULL, + jobId VARCHAR(36) NOT NULL, + serverId VARCHAR(255) NOT NULL, + acquiredAt BIGINT NOT NULL, + lastHeartbeat BIGINT NOT NULL, + expiresAt BIGINT NOT NULL, + PRIMARY KEY (lockKey) +); + +CREATE TABLE IF NOT EXISTS rdf_index_server_stats ( + id VARCHAR(36) NOT NULL, + jobId VARCHAR(36) NOT NULL, + serverId VARCHAR(256) NOT NULL, + entityType VARCHAR(128) NOT NULL, + processedRecords BIGINT DEFAULT 0, + successRecords BIGINT DEFAULT 0, + failedRecords BIGINT DEFAULT 0, + partitionsCompleted INT DEFAULT 0, + partitionsFailed INT DEFAULT 0, + lastUpdatedAt BIGINT NOT NULL, + PRIMARY KEY (id), + UNIQUE (jobId, serverId, entityType) +); + +CREATE INDEX IF NOT EXISTS idx_rdf_index_server_stats_job_id ON rdf_index_server_stats(jobId); diff --git a/docker/development/docker-compose-fuseki.yml b/docker/development/docker-compose-fuseki.yml index 14d7195a33a7..23d9daed30a6 100644 --- a/docker/development/docker-compose-fuseki.yml +++ b/docker/development/docker-compose-fuseki.yml @@ -1,27 +1,76 @@ version: "3.9" +# Compose override for RDF-enabled local stacks. +# Use together with docker-compose.yml or docker-compose-postgres.yml. services: + execute-migrate-all: + environment: + RDF_ENABLED: ${RDF_ENABLED:-true} + RDF_STORAGE_TYPE: ${RDF_STORAGE_TYPE:-FUSEKI} + RDF_ENDPOINT: ${RDF_ENDPOINT:-http://fuseki:3030/openmetadata} + RDF_REMOTE_USERNAME: ${RDF_REMOTE_USERNAME:-admin} + RDF_REMOTE_PASSWORD: ${RDF_REMOTE_PASSWORD:-admin} + RDF_BASE_URI: ${RDF_BASE_URI:-https://open-metadata.org/} + RDF_JSONLD_ENABLED: ${RDF_JSONLD_ENABLED:-true} + RDF_SPARQL_ENABLED: ${RDF_SPARQL_ENABLED:-true} + RDF_DATASET: ${RDF_DATASET:-openmetadata} + depends_on: + fuseki: + condition: service_healthy + + openmetadata-server: + environment: + RDF_ENABLED: ${RDF_ENABLED:-true} + RDF_STORAGE_TYPE: ${RDF_STORAGE_TYPE:-FUSEKI} + RDF_ENDPOINT: ${RDF_ENDPOINT:-http://fuseki:3030/openmetadata} + RDF_REMOTE_USERNAME: ${RDF_REMOTE_USERNAME:-admin} + RDF_REMOTE_PASSWORD: ${RDF_REMOTE_PASSWORD:-admin} + RDF_BASE_URI: ${RDF_BASE_URI:-https://open-metadata.org/} + RDF_JSONLD_ENABLED: ${RDF_JSONLD_ENABLED:-true} + RDF_SPARQL_ENABLED: ${RDF_SPARQL_ENABLED:-true} + RDF_DATASET: ${RDF_DATASET:-openmetadata} + depends_on: + fuseki: + condition: service_healthy + fuseki: image: stain/jena-fuseki:5.0.0 container_name: openmetadata-fuseki hostname: fuseki ports: - "3030:3030" + networks: + - local_app_net environment: - ADMIN_PASSWORD=admin - - JVM_ARGS=-Xmx4g -Xms2g + - JVM_ARGS=${FUSEKI_JVM_ARGS:--Xmx1500m -Xms256m} - FUSEKI_BASE=/fuseki volumes: - fuseki-data:/fuseki deploy: resources: limits: - memory: 4G - reservations: memory: 2G + reservations: + memory: 256m + restart: "on-failure:3" + healthcheck: + test: "curl -s -f http://localhost:3030/\\$/ping > /dev/null || exit 1" + interval: 15s + timeout: 10s + retries: 20 + start_period: 60s # Create the database directory before starting Fuseki entrypoint: /bin/sh -c "mkdir -p /fuseki/databases/openmetadata && exec /docker-entrypoint.sh /jena-fuseki/fuseki-server --update --loc=/fuseki/databases/openmetadata /openmetadata" +networks: + local_app_net: + name: ometa_network + ipam: + driver: default + config: + - subnet: "172.16.239.0/24" + volumes: fuseki-data: driver: local diff --git a/docker/run_local_docker.sh b/docker/run_local_docker.sh index cec66a41114e..0fd66ee6df3e 100755 --- a/docker/run_local_docker.sh +++ b/docker/run_local_docker.sh @@ -10,284 +10,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -RED='\033[0;31m' - cd "$(dirname "${BASH_SOURCE[0]}")" || exit -helpFunction() -{ - echo "" - echo "Usage: $0 -m mode -d database" - echo "\t-m Running mode: [ui, no-ui]. Default [ui]\n" - echo "\t-d Database: [mysql, postgresql]. Default [mysql]\n" - echo "\t-s Skip maven build: [true, false]. Default [false]\n" - echo "\t-i Include ingestion: [true, false]. Default [true]\n" - echo "\t-x Open JVM debug port on 5005: [true, false]. Default [false]\n" - echo "\t-h For usage help\n" - echo "\t-r For Cleaning DB Volumes. [true, false]. Default [true]\n" - exit 1 # Exit script after printing help -} - -while getopts "m:d:s:i:x:r:h" opt -do - case "$opt" in - m ) mode="$OPTARG" ;; - d ) database="$OPTARG" ;; - s ) skipMaven="$OPTARG" ;; - i ) includeIngestion="$OPTARG" ;; - x ) debugOM="$OPTARG" ;; - r ) cleanDbVolumes="$OPTARG" ;; - h ) helpFunction ;; - ? ) helpFunction ;; - esac -done - -mode="${mode:=ui}" -database="${database:=mysql}" -skipMaven="${skipMaven:=false}" -includeIngestion="${includeIngestion:=true}" -debugOM="${debugOM:=false}" -authorizationToken="eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg" -cleanDbVolumes="${cleanDbVolumes:=true}" - -echo "Running local docker using mode [$mode] database [$database] and skipping maven build [$skipMaven] with cleanDB as [$cleanDbVolumes] and including ingestion [$includeIngestion]" - -cd ../ - -echo "Stopping any previous Local Docker Containers" -docker compose -f docker/development/docker-compose-postgres.yml down --remove-orphans -docker compose -f docker/development/docker-compose.yml down --remove-orphans - -if [[ $skipMaven == "false" ]]; then - if [[ $mode == "no-ui" ]]; then - echo "Maven Build - Skipping Tests and UI" - mvn -DskipTests -DonlyBackend clean package -pl !openmetadata-ui - else - echo "Maven Build - Skipping Tests" - mvn -DskipTests clean package - fi -else - echo "Skipping Maven Build" -fi - -RESULT=$? -if [ $RESULT -ne 0 ]; then - echo "Failed to run Maven build!" - exit 1 -fi - -if [[ $debugOM == "true" ]]; then - export OPENMETADATA_DEBUG=true -fi - -if [[ $cleanDbVolumes == "true" ]] -then - if [[ -d "$PWD/docker/development/docker-volume/" ]] - then - rm -rf $PWD/docker/development/docker-volume - fi -fi - -if [[ $includeIngestion == "true" ]]; then - if [[ $VIRTUAL_ENV == "" ]]; - then - echo "Please Use Virtual Environment and make sure to generate Pydantic Models"; - else - echo "Generating Pydantic Models"; - make install_dev generate - fi -else - echo "Skipping Pydantic Models generation (ingestion disabled)" -fi - - -echo "Starting Local Docker Containers" - -if [[ $database == "postgresql" ]]; then - COMPOSE_FILE="docker/development/docker-compose-postgres.yml" - DB_SERVICE="postgresql" - SEARCH_SERVICE="opensearch" -elif [[ $database == "mysql" ]]; then - COMPOSE_FILE="docker/development/docker-compose.yml" - DB_SERVICE="mysql" - SEARCH_SERVICE="elasticsearch" -else - echo "Invalid database type: $database" - exit 1 -fi - -if [[ $includeIngestion == "true" ]]; then - echo "Building all services including ingestion (dependency: ${INGESTION_DEPENDENCY:-all})" - docker compose -f $COMPOSE_FILE build --build-arg INGESTION_DEPENDENCY="${INGESTION_DEPENDENCY:-all}" && docker compose -f $COMPOSE_FILE up -d -else - echo "Building services without ingestion" - docker compose -f $COMPOSE_FILE build $SEARCH_SERVICE $DB_SERVICE execute-migrate-all openmetadata-server && \ - docker compose -f $COMPOSE_FILE up -d $SEARCH_SERVICE $DB_SERVICE execute-migrate-all openmetadata-server -fi - -RESULT=$? -if [ $RESULT -ne 0 ]; then - echo "Failed to start Docker instances!" - exit 1 -fi - -until curl -s -f "http://localhost:9200/_cat/indices/openmetadata_team_search_index"; do - echo 'Checking if Elastic Search instance is up...\n' - sleep 5 -done - -if [[ $includeIngestion == "true" ]]; then - # Function to get OAuth access token for Airflow API - get_airflow_token() { - local token_response=$(curl -s -X POST 'http://localhost:8080/auth/token' \ - -H 'Content-Type: application/json' \ - -d '{"username": "admin", "password": "admin"}') - - local access_token=$(echo "$token_response" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('access_token', ''))" 2>/dev/null || echo "") - - if [ -z "$access_token" ]; then - echo "✗ Failed to get access token" >&2 - echo " Response: ${token_response}" >&2 - return 1 - fi - - echo "$access_token" - } - - # Wait for Airflow API to be ready and get initial token - echo "Waiting for Airflow API to be ready..." - until AIRFLOW_ACCESS_TOKEN=$(get_airflow_token) 2>/dev/null && [ -n "$AIRFLOW_ACCESS_TOKEN" ]; do - echo 'Checking if Airflow API is reachable...' - sleep 5 - done - echo "✓ Airflow API is ready, token obtained" - - # Check if sample_data DAG is available - echo "Checking if Sample Data DAG is available..." - until curl -s -f -H "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" "http://localhost:8080/api/v2/dags/sample_data" >/dev/null 2>&1; do - # Check for import errors - IMPORT_ERRORS=$(curl -s -H "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" "http://localhost:8080/api/v2/importErrors" 2>/dev/null) - if [ -n "$IMPORT_ERRORS" ]; then - echo "$IMPORT_ERRORS" | grep "/airflow_sample_data.py" > /dev/null 2>&1 - if [ "$?" == "0" ]; then - echo -e "${RED}Airflow found an error importing \`sample_data\` DAG" - echo "$IMPORT_ERRORS" | python3 -c "import sys, json; data=json.load(sys.stdin); [print(json.dumps(e, indent=2)) for e in data.get('import_errors', []) if e.get('filename', '').endswith('airflow_sample_data.py')]" 2>/dev/null || echo "$IMPORT_ERRORS" - exit 1 - fi - fi - echo 'Checking if Sample Data DAG is reachable...' - sleep 5 - # Refresh token if needed (tokens expire after 24h) - AIRFLOW_ACCESS_TOKEN=$(get_airflow_token) 2>/dev/null - done - echo "✓ Sample Data DAG is available" -fi - -until curl -s -f --header "Authorization: Bearer $authorizationToken" "http://localhost:8585/api/v1/tables"; do - echo 'Checking if OM Server is reachable...\n' - sleep 5 -done - -if [[ $includeIngestion == "true" ]]; then - # Function to unpause DAG using Airflow API with OAuth Bearer token - unpause_dag() { - local dag_id=$1 - echo "Unpausing DAG: ${dag_id}" - - # Get fresh token if not already set - if [ -z "$AIRFLOW_ACCESS_TOKEN" ]; then - AIRFLOW_ACCESS_TOKEN=$(get_airflow_token) - if [ -z "$AIRFLOW_ACCESS_TOKEN" ]; then - return 1 - fi - echo "✓ OAuth token obtained" - fi - - response=$(curl -s -w "\n%{http_code}" --location --request PATCH "http://localhost:8080/api/v2/dags/${dag_id}" \ - --header "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" \ - --header 'Content-Type: application/json' \ - --data-raw '{"is_paused": false}') - - http_code=$(echo "$response" | tail -n1) - body=$(echo "$response" | sed '$d') - - if [ "$http_code" = "200" ] || [ "$http_code" = "201" ]; then - echo "✓ Successfully unpaused ${dag_id}" - else - echo "✗ Failed to unpause ${dag_id} (HTTP ${http_code})" - echo " Response: ${body}" - # Token might be expired, try refreshing once - if [ "$http_code" = "401" ]; then - echo " Refreshing token and retrying..." - AIRFLOW_ACCESS_TOKEN=$(get_airflow_token) - if [ -n "$AIRFLOW_ACCESS_TOKEN" ]; then - response=$(curl -s -w "\n%{http_code}" --location --request PATCH "http://localhost:8080/api/v2/dags/${dag_id}" \ - --header "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" \ - --header 'Content-Type: application/json' \ - --data-raw '{"is_paused": false}') - http_code=$(echo "$response" | tail -n1) - if [ "$http_code" = "200" ] || [ "$http_code" = "201" ]; then - echo "✓ Successfully unpaused ${dag_id} after retry" - fi - fi - fi - fi - } - - unpause_dag "sample_data" - unpause_dag "extended_sample_data" - - # Trigger sample_data DAG to run - echo "Triggering sample_data DAG..." - LOGICAL_DATE=$(date -u +"%Y-%m-%dT%H:%M:%SZ") - response=$(curl -s -w "\n%{http_code}" -X POST "http://localhost:8080/api/v2/dags/sample_data/dagRuns" \ - --header "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" \ - --header 'Content-Type: application/json' \ - --data-raw "{\"logical_date\": \"$LOGICAL_DATE\"}") - - http_code=$(echo "$response" | tail -n1) - if [ "$http_code" = "200" ] || [ "$http_code" = "201" ]; then - echo "✓ Successfully triggered sample_data DAG" - else - echo "⚠ Could not trigger sample_data DAG (HTTP ${http_code})" - echo " Response: $(echo "$response" | sed '$d')" - echo " Note: DAG may run automatically on schedule" - fi - - echo 'Validate sample data DAG...' - sleep 5 - # This validates the sample data DAG flow - make install - - # Run validation with timeout to avoid hanging indefinitely - echo "Running DAG validation (this may take a few minutes)..." - timeout 300 python docker/validate_compose.py || { - exit_code=$? - if [ $exit_code -eq 124 ]; then - echo "⚠ Warning: DAG validation timed out after 5 minutes" - echo " The DAG may still be running. Check Airflow UI at http://localhost:8080" - else - echo "⚠ Warning: DAG validation failed with exit code $exit_code" - fi - echo " Continuing with remaining setup..." - } - - sleep 5 - unpause_dag "sample_usage" - sleep 5 - unpause_dag "index_metadata" - sleep 2 - unpause_dag "sample_lineage" -else - echo "Skipping Airflow DAG setup (ingestion disabled)" -fi - -echo "✔running reindexing" -# Trigger ElasticSearch ReIndexing from UI -curl --location --request POST 'http://localhost:8585/api/v1/apps/trigger/SearchIndexingApplication' \ ---header 'Authorization: Bearer eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg' - -sleep 60 # Sleep for 60 seconds to make sure the elasticsearch reindexing from UI finishes -tput setaf 2 -echo "✔ OpenMetadata is up and running" +source ./run_local_docker_common.sh +run_local_docker_main "$@" diff --git a/docker/run_local_docker_common.sh b/docker/run_local_docker_common.sh new file mode 100755 index 000000000000..6de60d23012b --- /dev/null +++ b/docker/run_local_docker_common.sh @@ -0,0 +1,614 @@ +#!/bin/bash +# Copyright 2021 Collate +# 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. + +RED='\033[0;31m' +RUN_LOCAL_DOCKER_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +authorizationToken="eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg" + +helpFunction() { + echo "" + echo "Usage: $0 -m mode -d database" + echo "\t-m Running mode: [ui, no-ui]. Default [ui]\n" + echo "\t-d Database: [mysql, postgresql]. Default [mysql]\n" + echo "\t-s Skip maven build: [true, false]. Default [false]\n" + echo "\t-i Include ingestion: [true, false]. Default [true]\n" + echo "\t-x Open JVM debug port on 5005: [true, false]. Default [false]\n" + echo "\t-h For usage help\n" + echo "\t-r For Cleaning DB Volumes. [true, false]. Default [true]\n" + exit 1 +} + +current_time_ms() { + python3 -c "import time; print(int(time.time() * 1000))" +} + +run_with_timeout() { + local secs=$1 + shift + if command -v timeout &>/dev/null; then + timeout "$secs" "$@" + elif command -v gtimeout &>/dev/null; then + gtimeout "$secs" "$@" + else + "$@" + fi +} + +parse_app_run_status_line() { + local payload=$1 + local payload_shape=$2 + local threshold_ms=$3 + local tolerance_ms=$4 + + APP_RUN_BODY="$payload" python3 - "$payload_shape" "$threshold_ms" "$tolerance_ms" <<'PY' +import json +import os +import sys + +payload_shape = sys.argv[1] +threshold = int(sys.argv[2]) +tolerance = int(sys.argv[3]) + +try: + payload = json.loads(os.environ["APP_RUN_BODY"]) +except json.JSONDecodeError: + print("invalid") + sys.exit(0) + +if payload_shape == "latest": + record = payload if isinstance(payload, dict) else None +else: + records = payload.get("data") if isinstance(payload, dict) else None + record = records[0] if records else None + +if not isinstance(record, dict) or not record: + print("missing") + sys.exit(0) + +timestamp = int(record.get("timestamp") or 0) +start_time = int(record.get("startTime") or 0) +status = str(record.get("status") or "").lower() +execution_time = record.get("executionTime") + +marker = start_time if start_time > 0 else timestamp +success_statuses = {"completed", "success"} +failure_statuses = {"activeerror", "failed", "stopped"} +active_statuses = {"running", "started", "pending", "active", "stopinprogress"} +has_execution_time = execution_time not in (None, "", 0, "0") + +if marker + tolerance < threshold: + print(f"stale:{status}:{marker}") +elif status in success_statuses: + print(f"success:{status}:{marker}") +elif status in failure_statuses or (has_execution_time and status not in active_statuses): + print(f"failure:{status}:{marker}") +elif status in active_statuses: + print(f"active:{status}:{marker}") +else: + print(f"seen:{status}:{marker}") +PY +} + +wait_for_app_availability() { + local app_name=$1 + local timeout_seconds=${2:-120} + local deadline=$((SECONDS + timeout_seconds)) + + while [ $SECONDS -lt $deadline ]; do + local http_code + http_code=$(curl -s -o /dev/null -w "%{http_code}" \ + --header "Authorization: Bearer $authorizationToken" \ + "http://localhost:8585/api/v1/apps/name/${app_name}") + + if [ "$http_code" = "200" ]; then + echo "✓ App ${app_name} is available" + return 0 + fi + + echo "Waiting for app ${app_name} to become available..." + sleep 5 + done + + echo "✗ App ${app_name} did not become available within ${timeout_seconds}s" + return 1 +} + +ensure_app_installed() { + local app_name=$1 + local response + local http_code + local body + + response=$(curl -s -w "\n%{http_code}" \ + --header "Authorization: Bearer $authorizationToken" \ + "http://localhost:8585/api/v1/apps/name/${app_name}") + http_code=$(echo "$response" | tail -n1) + body=$(echo "$response" | sed '$d') + + if [ "$http_code" = "200" ]; then + echo "✓ App ${app_name} is already installed" + return 0 + fi + + if [ "$http_code" != "404" ]; then + echo "✗ Failed to inspect app ${app_name} (HTTP ${http_code})" + echo " Response: ${body}" + return 1 + fi + + echo "App ${app_name} is not installed. Installing it from the marketplace definition..." + local marketplace_response + marketplace_response=$(curl -s -f \ + --header "Authorization: Bearer $authorizationToken" \ + "http://localhost:8585/api/v1/apps/marketplace/name/${app_name}") || { + echo "✗ Could not fetch marketplace definition for ${app_name}" + return 1 + } + + local create_payload + create_payload=$(MARKETPLACE_RESPONSE="$marketplace_response" python3 - <<'PY' +import json +import os + +definition = json.loads(os.environ["MARKETPLACE_RESPONSE"]) +payload = { + "name": definition["name"], + "displayName": definition.get("displayName"), + "description": definition.get("description"), + "appConfiguration": definition.get("appConfiguration", {}), + "appSchedule": {"scheduleTimeline": "None"}, + "supportsInterrupt": definition.get("supportsInterrupt", False), +} +print(json.dumps(payload)) +PY +) + + response=$(curl -s -w "\n%{http_code}" --location --request POST 'http://localhost:8585/api/v1/apps' \ + --header "Authorization: Bearer $authorizationToken" \ + --header 'Content-Type: application/json' \ + --data-raw "$create_payload") + http_code=$(echo "$response" | tail -n1) + body=$(echo "$response" | sed '$d') + + if [ "$http_code" = "200" ] || [ "$http_code" = "201" ] || [ "$http_code" = "409" ]; then + wait_for_app_availability "$app_name" 120 + return $? + fi + + echo "✗ Failed to install app ${app_name} (HTTP ${http_code})" + echo " Response: ${body}" + return 1 +} + +wait_for_app_run_completion() { + local app_name=$1 + local trigger_timestamp_ms=$2 + local timeout_seconds=${3:-${APP_RUN_WAIT_TIMEOUT_SECONDS:-300}} + local freshness_tolerance_ms=${APP_RUN_FRESHNESS_TOLERANCE_MS:-5000} + local deadline=$((SECONDS + timeout_seconds)) + + while [ $SECONDS -lt $deadline ]; do + local latest_run_response + local fallback_status_response + local http_code + local fallback_http_code + local body + local fallback_body + local status_line + local fallback_status_line + + latest_run_response=$(curl -s -w "\n%{http_code}" \ + --header "Authorization: Bearer $authorizationToken" \ + "http://localhost:8585/api/v1/apps/name/${app_name}/runs/latest") + http_code=$(printf "%s" "$latest_run_response" | tail -n1) + body=$(printf "%s" "$latest_run_response" | sed '$d') + + case "$http_code" in + 200) + status_line=$(parse_app_run_status_line "$body" "latest" "$trigger_timestamp_ms" "$freshness_tolerance_ms") + ;; + 204) + status_line="missing" + ;; + *) + status_line="endpoint_error:${http_code}" + ;; + esac + + if [[ "$status_line" == stale:* || "$status_line" == "missing" || "$status_line" == "invalid" || "$status_line" == endpoint_error:* ]]; then + fallback_status_response=$(curl -s -w "\n%{http_code}" \ + --header "Authorization: Bearer $authorizationToken" \ + "http://localhost:8585/api/v1/apps/name/${app_name}/status?offset=0&limit=1") + fallback_http_code=$(printf "%s" "$fallback_status_response" | tail -n1) + fallback_body=$(printf "%s" "$fallback_status_response" | sed '$d') + + case "$fallback_http_code" in + 200) + fallback_status_line=$(parse_app_run_status_line "$fallback_body" "list" "$trigger_timestamp_ms" "$freshness_tolerance_ms") + ;; + 204) + fallback_status_line="missing" + ;; + *) + fallback_status_line="endpoint_error:${fallback_http_code}" + ;; + esac + + if [[ "$fallback_status_line" != stale:* && "$fallback_status_line" != "missing" && "$fallback_status_line" != "invalid" && "$fallback_status_line" != endpoint_error:* ]]; then + status_line="$fallback_status_line" + body="$fallback_body" + elif [[ "$status_line" == endpoint_error:* && "$fallback_status_line" == endpoint_error:* ]]; then + echo "✗ Failed to read run status for ${app_name} from both app run endpoints" + echo " runs/latest response: ${body}" + echo " status response: ${fallback_body}" + return 1 + elif [[ "$status_line" == "missing" || "$status_line" == "invalid" || "$status_line" == endpoint_error:* ]]; then + status_line="$fallback_status_line" + body="$fallback_body" + fi + fi + + case "$status_line" in + success:completed:*|success:success:*) + echo "✓ ${app_name} completed successfully" + return 0 + ;; + failure:*) + echo "✗ ${app_name} finished with status ${status_line#failure:}" + echo " Response: ${body}" + return 1 + ;; + active:*) + echo "Waiting for ${app_name} to finish (${status_line#active:})..." + ;; + seen:*) + echo "Waiting for ${app_name}; latest status was ${status_line#seen:}" + ;; + stale:*|missing|invalid) + echo "Waiting for a fresh run record for ${app_name}..." + ;; + *) + echo "Waiting for ${app_name}; latest status was ${status_line}" + ;; + esac + + sleep 5 + done + + echo "✗ Timed out waiting for ${app_name} to finish within ${timeout_seconds}s" + return 1 +} + +trigger_app_and_wait() { + local app_name=$1 + local payload=${2:-} + local timeout_seconds=${3:-${APP_RUN_WAIT_TIMEOUT_SECONDS:-300}} + local trigger_timestamp_ms + local response + local http_code + local body + + trigger_timestamp_ms=$(current_time_ms) + if [ -n "$payload" ]; then + response=$(curl -s -w "\n%{http_code}" --location --request POST "http://localhost:8585/api/v1/apps/trigger/${app_name}" \ + --header "Authorization: Bearer $authorizationToken" \ + --header 'Content-Type: application/json' \ + --data-raw "$payload") + else + response=$(curl -s -w "\n%{http_code}" --location --request POST "http://localhost:8585/api/v1/apps/trigger/${app_name}" \ + --header "Authorization: Bearer $authorizationToken") + fi + + http_code=$(echo "$response" | tail -n1) + body=$(echo "$response" | sed '$d') + + if [ "$http_code" != "200" ] && [ "$http_code" != "201" ] && [ "$http_code" != "202" ]; then + echo "✗ Failed to trigger ${app_name} (HTTP ${http_code})" + echo " Response: ${body}" + return 1 + fi + + wait_for_app_run_completion "$app_name" "$trigger_timestamp_ms" "$timeout_seconds" +} + +run_local_docker_main() { + local mode database skipMaven includeIngestion debugOM cleanDbVolumes + local COMPOSE_FILE DB_SERVICE SEARCH_SERVICE + local response http_code body + local validation_timeout_seconds sample_data_validation_failed + local AIRFLOW_ACCESS_TOKEN IMPORT_ERRORS LOGICAL_DATE + local extra_compose_file + local -a COMPOSE_ARGS additional_up_services + + OPTIND=1 + while getopts "m:d:s:i:x:r:h" opt; do + case "$opt" in + m ) mode="$OPTARG" ;; + d ) database="$OPTARG" ;; + s ) skipMaven="$OPTARG" ;; + i ) includeIngestion="$OPTARG" ;; + x ) debugOM="$OPTARG" ;; + r ) cleanDbVolumes="$OPTARG" ;; + h ) helpFunction ;; + ? ) helpFunction ;; + esac + done + + mode="${mode:=ui}" + database="${database:=mysql}" + skipMaven="${skipMaven:=false}" + includeIngestion="${includeIngestion:=true}" + debugOM="${debugOM:=false}" + cleanDbVolumes="${cleanDbVolumes:=true}" + export APP_RUN_WAIT_TIMEOUT_SECONDS="${APP_RUN_WAIT_TIMEOUT_SECONDS:-300}" + + echo "Running local docker using mode [$mode] database [$database] and skipping maven build [$skipMaven] with cleanDB as [$cleanDbVolumes] and including ingestion [$includeIngestion]" + + cd "$RUN_LOCAL_DOCKER_DIR/.." || exit 1 + + echo "Stopping any previous Local Docker Containers" + docker compose -f docker/development/docker-compose-postgres.yml down --remove-orphans + docker compose -f docker/development/docker-compose.yml down --remove-orphans + if [[ -n "${OM_EXTRA_COMPOSE_FILES:-}" ]]; then + for extra_compose_file in $OM_EXTRA_COMPOSE_FILES; do + docker compose -f docker/development/docker-compose-postgres.yml -f "$extra_compose_file" down --remove-orphans + docker compose -f docker/development/docker-compose.yml -f "$extra_compose_file" down --remove-orphans + done + fi + + if [[ $skipMaven == "false" ]]; then + if [[ $mode == "no-ui" ]]; then + echo "Maven Build - Skipping Tests and UI" + mvn -DskipTests -DonlyBackend clean package -pl !openmetadata-ui + else + echo "Maven Build - Skipping Tests" + mvn -DskipTests clean package + fi + else + echo "Skipping Maven Build" + fi + + if [ $? -ne 0 ]; then + echo "Failed to run Maven build!" + exit 1 + fi + + if [[ $debugOM == "true" ]]; then + export OPENMETADATA_DEBUG=true + fi + + if [[ $cleanDbVolumes == "true" ]]; then + if [[ -d "$PWD/docker/development/docker-volume/" ]]; then + if ! rm -rf "$PWD/docker/development/docker-volume"; then + if command -v sudo >/dev/null 2>&1 && sudo -n true >/dev/null 2>&1; then + sudo rm -rf "$PWD/docker/development/docker-volume" + else + echo "Warning: failed to remove $PWD/docker/development/docker-volume; continuing may reuse stale database state" + fi + fi + fi + fi + + if [[ $includeIngestion == "true" ]]; then + if [[ $VIRTUAL_ENV == "" ]]; then + echo "Please Use Virtual Environment and make sure to generate Pydantic Models" + else + echo "Generating Pydantic Models" + make install_dev generate + fi + else + echo "Skipping Pydantic Models generation (ingestion disabled)" + fi + + echo "Starting Local Docker Containers" + + if [[ $database == "postgresql" ]]; then + COMPOSE_FILE="docker/development/docker-compose-postgres.yml" + DB_SERVICE="postgresql" + SEARCH_SERVICE="opensearch" + elif [[ $database == "mysql" ]]; then + COMPOSE_FILE="docker/development/docker-compose.yml" + DB_SERVICE="mysql" + SEARCH_SERVICE="elasticsearch" + else + echo "Invalid database type: $database" + exit 1 + fi + + COMPOSE_ARGS=(-f "$COMPOSE_FILE") + if [[ -n "${OM_EXTRA_COMPOSE_FILES:-}" ]]; then + for extra_compose_file in $OM_EXTRA_COMPOSE_FILES; do + COMPOSE_ARGS+=(-f "$extra_compose_file") + done + fi + + if [[ $includeIngestion == "true" ]]; then + echo "Building all services including ingestion (dependency: ${INGESTION_DEPENDENCY:-all})" + docker compose "${COMPOSE_ARGS[@]}" build --build-arg INGESTION_DEPENDENCY="${INGESTION_DEPENDENCY:-all}" && \ + docker compose "${COMPOSE_ARGS[@]}" up -d + else + echo "Building services without ingestion" + docker compose "${COMPOSE_ARGS[@]}" build $SEARCH_SERVICE $DB_SERVICE execute-migrate-all openmetadata-server || exit 1 + if [[ -n "${OM_ADDITIONAL_UP_SERVICES:-}" ]]; then + for extra_compose_file in $OM_ADDITIONAL_UP_SERVICES; do + additional_up_services+=("$extra_compose_file") + done + fi + docker compose "${COMPOSE_ARGS[@]}" up -d "${additional_up_services[@]}" $SEARCH_SERVICE $DB_SERVICE execute-migrate-all openmetadata-server + fi + + if [ $? -ne 0 ]; then + echo "Failed to start Docker instances!" + exit 1 + fi + + until curl -s -f "http://localhost:9200/_cat/indices/openmetadata_team_search_index"; do + echo 'Checking if Elastic Search instance is up...\n' + sleep 5 + done + + if [[ $includeIngestion == "true" ]]; then + get_airflow_token() { + local token_response + local access_token + + token_response=$(curl -s -X POST 'http://localhost:8080/auth/token' \ + -H 'Content-Type: application/json' \ + -d '{"username": "admin", "password": "admin"}') + + access_token=$(echo "$token_response" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('access_token', ''))" 2>/dev/null || echo "") + + if [ -z "$access_token" ]; then + echo "✗ Failed to get access token" >&2 + echo " Response: ${token_response}" >&2 + return 1 + fi + + echo "$access_token" + } + + echo "Waiting for Airflow API to be ready..." + until AIRFLOW_ACCESS_TOKEN=$(get_airflow_token) 2>/dev/null && [ -n "$AIRFLOW_ACCESS_TOKEN" ]; do + echo 'Checking if Airflow API is reachable...' + sleep 5 + done + echo "✓ Airflow API is ready, token obtained" + + echo "Checking if Sample Data DAG is available..." + until curl -s -f -H "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" "http://localhost:8080/api/v2/dags/sample_data" >/dev/null 2>&1; do + IMPORT_ERRORS=$(curl -s -H "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" "http://localhost:8080/api/v2/importErrors" 2>/dev/null) + if [ -n "$IMPORT_ERRORS" ]; then + echo "$IMPORT_ERRORS" | grep "/airflow_sample_data.py" > /dev/null 2>&1 + if [ "$?" == "0" ]; then + echo -e "${RED}Airflow found an error importing \`sample_data\` DAG" + echo "$IMPORT_ERRORS" | python3 -c "import sys, json; data=json.load(sys.stdin); [print(json.dumps(e, indent=2)) for e in data.get('import_errors', []) if e.get('filename', '').endswith('airflow_sample_data.py')]" 2>/dev/null || echo "$IMPORT_ERRORS" + exit 1 + fi + fi + echo 'Checking if Sample Data DAG is reachable...' + sleep 5 + AIRFLOW_ACCESS_TOKEN=$(get_airflow_token) 2>/dev/null + done + echo "✓ Sample Data DAG is available" + fi + + until curl -s -f --header "Authorization: Bearer $authorizationToken" "http://localhost:8585/api/v1/tables"; do + echo 'Checking if OM Server is reachable...\n' + sleep 5 + done + + if [[ $includeIngestion == "true" ]]; then + unpause_dag() { + local dag_id=$1 + + echo "Unpausing DAG: ${dag_id}" + if [ -z "$AIRFLOW_ACCESS_TOKEN" ]; then + AIRFLOW_ACCESS_TOKEN=$(get_airflow_token) + if [ -z "$AIRFLOW_ACCESS_TOKEN" ]; then + return 1 + fi + echo "✓ OAuth token obtained" + fi + + response=$(curl -s -w "\n%{http_code}" --location --request PATCH "http://localhost:8080/api/v2/dags/${dag_id}" \ + --header "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" \ + --header 'Content-Type: application/json' \ + --data-raw '{"is_paused": false}') + + http_code=$(echo "$response" | tail -n1) + body=$(echo "$response" | sed '$d') + + if [ "$http_code" = "200" ] || [ "$http_code" = "201" ]; then + echo "✓ Successfully unpaused ${dag_id}" + else + echo "✗ Failed to unpause ${dag_id} (HTTP ${http_code})" + echo " Response: ${body}" + if [ "$http_code" = "401" ]; then + echo " Refreshing token and retrying..." + AIRFLOW_ACCESS_TOKEN=$(get_airflow_token) + if [ -n "$AIRFLOW_ACCESS_TOKEN" ]; then + response=$(curl -s -w "\n%{http_code}" --location --request PATCH "http://localhost:8080/api/v2/dags/${dag_id}" \ + --header "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" \ + --header 'Content-Type: application/json' \ + --data-raw '{"is_paused": false}') + http_code=$(echo "$response" | tail -n1) + if [ "$http_code" = "200" ] || [ "$http_code" = "201" ]; then + echo "✓ Successfully unpaused ${dag_id} after retry" + fi + fi + fi + fi + } + + unpause_dag "sample_data" + unpause_dag "extended_sample_data" + + echo "Triggering sample_data DAG..." + LOGICAL_DATE=$(date -u +"%Y-%m-%dT%H:%M:%SZ") + response=$(curl -s -w "\n%{http_code}" -X POST "http://localhost:8080/api/v2/dags/sample_data/dagRuns" \ + --header "Authorization: Bearer $AIRFLOW_ACCESS_TOKEN" \ + --header 'Content-Type: application/json' \ + --data-raw "{\"logical_date\": \"$LOGICAL_DATE\"}") + + http_code=$(echo "$response" | tail -n1) + if [ "$http_code" = "200" ] || [ "$http_code" = "201" ]; then + echo "✓ Successfully triggered sample_data DAG" + else + echo "⚠ Could not trigger sample_data DAG (HTTP ${http_code})" + echo " Response: $(echo "$response" | sed '$d')" + echo " Note: DAG may run automatically on schedule" + fi + + echo 'Validate sample data DAG...' + sleep 5 + make install + + echo "Running DAG validation (this may take a few minutes)..." + sample_data_validation_failed=false + validation_timeout_seconds="${VALIDATION_TIMEOUT_SECONDS:-300}" + + run_with_timeout "$validation_timeout_seconds" python docker/validate_compose.py || { + local exit_code=$? + sample_data_validation_failed=true + if [ $exit_code -eq 124 ]; then + echo "⚠ Warning: DAG validation timed out after ${validation_timeout_seconds} seconds" + echo " The DAG may still be running. Check Airflow UI at http://localhost:8080" + else + echo "⚠ Warning: DAG validation failed with exit code $exit_code" + fi + echo " Continuing with remaining setup..." + } + + if [[ "${STRICT_DAG_VALIDATION:-false}" == "true" && "$sample_data_validation_failed" == "true" ]]; then + echo "✗ Startup requires sample data ingestion to complete before continuing." + exit 1 + fi + + sleep 5 + unpause_dag "sample_usage" + sleep 5 + unpause_dag "index_metadata" + sleep 2 + unpause_dag "sample_lineage" + else + echo "Skipping Airflow DAG setup (ingestion disabled)" + fi + + echo "✔running reindexing" + ensure_app_installed "SearchIndexingApplication" + if ! trigger_app_and_wait "SearchIndexingApplication" "" "$APP_RUN_WAIT_TIMEOUT_SECONDS"; then + exit 1 + fi + + tput setaf 2 + echo "✔ OpenMetadata is up and running" +} diff --git a/docker/run_local_docker_rdf.sh b/docker/run_local_docker_rdf.sh index 9f7aef4e14d8..f411e16083aa 100755 --- a/docker/run_local_docker_rdf.sh +++ b/docker/run_local_docker_rdf.sh @@ -12,214 +12,95 @@ cd "$(dirname "${BASH_SOURCE[0]}")" || exit -helpFunction() -{ +helpFunction() { echo "" - echo "Usage: $0 -m mode -d database" - echo "\t-m Running mode: [ui, no-ui]. Default [ui]" - echo "\t-d Database: [mysql, postgresql]. Default [mysql]" - echo "\t-s Skip maven build: [true, false]. Default [false]" - echo "\t-x Open JVM debug port on 5005: [true, false]. Default [false]" - echo "\t-h For usage help" - echo "\t-r For Cleaning DB Volumes. [true, false]. Default [true]" + echo "Usage: $0 [run_local_docker.sh args]" echo "\t-f Start Fuseki for RDF support: [true, false]. Default [true]" - exit 1 # Exit script after printing help + echo "\t-h For usage help" + exit 1 } -while getopts "m:d:s:x:r:f:h" opt -do - case "$opt" in - m ) mode="$OPTARG" ;; - d ) database="$OPTARG" ;; - s ) skipMaven="$OPTARG" ;; - x ) debugOM="$OPTARG" ;; - r ) cleanDbVolumes="$OPTARG" ;; - f ) startFuseki="$OPTARG" ;; - h ) helpFunction ;; - ? ) helpFunction ;; - esac +startFuseki=true +filtered_args=() + +while [[ $# -gt 0 ]]; do + case "$1" in + -f) + if [[ $# -lt 2 ]]; then + helpFunction + fi + startFuseki="$2" + shift 2 + ;; + -h) + helpFunction + ;; + *) + filtered_args+=("$1") + shift + ;; + esac done -mode="${mode:=ui}" -database="${database:=mysql}" -skipMaven="${skipMaven:=false}" -debugOM="${debugOM:=false}" -authorizationToken="eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg" -cleanDbVolumes="${cleanDbVolumes:=true}" -startFuseki="${startFuseki:=true}" - -echo "Running local docker using mode [$mode] database [$database] and skipping maven build [$skipMaven] with cleanDB as [$cleanDbVolumes] and Fuseki [$startFuseki]" - -cd ../ - -echo "Stopping any previous Local Docker Containers" -docker compose -f docker/development/docker-compose-postgres.yml down --remove-orphans -docker compose -f docker/development/docker-compose.yml down --remove-orphans -docker compose -f docker/development/docker-compose-fuseki.yml down --remove-orphans - -if [[ $skipMaven == "false" ]]; then - if [[ $mode == "no-ui" ]]; then - echo "Maven Build - Skipping Tests and UI" - mvn -DskipTests -DonlyBackend clean package -pl !openmetadata-ui - else - echo "Maven Build - Skipping Tests" - mvn -DskipTests clean package - fi -else - echo "Skipping Maven Build" -fi - -RESULT=$? -if [ $RESULT -ne 0 ]; then - echo "Failed to run Maven build!" - exit 1 -fi - -if [[ $debugOM == "true" ]]; then - export OPENMETADATA_DEBUG=true -fi - -if [[ $cleanDbVolumes == "true" ]] -then - if [[ -d "$PWD/docker/development/docker-volume/" ]] - then - rm -rf $PWD/docker/development/docker-volume - fi -fi - -if [[ $VIRTUAL_ENV == "" ]]; -then - echo "Please Use Virtual Environment and make sure to generate Pydantic Models"; -else - echo "Generating Pydantic Models"; - make install_dev generate -fi - -# Start Fuseki if requested if [[ $startFuseki == "true" ]]; then - echo "Starting Apache Jena Fuseki for RDF support" - docker compose -f docker/development/docker-compose-fuseki.yml up -d - - # Wait for Fuseki to be ready - until curl -s -f "http://localhost:3030/$/ping" > /dev/null 2>&1; do - echo 'Waiting for Fuseki to start...' - sleep 5 - done - echo "✔ Fuseki is ready" - - # Set RDF environment variables - export RDF_ENABLED=true - export RDF_STORAGE_TYPE=FUSEKI - export RDF_BASE_URI="https://open-metadata.org/" - export RDF_ENDPOINT="http://localhost:3030/openmetadata" - export RDF_REMOTE_USERNAME="admin" - export RDF_REMOTE_PASSWORD="admin" - export RDF_DATASET="openmetadata" -fi - -echo "Starting Local Docker Containers" -echo "Using ingestion dependency: ${INGESTION_DEPENDENCY:-all}" - -if [[ $database == "postgresql" ]]; then - docker compose -f docker/development/docker-compose-postgres.yml build --build-arg INGESTION_DEPENDENCY="${INGESTION_DEPENDENCY:-all}" && docker compose -f docker/development/docker-compose-postgres.yml up -d -elif [[ $database == "mysql" ]]; then - docker compose -f docker/development/docker-compose.yml build --build-arg INGESTION_DEPENDENCY="${INGESTION_DEPENDENCY:-all}" && docker compose -f docker/development/docker-compose.yml up -d + export RDF_ENABLED=true + export RDF_AUTO_REINDEX=true + export RDF_STORAGE_TYPE="${RDF_STORAGE_TYPE:-FUSEKI}" + export RDF_ENDPOINT="${RDF_ENDPOINT:-http://fuseki:3030/openmetadata}" + export RDF_REMOTE_USERNAME="${RDF_REMOTE_USERNAME:-admin}" + export RDF_REMOTE_PASSWORD="${RDF_REMOTE_PASSWORD:-admin}" + export RDF_BASE_URI="${RDF_BASE_URI:-https://open-metadata.org/}" + export RDF_DATASET="${RDF_DATASET:-openmetadata}" + # RDF listeners slow down sample-data ingestion enough that the default 5-minute + # validation window is too aggressive for CI. + export VALIDATE_COMPOSE_MAX_RETRIES="${VALIDATE_COMPOSE_MAX_RETRIES:-60}" + export VALIDATE_COMPOSE_DAG_RUN_RETRIES="${VALIDATE_COMPOSE_DAG_RUN_RETRIES:-120}" + export VALIDATE_COMPOSE_RETRY_INTERVAL_SECONDS="${VALIDATE_COMPOSE_RETRY_INTERVAL_SECONDS:-10}" + export VALIDATE_COMPOSE_DAG_RUN_POLL_SECONDS="${VALIDATE_COMPOSE_DAG_RUN_POLL_SECONDS:-5}" + export VALIDATION_TIMEOUT_SECONDS="${VALIDATION_TIMEOUT_SECONDS:-900}" + export APP_RUN_WAIT_TIMEOUT_SECONDS="${APP_RUN_WAIT_TIMEOUT_SECONDS:-900}" + export STRICT_DAG_VALIDATION=true + export OM_EXTRA_COMPOSE_FILES="docker/development/docker-compose-fuseki.yml" + export OM_ADDITIONAL_UP_SERVICES="fuseki" else - echo "Invalid database type: $database" - exit 1 + export RDF_ENABLED=false + export RDF_AUTO_REINDEX=false + unset RDF_STORAGE_TYPE + unset RDF_ENDPOINT + unset RDF_REMOTE_USERNAME + unset RDF_REMOTE_PASSWORD + unset RDF_BASE_URI + unset RDF_DATASET + unset VALIDATE_COMPOSE_MAX_RETRIES + unset VALIDATE_COMPOSE_DAG_RUN_RETRIES + unset VALIDATE_COMPOSE_RETRY_INTERVAL_SECONDS + unset VALIDATE_COMPOSE_DAG_RUN_POLL_SECONDS + export VALIDATION_TIMEOUT_SECONDS="${VALIDATION_TIMEOUT_SECONDS:-300}" + export APP_RUN_WAIT_TIMEOUT_SECONDS="${APP_RUN_WAIT_TIMEOUT_SECONDS:-300}" + unset STRICT_DAG_VALIDATION + unset OM_EXTRA_COMPOSE_FILES + unset OM_ADDITIONAL_UP_SERVICES fi -RESULT=$? -if [ $RESULT -ne 0 ]; then - echo "Failed to start Docker instances!" - exit 1 -fi - -until curl -s -f "http://localhost:9200/_cat/indices/openmetadata_team_search_index"; do - echo 'Checking if Elastic Search instance is up...' - sleep 5 -done - -until curl -s -f --header 'Authorization: Basic YWRtaW46YWRtaW4=' "http://localhost:8080/api/v1/dags/sample_data"; do - echo 'Checking if Sample Data DAG is reachable...' - sleep 5 -done - -until curl -s -f --header "Authorization: Bearer $authorizationToken" "http://localhost:8585/api/v1/tables"; do - echo 'Checking if OM Server is reachable...' - sleep 5 -done - -curl --location --request PATCH 'localhost:8080/api/v1/dags/sample_data' \ - --header 'Authorization: Basic YWRtaW46YWRtaW4=' \ - --header 'Content-Type: application/json' \ - --data-raw '{ - "is_paused": false - }' - -curl --location --request PATCH 'localhost:8080/api/v1/dags/extended_sample_data' \ - --header 'Authorization: Basic YWRtaW46YWRtaW4=' \ - --header 'Content-Type: application/json' \ - --data-raw '{ - "is_paused": false - }' - -echo 'Validate sample data DAG...' -sleep 5 -# This validates the sample data DAG flow -make install -python docker/validate_compose.py - -sleep 5 -curl --location --request PATCH 'localhost:8080/api/v1/dags/sample_usage' \ - --header 'Authorization: Basic YWRtaW46YWRtaW4=' \ - --header 'Content-Type: application/json' \ - --data-raw '{ - "is_paused": false - }' -sleep 5 -curl --location --request PATCH 'localhost:8080/api/v1/dags/index_metadata' \ - --header 'Authorization: Basic YWRtaW46YWRtaW4=' \ - --header 'Content-Type: application/json' \ - --data-raw '{ - "is_paused": false - }' -sleep 2 -curl --location --request PATCH 'localhost:8080/api/v1/dags/sample_lineage' \ - --header 'Authorization: Basic YWRtaW46YWRtaW4=' \ - --header 'Content-Type: application/json' \ - --data-raw '{ - "is_paused": false - }' - -echo "✔running reindexing" -# Trigger ElasticSearch ReIndexing from UI -curl --location --request POST 'http://localhost:8585/api/v1/apps/trigger/SearchIndexingApplication' \ ---header 'Authorization: Bearer eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg' +source ./run_local_docker_common.sh -sleep 60 # Sleep for 60 seconds to make sure the elasticsearch reindexing from UI finishes +run_local_docker_main "${filtered_args[@]}" -# If RDF is enabled, trigger RDF indexing -if [[ $startFuseki == "true" ]]; then - echo "✔running RDF reindexing" - # Trigger RDF ReIndexing from UI - curl --location --request POST 'http://localhost:8585/api/v1/apps/trigger/RdfIndexApp' \ - --header 'Authorization: Bearer eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg' \ - --header 'Content-Type: application/json' \ - --data-raw '{ - "entities": ["all"], - "recreateIndex": true, - "batchSize": 100 - }' - - sleep 30 # Wait for RDF indexing to complete +if [[ $startFuseki != "true" || $RDF_AUTO_REINDEX != "true" ]]; then + exit 0 fi +ensure_app_installed "RdfIndexApp" +echo "✔running RDF reindexing" +trigger_app_and_wait "RdfIndexApp" '{ + "entities": [], + "recreateIndex": true, + "batchSize": 100, + "useDistributedIndexing": true, + "partitionSize": 10000 +}' "$APP_RUN_WAIT_TIMEOUT_SECONDS" + tput setaf 2 -echo "✔ OpenMetadata is up and running" -if [[ $startFuseki == "true" ]]; then - echo "✔ RDF/Knowledge Graph support is enabled" - echo " - Fuseki UI: http://localhost:3030" - echo " - SPARQL endpoint: http://localhost:3030/openmetadata/sparql" -fi -echo "" \ No newline at end of file +echo "✔ RDF/Knowledge Graph support is enabled" +echo " - Fuseki UI: http://localhost:3030" +echo " - SPARQL endpoint: http://localhost:3030/openmetadata/sparql" diff --git a/docker/validate_compose.py b/docker/validate_compose.py index c255b23b3d4d..a82fe6179a85 100644 --- a/docker/validate_compose.py +++ b/docker/validate_compose.py @@ -1,3 +1,4 @@ +import os import time from pprint import pprint from typing import Optional, Tuple @@ -13,6 +14,21 @@ PASSWORD = "admin" _access_token: Optional[str] = None +_last_dag_logs_supported: Optional[bool] = None + + +def get_env_int(name: str, default: int) -> int: + value = os.getenv(name) + if value is None: + return default + + try: + return int(value) + except ValueError: + log_ansi_encoded_string( + message=f"Invalid integer for {name}: {value}. Falling back to {default}." + ) + return default def get_access_token() -> str: @@ -50,7 +66,8 @@ def get_last_run_info() -> Tuple[str, str]: """ Make sure we can pick up the latest run info """ - max_retries = 30 + max_retries = get_env_int("VALIDATE_COMPOSE_DAG_RUN_RETRIES", 30) + poll_interval_seconds = get_env_int("VALIDATE_COMPOSE_DAG_RUN_POLL_SECONDS", 5) retries = 0 while retries < max_retries: @@ -89,7 +106,7 @@ def get_last_run_info() -> Tuple[str, str]: _access_token = None log_ansi_encoded_string(message=f"Error getting DAG runs: {e}") - time.sleep(5) + time.sleep(poll_interval_seconds) retries += 1 return None, None @@ -99,19 +116,33 @@ def print_last_run_logs() -> None: """ Show the logs """ + global _last_dag_logs_supported + try: - logs = requests.get( + response = requests.get( f"{AIRFLOW_URL}/api/v2/openmetadata/last_dag_logs?dag_id=sample_data&task_id=ingest_using_recipe", headers=get_auth_headers(), timeout=REQUESTS_TIMEOUT - ).text - pprint(logs) + ) + + if response.status_code == 404: + if _last_dag_logs_supported is not False: + log_ansi_encoded_string( + message="Airflow last_dag_logs route is unavailable. Skipping task log fetch." + ) + _last_dag_logs_supported = False + return + + response.raise_for_status() + _last_dag_logs_supported = True + pprint(response.text) except Exception as e: log_ansi_encoded_string(message=f"Could not fetch logs: {e}") def main(): - max_retries = 15 + max_retries = get_env_int("VALIDATE_COMPOSE_MAX_RETRIES", 15) + retry_interval_seconds = get_env_int("VALIDATE_COMPOSE_RETRY_INTERVAL_SECONDS", 10) retries = 0 while retries < max_retries: @@ -121,7 +152,7 @@ def main(): log_ansi_encoded_string( message="Waiting for DAG run to start...", ) - time.sleep(10) + time.sleep(retry_interval_seconds) retries += 1 continue @@ -135,7 +166,7 @@ def main(): message=f"DAG run [{dag_run_id}] is {state}. Waiting for completion...", ) print_last_run_logs() - time.sleep(10) + time.sleep(retry_interval_seconds) retries += 1 elif state == "failed": log_ansi_encoded_string(message=f"DAG run [{dag_run_id}] FAILED!") @@ -146,7 +177,7 @@ def main(): message=f"Waiting for sample data ingestion. Current state: {state}", ) print_last_run_logs() - time.sleep(10) + time.sleep(retry_interval_seconds) retries += 1 if retries == max_retries: diff --git a/docs/rdf-local-development.md b/docs/rdf-local-development.md index 44b166eaea2f..c46d21bddcd4 100644 --- a/docs/rdf-local-development.md +++ b/docs/rdf-local-development.md @@ -1,6 +1,6 @@ # RDF/Apache Jena Local Development Guide -This guide documents how to set up RDF/Knowledge Graph support for local development with OpenMetadata running in IntelliJ IDEA and Apache Jena Fuseki running in Docker. +This guide documents how to set up RDF/Knowledge Graph support for local development with OpenMetadata and Apache Jena Fuseki. ## Overview @@ -28,16 +28,34 @@ OpenMetadata supports RDF (Resource Description Framework) for knowledge graph c ## Quick Start -### Step 1: Start Apache Jena Fuseki +### Step 1: Choose the Right Startup Mode -Start the Fuseki triple store using Docker Compose: +The standard local Docker flow does not enable RDF or start Fuseki: ```bash cd /path/to/OpenMetadata -docker compose -f docker/development/docker-compose-fuseki.yml up -d +./docker/run_local_docker.sh -d mysql +``` + +For PostgreSQL-based development: + +```bash +./docker/run_local_docker.sh -d postgresql ``` -This starts Fuseki with: +Use the RDF-specific startup script when you want the full Docker stack with Fuseki enabled: + +```bash +./docker/run_local_docker_rdf.sh -d mysql +``` + +For PostgreSQL-based RDF development: + +```bash +./docker/run_local_docker_rdf.sh -d postgresql +``` + +This RDF startup path starts OpenMetadata, the backing database, search, ingestion services, and Fuseki with: - **Port**: 3030 - **Admin Password**: admin - **Dataset**: openmetadata @@ -59,7 +77,17 @@ The Fuseki web UI is available at `http://localhost:3030` with credentials: ### Step 3: Configure IntelliJ Run Configuration -Create or modify your IntelliJ run configuration for `OpenMetadataApplication` with these environment variables: +If you are running the full RDF Docker stack with `run_local_docker_rdf.sh`, the Docker services already receive the RDF environment variables automatically. + +If you want to run the OpenMetadata server directly from IntelliJ while keeping Fuseki in Docker, start Fuseki separately: + +```bash +docker compose -f docker/development/docker-compose.yml -f docker/development/docker-compose-fuseki.yml up -d fuseki +``` + +If your local backend uses PostgreSQL, swap `docker-compose.yml` for `docker-compose-postgres.yml`. + +Create or modify your IntelliJ run configuration for `OpenMetadataApplication` with these environment variables only when you want to run the OpenMetadata server directly from IntelliJ while keeping Fuseki in Docker: ``` RDF_ENABLED=true @@ -211,7 +239,7 @@ Trigger the RDF indexing application to populate the triple store with existing curl -X POST \ -H "Authorization: Bearer " \ -H "Content-Type: application/json" \ - -d '{"entities": ["all"], "recreateIndex": true, "batchSize": 100}' \ + -d '{"entities": [], "recreateIndex": true, "batchSize": 100}' \ http://localhost:8585/api/v1/apps/trigger/RdfIndexApp ``` diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/server/auth/handlers/McpCallbackServlet.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/server/auth/handlers/McpCallbackServlet.java index 53d246d6d3c7..278e5999c57b 100644 --- a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/server/auth/handlers/McpCallbackServlet.java +++ b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/server/auth/handlers/McpCallbackServlet.java @@ -18,12 +18,13 @@ import jakarta.servlet.http.HttpSession; import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.OutputStreamWriter; +import java.io.PrintWriter; +import java.nio.charset.Charset; import java.nio.charset.StandardCharsets; import java.util.List; import java.util.Map; import java.util.TreeMap; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; import lombok.extern.slf4j.Slf4j; import org.openmetadata.mcp.server.auth.provider.UserSSOOAuthProvider; import org.openmetadata.mcp.server.auth.repository.McpPendingAuthRequestRepository; @@ -104,7 +105,7 @@ private String resolveBaseUrl() { return mcpConfig.getBaseUrl(); } } catch (Exception e) { - LOG.warn("Failed to get base URL from MCP config: {}", e.getMessage()); + LOG.warn("Failed to get base URL from MCP config", e); } try { SystemRepository systemRepository = Entity.getSystemRepository(); @@ -119,8 +120,13 @@ private String resolveBaseUrl() { } } } catch (Exception e) { - LOG.warn("Could not get base URL from system settings: {}", e.getMessage()); + LOG.warn("Could not get base URL from system settings", e); } + + LOG.error( + "No base URL configured in MCP settings or system settings. " + + "Falling back to http://localhost:8585 — this is only suitable for local development. " + + "Configure a proper base URL for production deployments."); return "http://localhost:8585"; } @@ -296,54 +302,21 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) AuthenticationCodeFlowHandler.SESSION_REDIRECT_URI, baseUrl + "/mcp/callback"); LOG.debug("Set session SSO callback URL to: {}", ssoCallbackUrl); - AtomicBoolean handlerWroteError = new AtomicBoolean(false); - AtomicInteger capturedStatus = new AtomicInteger(0); - ByteArrayOutputStream errorSink = new ByteArrayOutputStream(); - HttpServletResponseWrapper responseWrapper = - new HttpServletResponseWrapper(response) { - @Override - public void sendRedirect(String location) throws IOException { - LOG.info("Intercepted redirect from SSO handler"); - LOG.debug("Redirect target: {}", location); - } - - @Override - public void setStatus(int sc) { - if (sc >= 400) { - capturedStatus.set(sc); - handlerWroteError.set(true); - } - } - - @Override - public ServletOutputStream getOutputStream() throws IOException { - if (handlerWroteError.get()) { - return new ServletOutputStream() { - @Override - public void write(int b) { - errorSink.write(b); - } - - @Override - public boolean isReady() { - return true; - } - - @Override - public void setWriteListener(WriteListener writeListener) {} - }; - } - return super.getOutputStream(); - } - }; - - ssoHandler.handleCallback(request, responseWrapper); - - if (handlerWroteError.get()) { - String errorBody = errorSink.toString(StandardCharsets.UTF_8); - LOG.error("SSO token exchange failed with HTTP {}: {}", capturedStatus.get(), errorBody); + BufferedServletResponseWrapper callbackResponse = + new BufferedServletResponseWrapper(response); + ssoHandler.handleCallback(request, callbackResponse); + + if (callbackResponse.getStatusCode() >= HttpServletResponse.SC_BAD_REQUEST) { + String errorBody = callbackResponse.getCapturedBody(); + LOG.error( + "SSO token exchange failed with HTTP {}: {}", + callbackResponse.getStatusCode(), + errorBody); throw new IllegalStateException( - "SSO provider token exchange failed (HTTP " + capturedStatus.get() + "): " + errorBody); + "SSO provider token exchange failed (HTTP " + + callbackResponse.getStatusCode() + + "): " + + errorBody); } OidcCredentials credentials = (OidcCredentials) session.getAttribute(OIDC_CREDENTIAL_PROFILE); @@ -386,8 +359,15 @@ public void setWriteListener(WriteListener writeListener) {} LOG.debug("Extracted user identity from SSO callback"); - userSSOProvider.handleSSOCallbackWithDbState( - request, response, userName, email, "mcp:" + pendingRequest.authRequestId()); + processBufferedCallbackResponse( + response, + wrappedResponse -> + userSSOProvider.handleSSOCallbackWithDbState( + request, + wrappedResponse, + userName, + email, + "mcp:" + pendingRequest.authRequestId())); LOG.info("MCP OAuth SSO callback completed successfully"); @@ -473,9 +453,195 @@ private void handleDirectIdTokenFlow( session.removeAttribute("mcp.auth.request.id"); - userSSOProvider.handleSSOCallbackWithDbState( - request, response, userName, email, "mcp:" + authRequestId); + processBufferedCallbackResponse( + response, + wrappedResponse -> + userSSOProvider.handleSSOCallbackWithDbState( + request, wrappedResponse, userName, email, "mcp:" + authRequestId)); LOG.info("MCP OAuth direct ID token flow completed successfully"); } + + private void processBufferedCallbackResponse( + HttpServletResponse response, BufferedResponseAction action) throws Exception { + BufferedServletResponseWrapper bufferedResponse = new BufferedServletResponseWrapper(response); + action.execute(bufferedResponse); + bufferedResponse.commitTo(response); + } + + @FunctionalInterface + private interface BufferedResponseAction { + void execute(HttpServletResponse response) throws Exception; + } + + private static final class BufferedServletResponseWrapper extends HttpServletResponseWrapper { + private final ByteArrayOutputStream buffer = new ByteArrayOutputStream(); + private final ServletOutputStream outputStream = new BufferedServletOutputStream(buffer); + private PrintWriter writer; + private Integer statusCode = HttpServletResponse.SC_OK; + private String contentType; + private String characterEncoding = StandardCharsets.UTF_8.name(); + private String redirectLocation; + private boolean committed; + private boolean outputStreamRequested; + private boolean writerRequested; + + BufferedServletResponseWrapper(HttpServletResponse response) { + super(response); + } + + int getStatusCode() { + return statusCode; + } + + String getCapturedBody() throws IOException { + flushBuffer(); + return buffer.toString(StandardCharsets.UTF_8); + } + + void commitTo(HttpServletResponse response) throws IOException { + flushBuffer(); + if (redirectLocation != null) { + response.sendRedirect(redirectLocation); + return; + } + + response.setStatus(statusCode); + if (characterEncoding != null) { + response.setCharacterEncoding(characterEncoding); + } + if (contentType != null) { + response.setContentType(contentType); + } + if (buffer.size() > 0) { + response.getOutputStream().write(buffer.toByteArray()); + } + response.flushBuffer(); + } + + @Override + public void setStatus(int sc) { + statusCode = sc; + } + + @Override + public void sendError(int sc) { + statusCode = sc; + committed = true; + } + + @Override + public void sendError(int sc, String msg) throws IOException { + statusCode = sc; + committed = true; + if (msg != null) { + writeToBuffer(msg); + } + } + + @Override + public void sendRedirect(String location) { + redirectLocation = location; + committed = true; + } + + @Override + public void setContentType(String type) { + contentType = type; + } + + @Override + public String getContentType() { + return contentType; + } + + @Override + public void setCharacterEncoding(String charset) { + if (!writerRequested) { + characterEncoding = charset; + } + } + + @Override + public String getCharacterEncoding() { + return characterEncoding; + } + + @Override + public ServletOutputStream getOutputStream() { + if (writerRequested) { + throw new IllegalStateException("getWriter() has already been called on this response"); + } + outputStreamRequested = true; + return outputStream; + } + + @Override + public PrintWriter getWriter() { + if (outputStreamRequested) { + throw new IllegalStateException( + "getOutputStream() has already been called on this response"); + } + if (writer == null) { + writer = + new PrintWriter( + new OutputStreamWriter( + buffer, + Charset.forName( + characterEncoding != null + ? characterEncoding + : StandardCharsets.UTF_8.name())), + true); + } + writerRequested = true; + return writer; + } + + @Override + public void flushBuffer() throws IOException { + if (writer != null) { + writer.flush(); + } + if (outputStreamRequested) { + outputStream.flush(); + } + committed = true; + } + + @Override + public boolean isCommitted() { + return committed; + } + + private void writeToBuffer(String value) throws IOException { + if (writer != null) { + writer.flush(); + } + buffer.write( + value.getBytes( + Charset.forName( + characterEncoding != null ? characterEncoding : StandardCharsets.UTF_8.name()))); + } + } + + private static final class BufferedServletOutputStream extends ServletOutputStream { + private final ByteArrayOutputStream buffer; + + private BufferedServletOutputStream(ByteArrayOutputStream buffer) { + this.buffer = buffer; + } + + @Override + public void write(int b) { + buffer.write(b); + } + + @Override + public boolean isReady() { + return true; + } + + @Override + public void setWriteListener(WriteListener writeListener) {} + } } diff --git a/openmetadata-mcp/src/test/java/org/openmetadata/mcp/server/auth/handlers/McpCallbackServletTest.java b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/server/auth/handlers/McpCallbackServletTest.java new file mode 100644 index 000000000000..038912f035d6 --- /dev/null +++ b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/server/auth/handlers/McpCallbackServletTest.java @@ -0,0 +1,78 @@ +package org.openmetadata.mcp.server.auth.handlers; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import jakarta.servlet.ServletOutputStream; +import jakarta.servlet.WriteListener; +import jakarta.servlet.http.HttpServletResponse; +import java.io.ByteArrayOutputStream; +import java.lang.reflect.Constructor; +import java.lang.reflect.Method; +import java.nio.charset.StandardCharsets; +import org.junit.jupiter.api.Test; + +class McpCallbackServletTest { + + @Test + void bufferedResponseSendErrorSupportsExistingOutputStreamUsage() throws Exception { + HttpServletResponse delegateResponse = mock(HttpServletResponse.class); + ByteArrayOutputStream committedBody = new ByteArrayOutputStream(); + + when(delegateResponse.getOutputStream()) + .thenReturn(new CapturingServletOutputStream(committedBody)); + + HttpServletResponse bufferedResponse = newBufferedResponse(delegateResponse); + bufferedResponse.getOutputStream().write("prefix:".getBytes(StandardCharsets.UTF_8)); + + assertDoesNotThrow( + () -> bufferedResponse.sendError(HttpServletResponse.SC_BAD_REQUEST, "invalid-request")); + + commitTo(bufferedResponse, delegateResponse); + + verify(delegateResponse).setStatus(HttpServletResponse.SC_BAD_REQUEST); + assertThat(committedBody.toString(StandardCharsets.UTF_8)).isEqualTo("prefix:invalid-request"); + } + + private static HttpServletResponse newBufferedResponse(HttpServletResponse response) + throws Exception { + Class wrapperClass = + Class.forName( + "org.openmetadata.mcp.server.auth.handlers.McpCallbackServlet$BufferedServletResponseWrapper"); + Constructor constructor = wrapperClass.getDeclaredConstructor(HttpServletResponse.class); + constructor.setAccessible(true); + return (HttpServletResponse) constructor.newInstance(response); + } + + private static void commitTo(HttpServletResponse bufferedResponse, HttpServletResponse response) + throws Exception { + Method method = + bufferedResponse.getClass().getDeclaredMethod("commitTo", HttpServletResponse.class); + method.setAccessible(true); + method.invoke(bufferedResponse, response); + } + + private static final class CapturingServletOutputStream extends ServletOutputStream { + private final ByteArrayOutputStream outputStream; + + private CapturingServletOutputStream(ByteArrayOutputStream outputStream) { + this.outputStream = outputStream; + } + + @Override + public void write(int b) { + outputStream.write(b); + } + + @Override + public boolean isReady() { + return true; + } + + @Override + public void setWriteListener(WriteListener writeListener) {} + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java index 4b0eab6398da..fe7716a8b9ca 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java @@ -83,6 +83,7 @@ import org.openmetadata.service.apps.ApplicationContext; import org.openmetadata.service.apps.ApplicationHandler; import org.openmetadata.service.apps.McpServerProvider; +import org.openmetadata.service.apps.bundles.rdf.distributed.RdfDistributedJobParticipant; import org.openmetadata.service.apps.bundles.searchIndex.distributed.DistributedJobParticipant; import org.openmetadata.service.apps.bundles.searchIndex.distributed.ServerIdentityResolver; import org.openmetadata.service.apps.scheduler.AppScheduler; @@ -379,6 +380,7 @@ public void run(OpenMetadataApplicationConfig catalogConfig, Environment environ // Register Distributed Job Participant for distributed search indexing registerDistributedJobParticipant(environment, jdbi, catalogConfig.getCacheConfig()); + registerDistributedRdfJobParticipant(environment, jdbi); // Register Event publishers registerEventPublisher(catalogConfig); @@ -1125,7 +1127,18 @@ protected void registerDistributedJobParticipant( "Registered DistributedJobParticipant for distributed search indexing using {}", notifierType); } catch (Exception e) { - LOG.warn("Failed to register DistributedJobParticipant: {}", e.getMessage()); + LOG.warn("Failed to register DistributedJobParticipant", e); + } + } + + protected void registerDistributedRdfJobParticipant(Environment environment, Jdbi jdbi) { + try { + CollectionDAO collectionDAO = jdbi.onDemand(CollectionDAO.class); + RdfDistributedJobParticipant participant = new RdfDistributedJobParticipant(collectionDAO); + environment.lifecycle().manage(participant); + LOG.info("Registered RdfDistributedJobParticipant for distributed RDF indexing"); + } catch (Exception e) { + LOG.warn("Failed to register RdfDistributedJobParticipant", e); } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/RdfBatchProcessor.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/RdfBatchProcessor.java new file mode 100644 index 000000000000..777205916d8e --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/RdfBatchProcessor.java @@ -0,0 +1,230 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.UUID; +import java.util.function.BooleanSupplier; +import java.util.stream.Collectors; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.entity.data.GlossaryTerm; +import org.openmetadata.schema.type.LineageDetails; +import org.openmetadata.schema.type.Relationship; +import org.openmetadata.schema.type.TermRelation; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.Entity; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.jdbi3.CollectionDAO.EntityRelationshipObject; +import org.openmetadata.service.rdf.RdfRepository; + +@Slf4j +public class RdfBatchProcessor { + public static final List ALL_RELATIONSHIPS = + java.util.Arrays.stream(Relationship.values()).map(Relationship::ordinal).toList(); + + public static final Set EXCLUDED_RELATIONSHIP_ENTITY_TYPES = + Set.of( + "changeEvent", + Entity.AUDIT_LOG, + Entity.WEB_ANALYTIC_EVENT, + "entityUsage", + "eventSubscription", + Entity.EVENT_SUBSCRIPTION, + "vote", + Entity.THREAD); + + public static final Set EXCLUDED_RELATIONSHIP_TYPES = + Set.of(Relationship.VOTED.ordinal(), Relationship.FOLLOWS.ordinal()); + + private final CollectionDAO collectionDAO; + private final RdfRepository rdfRepository; + + public RdfBatchProcessor(CollectionDAO collectionDAO, RdfRepository rdfRepository) { + this.collectionDAO = collectionDAO; + this.rdfRepository = rdfRepository; + } + + public BatchProcessingResult processEntities( + String entityType, List entities, BooleanSupplier stopRequested) { + if (entities == null || entities.isEmpty()) { + return new BatchProcessingResult(0, 0); + } + + BooleanSupplier effectiveStopRequested = stopRequested != null ? stopRequested : () -> false; + int successCount = 0; + int failedCount = 0; + List indexedEntities = new ArrayList<>(); + + for (EntityInterface entity : entities) { + if (effectiveStopRequested.getAsBoolean()) { + break; + } + try { + rdfRepository.createOrUpdate(entity); + indexedEntities.add(entity); + successCount++; + } catch (Exception e) { + LOG.error("Failed to index entity {} to RDF", entity.getId(), e); + failedCount++; + } + } + + if (!indexedEntities.isEmpty()) { + processBatchRelationships(entityType, indexedEntities); + if ("glossaryTerm".equals(entityType)) { + processGlossaryTermRelations(indexedEntities, effectiveStopRequested); + } + } + + return new BatchProcessingResult(successCount, failedCount); + } + + public void processBatchRelationships( + String entityType, List entities) { + if (entities == null || entities.isEmpty()) { + return; + } + + try { + List entityIds = + entities.stream().map(entity -> entity.getId().toString()).collect(Collectors.toList()); + + List outgoingRelationships = + collectionDAO + .relationshipDAO() + .findToBatchWithRelations(entityIds, entityType, ALL_RELATIONSHIPS); + + List incomingLineage = + collectionDAO + .relationshipDAO() + .findFromBatch( + entityIds, + Relationship.UPSTREAM.ordinal(), + org.openmetadata.schema.type.Include.ALL); + + List allRelationships = new ArrayList<>(); + + for (EntityRelationshipObject rel : outgoingRelationships) { + if (shouldSkipRelationship(rel)) { + continue; + } + + if (rel.getRelation() == Relationship.UPSTREAM.ordinal() && rel.getJson() != null) { + processLineageRelationship(rel); + } else { + if ("glossaryTerm".equals(entityType) + && rel.getRelation() == Relationship.RELATED_TO.ordinal() + && "glossaryTerm".equals(rel.getToEntity())) { + continue; + } + allRelationships.add(convertToEntityRelationship(rel)); + } + } + + for (EntityRelationshipObject rel : incomingLineage) { + if (shouldSkipRelationship(rel)) { + continue; + } + + if (rel.getJson() != null) { + processLineageRelationship(rel); + } else { + allRelationships.add(convertToEntityRelationship(rel)); + } + } + + if (!allRelationships.isEmpty()) { + rdfRepository.bulkAddRelationships(allRelationships); + } + } catch (Exception e) { + LOG.error("Failed to process batch relationships for entity type {}", entityType, e); + } + } + + public org.openmetadata.schema.type.EntityRelationship convertToEntityRelationship( + EntityRelationshipObject rel) { + return new org.openmetadata.schema.type.EntityRelationship() + .withFromEntity(rel.getFromEntity()) + .withFromId(UUID.fromString(rel.getFromId())) + .withToEntity(rel.getToEntity()) + .withToId(UUID.fromString(rel.getToId())) + .withRelation(rel.getRelation()) + .withRelationshipType(Relationship.values()[rel.getRelation()]); + } + + private boolean shouldSkipRelationship(EntityRelationshipObject rel) { + return EXCLUDED_RELATIONSHIP_ENTITY_TYPES.contains(rel.getToEntity()) + || EXCLUDED_RELATIONSHIP_ENTITY_TYPES.contains(rel.getFromEntity()) + || EXCLUDED_RELATIONSHIP_TYPES.contains(rel.getRelation()); + } + + void processLineageRelationship(EntityRelationshipObject rel) { + try { + UUID fromId = UUID.fromString(rel.getFromId()); + UUID toId = UUID.fromString(rel.getToId()); + LineageDetails lineageDetails = JsonUtils.readValue(rel.getJson(), LineageDetails.class); + rdfRepository.addLineageWithDetails( + rel.getFromEntity(), fromId, rel.getToEntity(), toId, lineageDetails); + } catch (Exception e) { + LOG.debug("Failed to parse lineage details, falling back to basic relationship", e); + try { + rdfRepository.addRelationship(convertToEntityRelationship(rel)); + } catch (Exception ex) { + LOG.debug("Failed to add basic lineage relationship", ex); + } + } + } + + void processGlossaryTermRelations( + List entities, BooleanSupplier stopRequested) { + List relations = new ArrayList<>(); + + for (EntityInterface entity : entities) { + if (stopRequested.getAsBoolean()) { + break; + } + + if (!(entity instanceof GlossaryTerm glossaryTerm)) { + continue; + } + + List relatedTerms = glossaryTerm.getRelatedTerms(); + if (relatedTerms == null || relatedTerms.isEmpty()) { + continue; + } + + UUID fromTermId = glossaryTerm.getId(); + for (TermRelation termRelation : relatedTerms) { + if (termRelation.getTerm() == null || termRelation.getTerm().getId() == null) { + continue; + } + + String relationType = + termRelation.getRelationType() != null ? termRelation.getRelationType() : "relatedTo"; + relations.add( + new RdfRepository.GlossaryTermRelationData( + fromTermId, termRelation.getTerm().getId(), relationType)); + } + } + + if (!relations.isEmpty()) { + rdfRepository.bulkAddGlossaryTermRelations(relations); + } + } + + public record BatchProcessingResult(int successCount, int failedCount) {} +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/RdfIndexApp.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/RdfIndexApp.java index 139261709d53..7f3f9809c9a9 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/RdfIndexApp.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/RdfIndexApp.java @@ -10,19 +10,22 @@ import jakarta.ws.rs.core.Response; import java.util.ArrayList; import java.util.HashSet; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Set; import java.util.UUID; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CancellationException; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; -import java.util.stream.Collectors; import lombok.Getter; import lombok.extern.slf4j.Slf4j; import org.openmetadata.schema.EntityInterface; @@ -30,20 +33,19 @@ import org.openmetadata.schema.entity.app.AppRunRecord; import org.openmetadata.schema.entity.app.FailureContext; import org.openmetadata.schema.entity.app.SuccessContext; -import org.openmetadata.schema.entity.data.GlossaryTerm; import org.openmetadata.schema.system.EntityStats; import org.openmetadata.schema.system.EventPublisherJob; import org.openmetadata.schema.system.IndexingError; import org.openmetadata.schema.system.Stats; import org.openmetadata.schema.system.StepStats; import org.openmetadata.schema.type.Include; -import org.openmetadata.schema.type.LineageDetails; -import org.openmetadata.schema.type.Relationship; -import org.openmetadata.schema.type.TermRelation; import org.openmetadata.schema.utils.JsonUtils; import org.openmetadata.schema.utils.ResultList; import org.openmetadata.service.Entity; import org.openmetadata.service.apps.AbstractNativeApplication; +import org.openmetadata.service.apps.bundles.rdf.distributed.DistributedRdfIndexExecutor; +import org.openmetadata.service.apps.bundles.rdf.distributed.RdfDistributedJobStatsAggregator; +import org.openmetadata.service.apps.bundles.rdf.distributed.RdfIndexJob; import org.openmetadata.service.exception.AppException; import org.openmetadata.service.jdbi3.CollectionDAO; import org.openmetadata.service.jdbi3.CollectionDAO.EntityRelationshipObject; @@ -66,29 +68,14 @@ public class RdfIndexApp extends AbstractNativeApplication { private static final int MAX_CONSUMER_THREADS = 5; private static final long WEBSOCKET_UPDATE_INTERVAL_MS = 2000; - private static final List ALL_RELATIONSHIPS = - java.util.Arrays.stream(Relationship.values()) - .map(Relationship::ordinal) - .collect(Collectors.toList()); - - // Entity types that should be excluded from RDF relationships as they don't provide - // meaningful semantic value (operational/audit entities) + private static final List ALL_RELATIONSHIPS = RdfBatchProcessor.ALL_RELATIONSHIPS; private static final Set EXCLUDED_RELATIONSHIP_ENTITY_TYPES = - Set.of( - "changeEvent", - "auditLog", - "webAnalyticEvent", - "entityUsage", - "eventSubscription", - "vote", - "THREAD"); - - // Relationship types that should be excluded from RDF as they don't provide - // meaningful semantic relationships (user interactions, not data relationships) + RdfBatchProcessor.EXCLUDED_RELATIONSHIP_ENTITY_TYPES; private static final Set EXCLUDED_RELATIONSHIP_TYPES = - Set.of(Relationship.VOTED.ordinal(), Relationship.FOLLOWS.ordinal()); + RdfBatchProcessor.EXCLUDED_RELATIONSHIP_TYPES; private final RdfRepository rdfRepository; + private final RdfBatchProcessor batchProcessor; private volatile boolean stopped = false; private volatile long lastWebSocketUpdate = 0; @@ -100,6 +87,7 @@ public class RdfIndexApp extends AbstractNativeApplication { private final AtomicReference rdfIndexStats = new AtomicReference<>(); private final AtomicBoolean producersDone = new AtomicBoolean(false); private BlockingQueue taskQueue; + private volatile DistributedRdfIndexExecutor distributedExecutor; record IndexingTask( String entityType, List entities, int offset, int retryCount) { @@ -115,6 +103,7 @@ boolean isPoisonPill() { public RdfIndexApp(CollectionDAO collectionDAO, SearchRepository searchRepository) { super(collectionDAO, searchRepository); this.rdfRepository = RdfRepository.getInstance(); + this.batchProcessor = new RdfBatchProcessor(collectionDAO, rdfRepository); } @Override @@ -160,9 +149,10 @@ public void execute(JobExecutionContext jobExecutionContext) { } try { - boolean containsAll = jobData.getEntities().contains(ALL); - if (containsAll) { - jobData.setEntities(getAll()); + jobData.setEntities(resolveEntityTypes(jobData.getEntities())); + if (jobData.getEntities().isEmpty()) { + throw new IllegalStateException( + "No repository-backed entity types configured for RDF indexing"); } LOG.info( @@ -178,7 +168,11 @@ public void execute(JobExecutionContext jobExecutionContext) { } updateJobStatus(EventPublisherJob.Status.RUNNING); - reIndexFromStartToEnd(); + if (Boolean.TRUE.equals(jobData.getUseDistributedIndexing())) { + reIndexDistributed(); + } else { + reIndexFromStartToEnd(); + } if (stopped) { updateJobStatus(EventPublisherJob.Status.STOPPED); @@ -208,6 +202,11 @@ private void initializeJob(JobExecutionContext jobExecutionContext) { rdfIndexStats.set(initializeTotalRecords(jobData.getEntities())); jobData.setStats(rdfIndexStats.get()); + if (Boolean.TRUE.equals(jobData.getUseDistributedIndexing())) { + sendUpdates(jobExecutionContext, true); + return; + } + int queueSize = jobData.getQueueSize() != null ? jobData.getQueueSize() : DEFAULT_QUEUE_SIZE; int effectiveQueueSize = calculateMemoryAwareQueueSize(queueSize); taskQueue = new LinkedBlockingQueue<>(effectiveQueueSize); @@ -236,6 +235,35 @@ private void clearRdfData() { } } + private void reIndexDistributed() throws InterruptedException { + int partitionSize = jobData.getPartitionSize() != null ? jobData.getPartitionSize() : 10000; + String createdBy = + getApp() != null && getApp().getName() != null ? getApp().getName() : "system"; + + distributedExecutor = new DistributedRdfIndexExecutor(collectionDAO, partitionSize); + distributedExecutor.performStartupRecovery(); + + RdfIndexJob distributedJob = + distributedExecutor.createJob(jobData.getEntities(), jobData, createdBy); + + ExecutorService distributedExecutionExecutor = + Executors.newSingleThreadExecutor( + Thread.ofVirtual().name("rdf-distributed-execution-", 0).factory()); + Future distributedExecution = + distributedExecutionExecutor.submit( + () -> { + distributedExecutor.execute(jobData); + return null; + }); + + try { + monitorDistributedJob(distributedJob.getId(), distributedExecution); + awaitDistributedExecution(distributedExecution); + } finally { + distributedExecutionExecutor.shutdownNow(); + } + } + private void reIndexFromStartToEnd() throws InterruptedException { long totalEntities = rdfIndexStats.get().getJobStats().getTotalRecords(); int numProducers = Math.clamp((int) (totalEntities / 5000), 2, MAX_PRODUCER_THREADS); @@ -266,10 +294,6 @@ private void reIndexFromStartToEnd() throws InterruptedException { } try { - // Clear entire RDF store before re-indexing to remove stale data - LOG.info("Clearing RDF store before re-indexing"); - rdfRepository.clearAll(); - processEntityTypes(); signalConsumersToStop(numConsumers); consumerLatch.await(); @@ -282,6 +306,73 @@ private void reIndexFromStartToEnd() throws InterruptedException { } } + private void monitorDistributedJob(UUID jobId, Future distributedExecution) + throws InterruptedException { + RdfDistributedJobStatsAggregator statsAggregator = new RdfDistributedJobStatsAggregator(); + + while (!stopped) { + RdfIndexJob latestJob = + distributedExecutor != null ? distributedExecutor.getJobWithFreshStats() : null; + if (latestJob != null) { + Stats aggregatedStats = statsAggregator.toStats(latestJob); + rdfIndexStats.set(aggregatedStats); + jobData.setStats(aggregatedStats); + sendUpdates(jobExecutionContext, false); + + if (latestJob.isTerminal()) { + if (latestJob.getStatus() + == org.openmetadata + .service + .apps + .bundles + .searchIndex + .distributed + .IndexJobStatus + .STOPPED) { + stopped = true; + } else if (latestJob.getStatus() + == org.openmetadata + .service + .apps + .bundles + .searchIndex + .distributed + .IndexJobStatus + .FAILED) { + jobData.setFailure( + new IndexingError() + .withErrorSource(IndexingError.ErrorSource.JOB) + .withMessage(latestJob.getErrorMessage())); + } + return; + } + } + + if (distributedExecution.isDone()) { + return; + } + + TimeUnit.SECONDS.sleep(2); + } + } + + private void awaitDistributedExecution(Future distributedExecution) + throws InterruptedException { + try { + distributedExecution.get(); + } catch (CancellationException e) { + if (!stopped) { + throw new RuntimeException("Distributed RDF execution was cancelled unexpectedly", e); + } + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException runtimeException) { + throw runtimeException; + } + throw new RuntimeException("Distributed RDF execution failed", cause); + } + } + private void runConsumer(int consumerId, CountDownLatch consumerLatch) { LOG.info("Consumer {} started", consumerId); try { @@ -310,220 +401,40 @@ private void processTask(IndexingTask task) { return; } - int successCount = 0; - int failedCount = 0; - try { - for (EntityInterface entity : entities) { - if (stopped) { - break; - } - try { - rdfRepository.createOrUpdate(entity); - successCount++; - } catch (Exception e) { - LOG.error("Failed to index entity {} to RDF", entity.getId(), e); - failedCount++; - } - } - - processBatchRelationships(entityType, entities); - - // Process glossary term relations if this is a glossaryTerm batch - if ("glossaryTerm".equals(entityType)) { - processGlossaryTermRelations(entities); - } + RdfBatchProcessor.BatchProcessingResult result = + batchProcessor.processEntities(entityType, entities, () -> stopped); StepStats currentStats = - new StepStats().withSuccessRecords(successCount).withFailedRecords(failedCount); + new StepStats() + .withSuccessRecords(result.successCount()) + .withFailedRecords(result.failedCount()); updateEntityStats(entityType, currentStats); sendUpdates(jobExecutionContext, false); } catch (Exception e) { LOG.error("Error processing batch for entity type {}", entityType, e); updateEntityStats( - entityType, - new StepStats() - .withSuccessRecords(successCount) - .withFailedRecords(entities.size() - successCount)); + entityType, new StepStats().withSuccessRecords(0).withFailedRecords(entities.size())); } } private void processBatchRelationships( String entityType, List entities) { - if (entities.isEmpty()) { - return; - } - - List entityIds = - entities.stream().map(e -> e.getId().toString()).collect(Collectors.toList()); - - try { - List outgoingRelationships = - collectionDAO - .relationshipDAO() - .findToBatchWithRelations(entityIds, entityType, ALL_RELATIONSHIPS); - - List incomingLineage = - collectionDAO - .relationshipDAO() - .findFromBatch(entityIds, Relationship.UPSTREAM.ordinal(), Include.ALL); - - List allRelationships = new ArrayList<>(); - - for (EntityRelationshipObject rel : outgoingRelationships) { - // Skip relationships to/from excluded entity types (changeEvent, auditLog, vote, etc.) - // These don't provide meaningful semantic value in the knowledge graph - if (EXCLUDED_RELATIONSHIP_ENTITY_TYPES.contains(rel.getToEntity()) - || EXCLUDED_RELATIONSHIP_ENTITY_TYPES.contains(rel.getFromEntity())) { - LOG.debug( - "Skipping relationship {} -> {} (excluded entity type: {} or {})", - rel.getFromId(), - rel.getToId(), - rel.getFromEntity(), - rel.getToEntity()); - continue; - } - - // Skip excluded relationship types (VOTED, FOLLOWS, etc.) - if (EXCLUDED_RELATIONSHIP_TYPES.contains(rel.getRelation())) { - LOG.debug( - "Skipping relationship {} -> {} (excluded relationship type: {})", - rel.getFromId(), - rel.getToId(), - rel.getRelation()); - continue; - } - - if (rel.getRelation() == Relationship.UPSTREAM.ordinal() && rel.getJson() != null) { - processLineageRelationship(rel); - } else { - // Skip glossary term RELATED_TO relationships - they're handled separately - // by processGlossaryTermRelations() with typed predicates - if ("glossaryTerm".equals(entityType) - && rel.getRelation() == Relationship.RELATED_TO.ordinal() - && "glossaryTerm".equals(rel.getToEntity())) { - LOG.debug( - "Skipping glossary term relation {} -> {} (handled by processGlossaryTermRelations)", - rel.getFromId(), - rel.getToId()); - continue; - } - allRelationships.add(convertToEntityRelationship(rel)); - } - } - - for (EntityRelationshipObject rel : incomingLineage) { - // Skip relationships to/from excluded entity types - if (EXCLUDED_RELATIONSHIP_ENTITY_TYPES.contains(rel.getToEntity()) - || EXCLUDED_RELATIONSHIP_ENTITY_TYPES.contains(rel.getFromEntity())) { - continue; - } - - // Skip excluded relationship types - if (EXCLUDED_RELATIONSHIP_TYPES.contains(rel.getRelation())) { - continue; - } - - if (rel.getJson() != null) { - processLineageRelationship(rel); - } else { - allRelationships.add(convertToEntityRelationship(rel)); - } - } - - if (!allRelationships.isEmpty()) { - rdfRepository.bulkAddRelationships(allRelationships); - LOG.debug( - "Bulk added {} relationships for {} entities", - allRelationships.size(), - entities.size()); - } - - } catch (Exception e) { - LOG.error("Failed to process batch relationships for entity type {}", entityType, e); - } + batchProcessor.processBatchRelationships(entityType, entities); } private void processLineageRelationship(EntityRelationshipObject rel) { - try { - UUID fromId = UUID.fromString(rel.getFromId()); - UUID toId = UUID.fromString(rel.getToId()); - LineageDetails lineageDetails = JsonUtils.readValue(rel.getJson(), LineageDetails.class); - rdfRepository.addLineageWithDetails( - rel.getFromEntity(), fromId, rel.getToEntity(), toId, lineageDetails); - LOG.debug( - "Added lineage with details from {}/{} to {}/{}", - rel.getFromEntity(), - fromId, - rel.getToEntity(), - toId); - } catch (Exception e) { - LOG.debug("Failed to parse lineage details, falling back to basic relationship", e); - try { - rdfRepository.addRelationship(convertToEntityRelationship(rel)); - } catch (Exception ex) { - LOG.debug("Failed to add basic lineage relationship", ex); - } - } + batchProcessor.processLineageRelationship(rel); } private void processGlossaryTermRelations(List entities) { - List relations = new ArrayList<>(); - - for (EntityInterface entity : entities) { - if (stopped) { - break; - } - - if (entity instanceof GlossaryTerm glossaryTerm) { - List relatedTerms = glossaryTerm.getRelatedTerms(); - if (relatedTerms != null && !relatedTerms.isEmpty()) { - UUID fromTermId = glossaryTerm.getId(); - LOG.info( - "Processing glossary term {} ({}) with {} relations", - glossaryTerm.getName(), - fromTermId, - relatedTerms.size()); - - for (TermRelation termRelation : relatedTerms) { - if (termRelation.getTerm() != null && termRelation.getTerm().getId() != null) { - UUID toTermId = termRelation.getTerm().getId(); - String relationType = - termRelation.getRelationType() != null - ? termRelation.getRelationType() - : "relatedTo"; - - LOG.info( - " Relation: {} -> {} (type: {}, raw: {})", - glossaryTerm.getName(), - termRelation.getTerm().getName(), - relationType, - termRelation.getRelationType()); - - relations.add( - new RdfRepository.GlossaryTermRelationData(fromTermId, toTermId, relationType)); - } - } - } - } - } - - if (!relations.isEmpty()) { - rdfRepository.bulkAddGlossaryTermRelations(relations); - LOG.info("Added {} glossary term relations to RDF store", relations.size()); - } + batchProcessor.processGlossaryTermRelations(entities, () -> stopped); } private org.openmetadata.schema.type.EntityRelationship convertToEntityRelationship( EntityRelationshipObject rel) { - return new org.openmetadata.schema.type.EntityRelationship() - .withFromEntity(rel.getFromEntity()) - .withFromId(UUID.fromString(rel.getFromId())) - .withToEntity(rel.getToEntity()) - .withToId(UUID.fromString(rel.getToId())) - .withRelation(rel.getRelation()) - .withRelationshipType(Relationship.values()[rel.getRelation()]); + return batchProcessor.convertToEntityRelationship(rel); } private void processEntityTypes() throws InterruptedException { @@ -755,6 +666,7 @@ public void updateRecordToDbAndNotify(JobExecutionContext jobExecutionContext) { appRecord.setSuccessContext( new SuccessContext().withAdditionalProperty("stats", jobData.getStats())); } + pushAppStatusUpdates(jobExecutionContext, appRecord, true); if (WebSocketManager.getInstance() != null) { String messageJson = JsonUtils.pojoToJson(appRecord); @@ -823,6 +735,9 @@ public void stop() { if (jobExecutor != null) { jobExecutor.shutdownNow(); } + if (distributedExecutor != null) { + distributedExecutor.stop(); + } LOG.info("RDF indexing job stopped successfully."); } @@ -837,6 +752,43 @@ protected void validateConfig(Map appConfig) { } private Set getAll() { - return new HashSet<>(Entity.getEntityList()); + return resolveEntityTypes(new HashSet<>(Entity.getEntityList())); + } + + private Set resolveEntityTypes(Set requestedEntities) { + Set entitiesToResolve = requestedEntities; + if (entitiesToResolve == null + || entitiesToResolve.isEmpty() + || entitiesToResolve.contains(ALL)) { + entitiesToResolve = new HashSet<>(Entity.getEntityList()); + } + + Set resolvedEntities = new LinkedHashSet<>(); + List skippedEntities = new ArrayList<>(); + for (String entityType : entitiesToResolve) { + if (entityType == null || entityType.isBlank() || ALL.equals(entityType)) { + continue; + } + if (isIndexableEntityType(entityType)) { + resolvedEntities.add(entityType); + } else { + skippedEntities.add(entityType); + } + } + + if (!skippedEntities.isEmpty()) { + LOG.info("Skipping RDF indexing for non repository-backed entity types: {}", skippedEntities); + } + + return resolvedEntities; + } + + private boolean isIndexableEntityType(String entityType) { + try { + Entity.getEntityRepository(entityType); + return true; + } catch (Exception e) { + return false; + } } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexCoordinator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexCoordinator.java new file mode 100644 index 000000000000..b30b0d9786c2 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexCoordinator.java @@ -0,0 +1,633 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import com.fasterxml.jackson.core.type.TypeReference; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.PartitionStatus; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.ServerIdentityResolver; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexJobDAO.RdfIndexJobRecord; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexPartitionDAO.RdfAggregatedStatsRecord; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexPartitionDAO.RdfEntityStatsRecord; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexPartitionDAO.RdfIndexPartitionRecord; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexPartitionDAO.RdfServerPartitionStatsRecord; + +@Slf4j +public class DistributedRdfIndexCoordinator { + private static final String REINDEX_LOCK_KEY = "RDF_REINDEX_LOCK"; + private static final long LOCK_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(5); + private static final long PARTITION_STALE_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(3); + private static final int MAX_PARTITION_RETRIES = 3; + private static final double IMMEDIATE_CLAIMABLE_PERCENT = 0.50; + private static final long PARTITION_RELEASE_WINDOW_MS = TimeUnit.SECONDS.toMillis(5); + + private final CollectionDAO collectionDAO; + private final RdfPartitionCalculator partitionCalculator; + private final String serverId; + private final AtomicLong lastClaimTimestamp = new AtomicLong(0); + + public DistributedRdfIndexCoordinator(CollectionDAO collectionDAO) { + this(collectionDAO, new RdfPartitionCalculator()); + } + + public DistributedRdfIndexCoordinator( + CollectionDAO collectionDAO, RdfPartitionCalculator partitionCalculator) { + this.collectionDAO = collectionDAO; + this.partitionCalculator = partitionCalculator; + this.serverId = ServerIdentityResolver.getInstance().getServerId(); + } + + public CollectionDAO getCollectionDAO() { + return collectionDAO; + } + + public boolean tryAcquireReindexLock(UUID jobId) { + long now = System.currentTimeMillis(); + return collectionDAO + .rdfReindexLockDAO() + .tryAcquireLock(REINDEX_LOCK_KEY, jobId.toString(), serverId, now, now + LOCK_TIMEOUT_MS); + } + + public boolean transferReindexLock(UUID fromJobId, UUID toJobId) { + long now = System.currentTimeMillis(); + return collectionDAO + .rdfReindexLockDAO() + .transferLock( + REINDEX_LOCK_KEY, + fromJobId.toString(), + toJobId.toString(), + serverId, + now, + now + LOCK_TIMEOUT_MS); + } + + public void refreshReindexLock(UUID jobId) { + long now = System.currentTimeMillis(); + collectionDAO + .rdfReindexLockDAO() + .updateHeartbeat(REINDEX_LOCK_KEY, jobId.toString(), now, now + LOCK_TIMEOUT_MS); + } + + public void releaseReindexLock(UUID jobId) { + collectionDAO.rdfReindexLockDAO().releaseLock(REINDEX_LOCK_KEY, jobId.toString()); + } + + public Optional getJob(UUID jobId) { + RdfIndexJobRecord record = collectionDAO.rdfIndexJobDAO().findById(jobId.toString()); + return Optional.ofNullable(record).map(this::toJob); + } + + public List getRecentJobs(List statuses, int limit) { + List statusNames = statuses.stream().map(Enum::name).toList(); + return collectionDAO.rdfIndexJobDAO().findByStatusesWithLimit(statusNames, limit).stream() + .map(this::toJob) + .toList(); + } + + public Optional getBlockingJob() { + List jobs = + getRecentJobs( + List.of(IndexJobStatus.READY, IndexJobStatus.RUNNING, IndexJobStatus.STOPPING), 1); + return jobs.stream().findFirst(); + } + + public RdfIndexJob createJob( + Set entities, EventPublisherJob jobConfiguration, String createdBy) { + UUID jobId = UUID.randomUUID(); + long now = System.currentTimeMillis(); + + Map entityStats = new HashMap<>(); + long totalRecords = 0; + for (String entityType : entities) { + long count = partitionCalculator.getEntityCount(entityType); + totalRecords += count; + entityStats.put( + entityType, + RdfIndexJob.EntityTypeStats.builder() + .entityType(entityType) + .totalRecords(count) + .processedRecords(0) + .successRecords(0) + .failedRecords(0) + .totalPartitions(0) + .completedPartitions(0) + .failedPartitions(0) + .build()); + } + + RdfIndexJob job = + RdfIndexJob.builder() + .id(jobId) + .status(IndexJobStatus.INITIALIZING) + .jobConfiguration(jobConfiguration) + .totalRecords(totalRecords) + .processedRecords(0) + .successRecords(0) + .failedRecords(0) + .entityStats(entityStats) + .createdBy(createdBy) + .createdAt(now) + .updatedAt(now) + .build(); + + collectionDAO + .rdfIndexJobDAO() + .insert( + jobId.toString(), + job.getStatus().name(), + JsonUtils.pojoToJson(jobConfiguration), + job.getTotalRecords(), + 0, + 0, + 0, + serializeEntityStats(entityStats), + createdBy, + now, + now); + return job; + } + + public RdfIndexJob initializePartitions(UUID jobId) { + RdfIndexJob job = + getJob(jobId).orElseThrow(() -> new IllegalStateException("RDF job not found: " + jobId)); + Set entityTypes = Set.copyOf(job.getJobConfiguration().getEntities()); + List partitions = + partitionCalculator.calculatePartitions(jobId, entityTypes); + long now = System.currentTimeMillis(); + int immediateCount = + Math.max(1, (int) Math.ceil(partitions.size() * IMMEDIATE_CLAIMABLE_PERCENT)); + + for (int i = 0; i < partitions.size(); i++) { + RdfIndexPartition partition = partitions.get(i); + long claimableAt; + if (i < immediateCount) { + claimableAt = now; + } else { + int remainingIndex = i - immediateCount; + int remainingCount = Math.max(1, partitions.size() - immediateCount); + claimableAt = now + (remainingIndex * PARTITION_RELEASE_WINDOW_MS) / remainingCount; + } + insertPartition(partition.withClaimableAt(claimableAt)); + } + + Map entityStats = new HashMap<>(job.getEntityStats()); + for (String entityType : entityTypes) { + int totalPartitions = + (int) + partitions.stream() + .filter(partition -> entityType.equals(partition.getEntityType())) + .count(); + RdfIndexJob.EntityTypeStats existing = entityStats.get(entityType); + if (existing != null) { + entityStats.put(entityType, existing.toBuilder().totalPartitions(totalPartitions).build()); + } + } + + long totalRecords = partitions.stream().mapToLong(RdfIndexPartition::getEstimatedCount).sum(); + RdfIndexJob updated = + job.toBuilder() + .status(IndexJobStatus.READY) + .totalRecords(totalRecords) + .entityStats(entityStats) + .updatedAt(System.currentTimeMillis()) + .build(); + updateJob(updated); + return updated; + } + + public RdfIndexPartition claimNextPartition(UUID jobId) { + long claimAt = nextClaimTimestamp(); + int updated = + collectionDAO + .rdfIndexPartitionDAO() + .claimNextPartitionAtomic(jobId.toString(), serverId, claimAt); + if (updated <= 0) { + return null; + } + + RdfIndexPartitionRecord record = + collectionDAO + .rdfIndexPartitionDAO() + .findLatestClaimedPartition(jobId.toString(), serverId, claimAt); + if (record == null) { + LOG.warn( + "Claimed RDF partition for job {} but could not retrieve the record; it may require stale recovery", + jobId); + return null; + } + + return toPartition(record); + } + + public void updatePartitionProgress(RdfIndexPartition partition) { + collectionDAO + .rdfIndexPartitionDAO() + .updateProgress( + partition.getId().toString(), + partition.getCursor(), + partition.getProcessedCount(), + partition.getSuccessCount(), + partition.getFailedCount(), + System.currentTimeMillis()); + } + + public void completePartition( + UUID partitionId, long cursor, long processedCount, long successCount, long failedCount) { + RdfIndexPartition partition = getPartition(partitionId); + long now = System.currentTimeMillis(); + collectionDAO + .rdfIndexPartitionDAO() + .update( + partitionId.toString(), + PartitionStatus.COMPLETED.name(), + cursor, + processedCount, + successCount, + failedCount, + partition.getAssignedServer(), + partition.getClaimedAt(), + partition.getStartedAt(), + now, + now, + null, + partition.getRetryCount()); + incrementServerStats(partition, processedCount, successCount, failedCount, 1, 0); + refreshAggregatedJob(jobIdFrom(partition)); + } + + public void failPartition( + UUID partitionId, + long cursor, + long processedCount, + long successCount, + long failedCount, + String errorMessage) { + RdfIndexPartition partition = getPartition(partitionId); + long now = System.currentTimeMillis(); + collectionDAO + .rdfIndexPartitionDAO() + .update( + partitionId.toString(), + PartitionStatus.FAILED.name(), + cursor, + processedCount, + successCount, + failedCount, + partition.getAssignedServer(), + partition.getClaimedAt(), + partition.getStartedAt(), + now, + now, + errorMessage, + partition.getRetryCount() + 1); + incrementServerStats(partition, processedCount, successCount, failedCount, 0, 1); + refreshAggregatedJob(jobIdFrom(partition)); + } + + public int reclaimStalePartitions(UUID jobId) { + long staleThreshold = System.currentTimeMillis() - PARTITION_STALE_TIMEOUT_MS; + int reclaimed = + collectionDAO + .rdfIndexPartitionDAO() + .reclaimStalePartitionsForRetry( + jobId.toString(), staleThreshold, MAX_PARTITION_RETRIES); + int failed = + collectionDAO + .rdfIndexPartitionDAO() + .failStalePartitionsExceedingRetries( + jobId.toString(), + staleThreshold, + MAX_PARTITION_RETRIES, + System.currentTimeMillis()); + if (reclaimed > 0 || failed > 0) { + LOG.info( + "Recovered RDF job {} partitions: reclaimed={}, failed={}", jobId, reclaimed, failed); + refreshAggregatedJob(jobId); + } + return reclaimed + failed; + } + + public void cancelPendingPartitions(UUID jobId) { + collectionDAO.rdfIndexPartitionDAO().cancelPendingPartitions(jobId.toString()); + refreshAggregatedJob(jobId); + } + + public void releaseServerPartitions(UUID jobId, String serverId, boolean stopJob, String reason) { + long now = System.currentTimeMillis(); + collectionDAO + .rdfIndexPartitionDAO() + .releaseProcessingPartitions( + jobId.toString(), + serverId, + stopJob ? PartitionStatus.CANCELLED.name() : PartitionStatus.PENDING.name(), + reason, + now, + stopJob ? now : null); + refreshAggregatedJob(jobId); + } + + public void updateJobStatus(UUID jobId, IndexJobStatus status, String errorMessage) { + RdfIndexJob job = + getJob(jobId).orElseThrow(() -> new IllegalStateException("RDF job not found: " + jobId)); + long now = System.currentTimeMillis(); + Long startedAt = job.getStartedAt(); + Long completedAt = job.getCompletedAt(); + + if (status == IndexJobStatus.RUNNING && startedAt == null) { + startedAt = now; + } + if (status == IndexJobStatus.STOPPED + || status == IndexJobStatus.COMPLETED + || status == IndexJobStatus.COMPLETED_WITH_ERRORS + || status == IndexJobStatus.FAILED) { + completedAt = completedAt != null ? completedAt : now; + } + + collectionDAO + .rdfIndexJobDAO() + .update( + jobId.toString(), + status.name(), + job.getProcessedRecords(), + job.getSuccessRecords(), + job.getFailedRecords(), + serializeEntityStats(job.getEntityStats()), + startedAt, + completedAt, + now, + errorMessage); + } + + public RdfIndexJob getJobWithAggregatedStats(UUID jobId) { + return refreshAggregatedJob(jobId); + } + + public boolean hasClaimableWork(UUID jobId) { + RdfIndexJob job = refreshAggregatedJob(jobId); + if (job == null || job.isTerminal()) { + return false; + } + + String id = jobId.toString(); + + return collectionDAO.rdfIndexPartitionDAO().countPendingPartitions(id) > 0 + || collectionDAO.rdfIndexPartitionDAO().countInFlightPartitions(id) > 0; + } + + public void performStartupRecovery() { + for (RdfIndexJob job : + getRecentJobs( + List.of(IndexJobStatus.READY, IndexJobStatus.RUNNING, IndexJobStatus.STOPPING), 20)) { + reclaimStalePartitions(job.getId()); + refreshAggregatedJob(job.getId()); + } + } + + private RdfIndexJob refreshAggregatedJob(UUID jobId) { + RdfIndexJob existing = getJob(jobId).orElse(null); + if (existing == null) { + return null; + } + + RdfAggregatedStatsRecord aggregate = + collectionDAO.rdfIndexPartitionDAO().getAggregatedStats(jobId.toString()); + Map entityStats = + collectionDAO.rdfIndexPartitionDAO().getEntityStats(jobId.toString()).stream() + .collect( + Collectors.toMap( + RdfEntityStatsRecord::entityType, + record -> + RdfIndexJob.EntityTypeStats.builder() + .entityType(record.entityType()) + .totalRecords(record.totalRecords()) + .processedRecords(record.processedRecords()) + .successRecords(record.successRecords()) + .failedRecords(record.failedRecords()) + .totalPartitions(record.totalPartitions()) + .completedPartitions(record.completedPartitions()) + .failedPartitions(record.failedPartitions()) + .build(), + (left, right) -> right, + HashMap::new)); + Map serverStats = + collectionDAO.rdfIndexPartitionDAO().getServerStats(jobId.toString()).stream() + .collect( + Collectors.toMap( + RdfServerPartitionStatsRecord::serverId, + record -> + RdfIndexJob.ServerStats.builder() + .serverId(record.serverId()) + .processedRecords(record.processedRecords()) + .successRecords(record.successRecords()) + .failedRecords(record.failedRecords()) + .totalPartitions(record.totalPartitions()) + .completedPartitions(record.completedPartitions()) + .processingPartitions(record.processingPartitions()) + .build(), + (left, right) -> right, + HashMap::new)); + + IndexJobStatus status = existing.getStatus(); + String errorMessage = existing.getErrorMessage(); + if (aggregate.pendingPartitions() == 0 && aggregate.processingPartitions() == 0) { + if (status == IndexJobStatus.STOPPING) { + status = IndexJobStatus.STOPPED; + } else if (aggregate.failedPartitions() > 0 || aggregate.failedRecords() > 0) { + status = IndexJobStatus.COMPLETED_WITH_ERRORS; + } else if (status == IndexJobStatus.READY || status == IndexJobStatus.RUNNING) { + status = IndexJobStatus.COMPLETED; + } + } else if (status == IndexJobStatus.READY) { + status = IndexJobStatus.RUNNING; + } + + Long completedAt = existing.getCompletedAt(); + if (completedAt == null + && (status == IndexJobStatus.COMPLETED + || status == IndexJobStatus.COMPLETED_WITH_ERRORS + || status == IndexJobStatus.FAILED + || status == IndexJobStatus.STOPPED)) { + completedAt = System.currentTimeMillis(); + } + + RdfIndexJob refreshed = + existing.toBuilder() + .status(status) + .processedRecords(aggregate.processedRecords()) + .successRecords(aggregate.successRecords()) + .failedRecords(aggregate.failedRecords()) + .entityStats(entityStats) + .serverStats(serverStats) + .updatedAt(System.currentTimeMillis()) + .errorMessage(errorMessage) + .completedAt(completedAt) + .build(); + + updateJob(refreshed); + return refreshed; + } + + private void incrementServerStats( + RdfIndexPartition partition, + long processedCount, + long successCount, + long failedCount, + int partitionsCompleted, + int partitionsFailed) { + String assignedServer = + partition.getAssignedServer() != null ? partition.getAssignedServer() : serverId; + collectionDAO + .rdfIndexServerStatsDAO() + .incrementStats( + UUID.randomUUID().toString(), + partition.getJobId().toString(), + assignedServer, + partition.getEntityType(), + processedCount, + successCount, + failedCount, + partitionsCompleted, + partitionsFailed, + System.currentTimeMillis()); + } + + private long nextClaimTimestamp() { + return lastClaimTimestamp.updateAndGet( + previous -> { + long now = System.currentTimeMillis(); + return now > previous ? now : previous + 1; + }); + } + + private void insertPartition(RdfIndexPartition partition) { + collectionDAO + .rdfIndexPartitionDAO() + .insert( + partition.getId().toString(), + partition.getJobId().toString(), + partition.getEntityType(), + partition.getPartitionIndex(), + partition.getRangeStart(), + partition.getRangeEnd(), + partition.getEstimatedCount(), + partition.getWorkUnits(), + partition.getPriority(), + partition.getStatus().name(), + partition.getCursor(), + partition.getClaimableAt()); + } + + private void updateJob(RdfIndexJob job) { + collectionDAO + .rdfIndexJobDAO() + .update( + job.getId().toString(), + job.getStatus().name(), + job.getProcessedRecords(), + job.getSuccessRecords(), + job.getFailedRecords(), + serializeEntityStats(job.getEntityStats()), + job.getStartedAt(), + job.getCompletedAt(), + job.getUpdatedAt(), + job.getErrorMessage()); + } + + private RdfIndexPartition getPartition(UUID partitionId) { + RdfIndexPartitionRecord record = + collectionDAO.rdfIndexPartitionDAO().findById(partitionId.toString()); + if (record == null) { + throw new IllegalStateException("RDF partition not found: " + partitionId); + } + return toPartition(record); + } + + private UUID jobIdFrom(RdfIndexPartition partition) { + return partition.getJobId(); + } + + private String serializeEntityStats(Map entityStats) { + return JsonUtils.pojoToJson(entityStats != null ? entityStats : Map.of()); + } + + private RdfIndexJob toJob(RdfIndexJobRecord record) { + Map entityStats = + record.stats() != null && !record.stats().isBlank() + ? JsonUtils.readValue( + record.stats(), new TypeReference>() {}) + : new HashMap<>(); + + EventPublisherJob jobConfiguration = + record.jobConfiguration() != null + ? JsonUtils.readValue(record.jobConfiguration(), EventPublisherJob.class) + : new EventPublisherJob(); + + return RdfIndexJob.builder() + .id(UUID.fromString(record.id())) + .status(IndexJobStatus.valueOf(record.status())) + .jobConfiguration(jobConfiguration) + .totalRecords(record.totalRecords()) + .processedRecords(record.processedRecords()) + .successRecords(record.successRecords()) + .failedRecords(record.failedRecords()) + .entityStats(entityStats) + .createdBy(record.createdBy()) + .createdAt(record.createdAt()) + .startedAt(record.startedAt()) + .completedAt(record.completedAt()) + .updatedAt(record.updatedAt()) + .errorMessage(record.errorMessage()) + .build(); + } + + private RdfIndexPartition toPartition(RdfIndexPartitionRecord record) { + return RdfIndexPartition.builder() + .id(UUID.fromString(record.id())) + .jobId(UUID.fromString(record.jobId())) + .entityType(record.entityType()) + .partitionIndex(record.partitionIndex()) + .rangeStart(record.rangeStart()) + .rangeEnd(record.rangeEnd()) + .estimatedCount(record.estimatedCount()) + .workUnits(record.workUnits()) + .priority(record.priority()) + .status(PartitionStatus.valueOf(record.status())) + .cursor(record.cursor()) + .processedCount(record.processedCount()) + .successCount(record.successCount()) + .failedCount(record.failedCount()) + .assignedServer(record.assignedServer()) + .claimedAt(record.claimedAt()) + .startedAt(record.startedAt()) + .completedAt(record.completedAt()) + .lastUpdateAt(record.lastUpdateAt()) + .lastError(record.lastError()) + .retryCount(record.retryCount()) + .claimableAt(record.claimableAt()) + .build(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexExecutor.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexExecutor.java new file mode 100644 index 000000000000..5f0b6a515970 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexExecutor.java @@ -0,0 +1,350 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.service.apps.bundles.rdf.RdfBatchProcessor; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.ServerIdentityResolver; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.rdf.RdfRepository; + +@Slf4j +public class DistributedRdfIndexExecutor { + private static final Set COORDINATED_JOBS = ConcurrentHashMap.newKeySet(); + private static final long LOCK_REFRESH_INTERVAL_MS = TimeUnit.MINUTES.toMillis(1); + private static final long STALE_CHECK_INTERVAL_MS = TimeUnit.SECONDS.toMillis(30); + private static final long CLAIM_RETRY_SLEEP_MS = 1000; + private static final long SHUTDOWN_TIMEOUT_SECONDS = 30; + + private final CollectionDAO collectionDAO; + private final DistributedRdfIndexCoordinator coordinator; + private final String serverId; + private final AtomicBoolean stopped = new AtomicBoolean(false); + private final AtomicBoolean localExecutionCleaned = new AtomicBoolean(true); + private final List activeWorkers = new CopyOnWriteArrayList<>(); + + @Getter private volatile RdfIndexJob currentJob; + private volatile ExecutorService workerExecutor; + private volatile Thread lockRefreshThread; + private volatile Thread staleReclaimerThread; + private volatile boolean coordinatorOwnedJob; + + public DistributedRdfIndexExecutor(CollectionDAO collectionDAO, int partitionSize) { + this( + collectionDAO, + new DistributedRdfIndexCoordinator( + collectionDAO, new RdfPartitionCalculator(partitionSize)), + ServerIdentityResolver.getInstance().getServerId()); + } + + DistributedRdfIndexExecutor( + CollectionDAO collectionDAO, DistributedRdfIndexCoordinator coordinator, String serverId) { + this.collectionDAO = collectionDAO; + this.coordinator = coordinator; + this.serverId = serverId; + } + + public static boolean isCoordinatingJob(UUID jobId) { + return COORDINATED_JOBS.contains(jobId); + } + + public void performStartupRecovery() { + coordinator.performStartupRecovery(); + } + + public RdfIndexJob createJob( + Set entities, EventPublisherJob jobConfiguration, String createdBy) { + Optional blockingJob = coordinator.getBlockingJob(); + if (blockingJob.isPresent()) { + throw new IllegalStateException( + "Another RDF reindex job is already active: " + blockingJob.get().getId()); + } + + UUID tempJobId = UUID.randomUUID(); + if (!coordinator.tryAcquireReindexLock(tempJobId)) { + throw new IllegalStateException("Failed to acquire RDF reindex lock"); + } + + try { + currentJob = coordinator.createJob(entities, jobConfiguration, createdBy); + currentJob = coordinator.initializePartitions(currentJob.getId()); + if (!coordinator.transferReindexLock(tempJobId, currentJob.getId())) { + throw new IllegalStateException("Failed to transfer RDF reindex lock to job"); + } + coordinatorOwnedJob = true; + return currentJob; + } catch (Exception e) { + coordinator.releaseReindexLock(tempJobId); + throw e; + } + } + + public void execute(EventPublisherJob jobConfiguration) throws InterruptedException { + if (currentJob == null) { + throw new IllegalStateException("RDF distributed job must be created before execution"); + } + + stopped.set(false); + localExecutionCleaned.set(false); + COORDINATED_JOBS.add(currentJob.getId()); + coordinator.updateJobStatus(currentJob.getId(), IndexJobStatus.RUNNING, null); + currentJob = coordinator.getJobWithAggregatedStats(currentJob.getId()); + if (currentJob == null) { + throw new IllegalStateException("Failed to load RDF distributed job state"); + } + + try { + startCoordinatorThreads(); + runWorkers(jobConfiguration, true); + finalizeCoordinatorJob(); + } finally { + cleanupCoordinatorExecution(); + } + } + + public void joinJob(RdfIndexJob job, EventPublisherJob jobConfiguration) + throws InterruptedException { + currentJob = job; + coordinatorOwnedJob = false; + stopped.set(false); + localExecutionCleaned.set(false); + runWorkers(jobConfiguration, false); + } + + public RdfIndexJob getJobWithFreshStats() { + if (currentJob == null) { + return null; + } + currentJob = coordinator.getJobWithAggregatedStats(currentJob.getId()); + return currentJob; + } + + public void stop() { + stopped.set(true); + + if (currentJob != null) { + if (coordinatorOwnedJob) { + coordinator.updateJobStatus(currentJob.getId(), IndexJobStatus.STOPPING, null); + coordinator.cancelPendingPartitions(currentJob.getId()); + coordinator.releaseServerPartitions(currentJob.getId(), serverId, true, "Stopped by user"); + } else { + coordinator.releaseServerPartitions( + currentJob.getId(), serverId, false, "Worker server stopped participating"); + } + } + + for (RdfPartitionWorker worker : activeWorkers) { + worker.stop(); + } + + cleanupLocalExecution(); + } + + private void runWorkers(EventPublisherJob jobConfiguration, boolean coordinatorMode) + throws InterruptedException { + activeWorkers.clear(); + + int workerCount = + Math.max( + 1, + Math.min( + jobConfiguration.getConsumerThreads() != null + ? jobConfiguration.getConsumerThreads() + : Runtime.getRuntime().availableProcessors(), + Runtime.getRuntime().availableProcessors() * 2)); + int batchSize = jobConfiguration.getBatchSize() != null ? jobConfiguration.getBatchSize() : 100; + RdfBatchProcessor batchProcessor = + new RdfBatchProcessor(collectionDAO, RdfRepository.getInstance()); + + workerExecutor = + Executors.newFixedThreadPool( + workerCount, + Thread.ofPlatform() + .name( + coordinatorMode + ? "rdf-distributed-coordinator-" + : "rdf-distributed-participant-", + 0) + .factory()); + try { + for (int i = 0; i < workerCount; i++) { + RdfPartitionWorker worker = new RdfPartitionWorker(coordinator, batchProcessor, batchSize); + activeWorkers.add(worker); + workerExecutor.submit(() -> workerLoop(worker)); + } + + workerExecutor.shutdown(); + workerExecutor.awaitTermination(Long.MAX_VALUE, TimeUnit.MILLISECONDS); + } finally { + activeWorkers.clear(); + if (workerExecutor != null && !workerExecutor.isShutdown()) { + shutdownWorkerExecutor(); + } + } + } + + private void workerLoop(RdfPartitionWorker worker) { + while (!stopped.get() && !Thread.currentThread().isInterrupted()) { + RdfIndexJob latestJob = getJobWithFreshStats(); + if (latestJob == null + || latestJob.isTerminal() + || latestJob.getStatus() == IndexJobStatus.STOPPING) { + return; + } + + RdfIndexPartition partition = coordinator.claimNextPartition(latestJob.getId()); + if (partition == null) { + try { + TimeUnit.MILLISECONDS.sleep(CLAIM_RETRY_SLEEP_MS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + continue; + } + + worker.processPartition(partition); + } + } + + private void finalizeCoordinatorJob() { + currentJob = coordinator.getJobWithAggregatedStats(currentJob.getId()); + if (currentJob == null) { + return; + } + + if (stopped.get()) { + coordinator.updateJobStatus(currentJob.getId(), IndexJobStatus.STOPPED, null); + } else if (!currentJob.isTerminal()) { + IndexJobStatus terminalStatus = + currentJob.getFailedRecords() > 0 + ? IndexJobStatus.COMPLETED_WITH_ERRORS + : IndexJobStatus.COMPLETED; + coordinator.updateJobStatus(currentJob.getId(), terminalStatus, currentJob.getErrorMessage()); + } + + currentJob = coordinator.getJobWithAggregatedStats(currentJob.getId()); + } + + private void startCoordinatorThreads() { + lockRefreshThread = + Thread.ofVirtual() + .name("rdf-lock-refresh-" + currentJob.getId().toString().substring(0, 8)) + .start( + () -> { + while (!stopped.get() && !Thread.currentThread().isInterrupted()) { + try { + coordinator.refreshReindexLock(currentJob.getId()); + TimeUnit.MILLISECONDS.sleep(LOCK_REFRESH_INTERVAL_MS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } catch (Exception e) { + LOG.warn("Failed to refresh RDF reindex lock for {}", currentJob.getId(), e); + } + } + }); + + staleReclaimerThread = + Thread.ofVirtual() + .name("rdf-stale-reclaimer-" + currentJob.getId().toString().substring(0, 8)) + .start( + () -> { + while (!stopped.get() && !Thread.currentThread().isInterrupted()) { + try { + coordinator.reclaimStalePartitions(currentJob.getId()); + TimeUnit.MILLISECONDS.sleep(STALE_CHECK_INTERVAL_MS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } catch (Exception e) { + LOG.warn( + "Failed to reclaim stale RDF partitions for {}", currentJob.getId(), e); + } + } + }); + } + + private void shutdownWorkerExecutor() { + if (workerExecutor == null || workerExecutor.isShutdown()) { + return; + } + + workerExecutor.shutdownNow(); + try { + if (!workerExecutor.awaitTermination(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + LOG.warn("Timed out waiting for RDF distributed workers to stop"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + workerExecutor = null; + } + } + + private void interruptThread(Thread thread) { + if (thread == null) { + return; + } + thread.interrupt(); + try { + thread.join(5_000); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + private void cleanupLocalExecution() { + if (!localExecutionCleaned.compareAndSet(false, true)) { + return; + } + + shutdownWorkerExecutor(); + interruptThread(lockRefreshThread); + interruptThread(staleReclaimerThread); + lockRefreshThread = null; + staleReclaimerThread = null; + activeWorkers.clear(); + } + + private void cleanupCoordinatorExecution() { + UUID jobId = currentJob != null ? currentJob.getId() : null; + + cleanupLocalExecution(); + + if (jobId != null && coordinatorOwnedJob) { + try { + coordinator.releaseReindexLock(jobId); + } catch (Exception e) { + LOG.warn("Failed to release RDF reindex lock for {}", jobId, e); + } + } + if (jobId != null) { + COORDINATED_JOBS.remove(jobId); + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobParticipant.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobParticipant.java new file mode 100644 index 000000000000..be848c537452 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobParticipant.java @@ -0,0 +1,151 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import io.dropwizard.lifecycle.Managed; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.ServerIdentityResolver; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.rdf.RdfRepository; + +@Slf4j +public class RdfDistributedJobParticipant implements Managed { + private static final long POLL_INTERVAL_MS = TimeUnit.SECONDS.toMillis(15); + + private final CollectionDAO collectionDAO; + private final String serverId; + private final DistributedRdfIndexCoordinator coordinator; + private final AtomicBoolean running = new AtomicBoolean(false); + private final AtomicBoolean participating = new AtomicBoolean(false); + + @Getter private UUID currentJobId; + + private volatile Thread pollThread; + private volatile Thread participantThread; + + public RdfDistributedJobParticipant(CollectionDAO collectionDAO) { + this.collectionDAO = collectionDAO; + this.serverId = ServerIdentityResolver.getInstance().getServerId(); + this.coordinator = new DistributedRdfIndexCoordinator(collectionDAO); + } + + @Override + public void start() { + RdfRepository rdfRepository = RdfRepository.getInstanceOrNull(); + if (rdfRepository == null || !rdfRepository.isEnabled()) { + LOG.info( + "Skipping RDF distributed participant registration because RDF is not initialized or disabled"); + return; + } + + if (running.compareAndSet(false, true)) { + pollThread = + Thread.ofVirtual() + .name("rdf-distributed-participant-poll") + .start( + () -> { + while (running.get() && !Thread.currentThread().isInterrupted()) { + try { + pollForJobs(); + TimeUnit.MILLISECONDS.sleep(POLL_INTERVAL_MS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } catch (Exception e) { + LOG.warn("Error polling for distributed RDF jobs", e); + } + } + }); + LOG.info("Started RDF distributed job participant on server {}", serverId); + } + } + + @Override + public void stop() { + if (running.compareAndSet(true, false)) { + interruptThread(pollThread); + interruptThread(participantThread); + LOG.info("Stopped RDF distributed job participant on server {}", serverId); + } + } + + private void pollForJobs() { + if (participating.get()) { + return; + } + + List activeJobs = + coordinator.getRecentJobs(List.of(IndexJobStatus.RUNNING, IndexJobStatus.STOPPING), 10); + for (RdfIndexJob job : activeJobs) { + if (job.isTerminal() + || job.getStatus() != IndexJobStatus.RUNNING + || DistributedRdfIndexExecutor.isCoordinatingJob(job.getId()) + || !coordinator.hasClaimableWork(job.getId())) { + continue; + } + + joinJob(job); + return; + } + } + + private void joinJob(RdfIndexJob job) { + if (!participating.compareAndSet(false, true)) { + return; + } + + currentJobId = job.getId(); + participantThread = + Thread.ofVirtual() + .name("rdf-distributed-participant-" + job.getId().toString().substring(0, 8)) + .start( + () -> { + try { + int partitionSize = + job.getJobConfiguration().getPartitionSize() != null + ? job.getJobConfiguration().getPartitionSize() + : 10000; + DistributedRdfIndexExecutor executor = + new DistributedRdfIndexExecutor(collectionDAO, partitionSize); + executor.joinJob(job, job.getJobConfiguration()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } catch (Exception e) { + LOG.warn("Failed to participate in RDF job {}", job.getId(), e); + } finally { + currentJobId = null; + participating.set(false); + participantThread = null; + } + }); + } + + private void interruptThread(Thread thread) { + if (thread == null) { + return; + } + thread.interrupt(); + try { + thread.join(5_000); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobStatsAggregator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobStatsAggregator.java new file mode 100644 index 000000000000..f2dfa9f9cf73 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobStatsAggregator.java @@ -0,0 +1,58 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import org.openmetadata.schema.system.EntityStats; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; + +public class RdfDistributedJobStatsAggregator { + public Stats toStats(RdfIndexJob job) { + Stats stats = new Stats(); + stats.setEntityStats(new EntityStats()); + + StepStats jobStats = + new StepStats() + .withTotalRecords(safeToInt(job.getTotalRecords())) + .withSuccessRecords(safeToInt(job.getSuccessRecords())) + .withFailedRecords(safeToInt(job.getFailedRecords())); + stats.setJobStats(jobStats); + + if (job.getEntityStats() != null) { + job.getEntityStats() + .forEach( + (entityType, entityStats) -> + stats + .getEntityStats() + .setAdditionalProperty( + entityType, + new StepStats() + .withTotalRecords(safeToInt(entityStats.getTotalRecords())) + .withSuccessRecords(safeToInt(entityStats.getSuccessRecords())) + .withFailedRecords(safeToInt(entityStats.getFailedRecords())))); + } + + return stats; + } + + private int safeToInt(long value) { + if (value > Integer.MAX_VALUE) { + return Integer.MAX_VALUE; + } + if (value < Integer.MIN_VALUE) { + return Integer.MIN_VALUE; + } + return (int) value; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfIndexJob.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfIndexJob.java new file mode 100644 index 000000000000..0f09b855d49c --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfIndexJob.java @@ -0,0 +1,83 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import java.util.Map; +import java.util.UUID; +import lombok.AllArgsConstructor; +import lombok.Builder; +import lombok.Data; +import lombok.NoArgsConstructor; +import lombok.With; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; + +@Data +@Builder(toBuilder = true) +@NoArgsConstructor +@AllArgsConstructor +@With +public class RdfIndexJob { + private UUID id; + private IndexJobStatus status; + private EventPublisherJob jobConfiguration; + private long totalRecords; + private long processedRecords; + private long successRecords; + private long failedRecords; + private Map entityStats; + private Map serverStats; + private String createdBy; + private long createdAt; + private Long startedAt; + private Long completedAt; + private long updatedAt; + private String errorMessage; + + public boolean isTerminal() { + return status == IndexJobStatus.COMPLETED + || status == IndexJobStatus.COMPLETED_WITH_ERRORS + || status == IndexJobStatus.FAILED + || status == IndexJobStatus.STOPPED; + } + + @Data + @Builder(toBuilder = true) + @NoArgsConstructor + @AllArgsConstructor + public static class EntityTypeStats { + private String entityType; + private long totalRecords; + private long processedRecords; + private long successRecords; + private long failedRecords; + private int totalPartitions; + private int completedPartitions; + private int failedPartitions; + } + + @Data + @Builder(toBuilder = true) + @NoArgsConstructor + @AllArgsConstructor + public static class ServerStats { + private String serverId; + private long processedRecords; + private long successRecords; + private long failedRecords; + private int totalPartitions; + private int completedPartitions; + private int processingPartitions; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfIndexPartition.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfIndexPartition.java new file mode 100644 index 000000000000..4cc041845718 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfIndexPartition.java @@ -0,0 +1,48 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import java.util.UUID; +import lombok.Builder; +import lombok.Data; +import lombok.With; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.PartitionStatus; + +@Data +@Builder(toBuilder = true) +@With +public class RdfIndexPartition { + private UUID id; + private UUID jobId; + private String entityType; + private int partitionIndex; + private long rangeStart; + private long rangeEnd; + private long estimatedCount; + private long workUnits; + private int priority; + private PartitionStatus status; + private long cursor; + private long processedCount; + private long successCount; + private long failedCount; + private String assignedServer; + private Long claimedAt; + private Long startedAt; + private Long completedAt; + private Long lastUpdateAt; + private String lastError; + private int retryCount; + private long claimableAt; +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionCalculator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionCalculator.java new file mode 100644 index 000000000000..9025f4cc6b51 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionCalculator.java @@ -0,0 +1,114 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.searchIndex.EntityPriority; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.PartitionStatus; + +@Slf4j +public class RdfPartitionCalculator { + + private static final int DEFAULT_PARTITION_SIZE = 10000; + private static final int MIN_PARTITION_SIZE = 1000; + private static final int MAX_PARTITION_SIZE = 50000; + + private static final Map ENTITY_COMPLEXITY_FACTORS = + Map.of( + "table", 1.5, + "dashboard", 1.3, + "pipeline", 1.2, + "mlmodel", 1.3, + "glossaryTerm", 1.1); + + private final int partitionSize; + + public RdfPartitionCalculator() { + this(DEFAULT_PARTITION_SIZE); + } + + public RdfPartitionCalculator(int partitionSize) { + this.partitionSize = Math.clamp(partitionSize, MIN_PARTITION_SIZE, MAX_PARTITION_SIZE); + } + + public List calculatePartitions(UUID jobId, Set entityTypes) { + List partitions = new ArrayList<>(); + for (String entityType : entityTypes) { + partitions.addAll(calculatePartitionsForEntity(jobId, entityType)); + } + return partitions; + } + + public List calculatePartitionsForEntity(UUID jobId, String entityType) { + long totalCount = getEntityCount(entityType); + if (totalCount <= 0) { + return List.of(); + } + + double complexityFactor = ENTITY_COMPLEXITY_FACTORS.getOrDefault(entityType, 1.0); + long adjustedPartitionSize = + Math.max(MIN_PARTITION_SIZE, (long) (partitionSize / complexityFactor)); + int priority = EntityPriority.getNumericPriority(entityType); + long numPartitions = (totalCount + adjustedPartitionSize - 1) / adjustedPartitionSize; + + List partitions = new ArrayList<>(); + for (int index = 0; index < numPartitions; index++) { + long rangeStart = index * adjustedPartitionSize; + long rangeEnd = Math.min(rangeStart + adjustedPartitionSize, totalCount); + long estimatedCount = rangeEnd - rangeStart; + partitions.add( + RdfIndexPartition.builder() + .id(UUID.randomUUID()) + .jobId(jobId) + .entityType(entityType) + .partitionIndex(index) + .rangeStart(rangeStart) + .rangeEnd(rangeEnd) + .estimatedCount(estimatedCount) + .workUnits((long) (estimatedCount * complexityFactor)) + .priority(priority) + .status(PartitionStatus.PENDING) + .cursor(rangeStart) + .processedCount(0) + .successCount(0) + .failedCount(0) + .retryCount(0) + .claimableAt(0) + .build()); + } + + LOG.info( + "Calculated {} RDF partitions for {} (totalRecords={}, partitionSize={})", + partitions.size(), + entityType, + totalCount, + adjustedPartitionSize); + return partitions; + } + + public long getEntityCount(String entityType) { + try { + return Entity.getEntityRepository(entityType).getDao().listTotalCount(); + } catch (Exception e) { + LOG.warn("Failed to fetch entity count for {}", entityType, e); + return 0; + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionWorker.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionWorker.java new file mode 100644 index 000000000000..a7130262a6a1 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionWorker.java @@ -0,0 +1,153 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; + +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.type.Include; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.rdf.RdfBatchProcessor; +import org.openmetadata.service.exception.SearchIndexException; +import org.openmetadata.service.jdbi3.ListFilter; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntitiesSource; + +@Slf4j +public class RdfPartitionWorker { + private static final long MAX_CURSOR_INITIALIZATION_OFFSET = (long) Integer.MAX_VALUE + 1L; + private static final int PROGRESS_UPDATE_INTERVAL = 100; + + private final DistributedRdfIndexCoordinator coordinator; + private final RdfBatchProcessor batchProcessor; + private final int batchSize; + private final AtomicBoolean stopped = new AtomicBoolean(false); + + public RdfPartitionWorker( + DistributedRdfIndexCoordinator coordinator, RdfBatchProcessor batchProcessor, int batchSize) { + this.coordinator = coordinator; + this.batchProcessor = batchProcessor; + this.batchSize = batchSize; + } + + public PartitionResult processPartition(RdfIndexPartition partition) { + String entityType = partition.getEntityType(); + long currentOffset = Math.max(partition.getCursor(), partition.getRangeStart()); + long processedCount = partition.getProcessedCount(); + long successCount = partition.getSuccessCount(); + long failedCount = partition.getFailedCount(); + + try { + String keysetCursor = initializeKeysetCursor(entityType, currentOffset); + while (currentOffset < partition.getRangeEnd() + && !stopped.get() + && !Thread.currentThread().isInterrupted()) { + int currentBatchSize = (int) Math.min(batchSize, partition.getRangeEnd() - currentOffset); + ResultList resultList = + readEntitiesKeyset(entityType, keysetCursor, currentBatchSize); + + if (resultList == null || listOrEmpty(resultList.getData()).isEmpty()) { + break; + } + + RdfBatchProcessor.BatchProcessingResult batchResult = + batchProcessor.processEntities(entityType, resultList.getData(), stopped::get); + int readerErrors = listOrEmpty(resultList.getErrors()).size(); + long batchProcessed = resultList.getData().size() + readerErrors; + + processedCount += batchProcessed; + successCount += batchResult.successCount(); + failedCount += batchResult.failedCount() + readerErrors; + currentOffset += batchProcessed; + + if (processedCount % PROGRESS_UPDATE_INTERVAL < batchProcessed) { + coordinator.updatePartitionProgress( + partition.toBuilder() + .cursor(currentOffset) + .processedCount(processedCount) + .successCount(successCount) + .failedCount(failedCount) + .build()); + } + + keysetCursor = resultList.getPaging() != null ? resultList.getPaging().getAfter() : null; + if (keysetCursor == null && currentOffset < partition.getRangeEnd()) { + keysetCursor = initializeKeysetCursor(entityType, currentOffset); + if (keysetCursor == null) { + break; + } + } + } + + if (stopped.get() || Thread.currentThread().isInterrupted()) { + return new PartitionResult(processedCount, successCount, failedCount, true, null); + } + + coordinator.completePartition( + partition.getId(), currentOffset, processedCount, successCount, failedCount); + return new PartitionResult(processedCount, successCount, failedCount, false, null); + } catch (Exception e) { + LOG.error("Failed to process RDF partition {}", partition.getId(), e); + coordinator.failPartition( + partition.getId(), + currentOffset, + processedCount, + successCount, + failedCount, + e.getMessage()); + return new PartitionResult(processedCount, successCount, failedCount, false, e.getMessage()); + } + } + + public void stop() { + stopped.set(true); + } + + private ResultList readEntitiesKeyset( + String entityType, String keysetCursor, int limit) throws SearchIndexException { + PaginatedEntitiesSource source = + new PaginatedEntitiesSource(entityType, limit, List.of("*"), 0); + return source.readNextKeyset(keysetCursor); + } + + private String initializeKeysetCursor(String entityType, long offset) { + if (offset <= 0) { + return null; + } + int cursorOffset = toCursorOffset(entityType, offset); + return Entity.getEntityRepository(entityType) + .getCursorAtOffset(new ListFilter(Include.ALL), cursorOffset); + } + + private int toCursorOffset(String entityType, long offset) { + long cursorOffset = offset - 1L; + if (cursorOffset > Integer.MAX_VALUE) { + throw new IllegalArgumentException( + String.format( + "Keyset cursor initialization for entityType %s does not support offsets above %d", + entityType, MAX_CURSOR_INITIALIZATION_OFFSET)); + } + return Math.toIntExact(cursorOffset); + } + + public record PartitionResult( + long processedCount, + long successCount, + long failedCount, + boolean stopped, + String errorMessage) {} +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java index 07b219a146d8..f032619dd66b 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java @@ -51,6 +51,7 @@ */ @Slf4j public class PartitionWorker { + private static final long MAX_CURSOR_INITIALIZATION_OFFSET = (long) Integer.MAX_VALUE + 1L; /** Time series entity types that need special handling */ private static final Set TIME_SERIES_ENTITIES = @@ -575,7 +576,7 @@ private String initializeKeysetCursor(String entityType, long offset) { return null; } if (!TIME_SERIES_ENTITIES.contains(entityType)) { - int cursorOffset = (int) offset - 1; + int cursorOffset = toCursorOffset(entityType, offset); ListFilter filter = new ListFilter(Include.ALL); String cursor = Entity.getEntityRepository(entityType).getCursorAtOffset(filter, cursorOffset); @@ -592,6 +593,17 @@ private String initializeKeysetCursor(String entityType, long offset) { } } + private int toCursorOffset(String entityType, long offset) { + long cursorOffset = offset - 1L; + if (cursorOffset > Integer.MAX_VALUE) { + throw new IllegalArgumentException( + String.format( + "Keyset cursor initialization for entityType %s does not support offsets above %d", + entityType, MAX_CURSOR_INITIALIZATION_OFFSET)); + } + return Math.toIntExact(cursorOffset); + } + /** * Write entities to the search index sink. * diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java index 561e405b7dc9..728e657626d5 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java @@ -498,6 +498,18 @@ public interface CollectionDAO { @CreateSqlObject SearchIndexServerStatsDAO searchIndexServerStatsDAO(); + @CreateSqlObject + RdfIndexJobDAO rdfIndexJobDAO(); + + @CreateSqlObject + RdfIndexPartitionDAO rdfIndexPartitionDAO(); + + @CreateSqlObject + RdfReindexLockDAO rdfReindexLockDAO(); + + @CreateSqlObject + RdfIndexServerStatsDAO rdfIndexServerStatsDAO(); + @CreateSqlObject AuditLogDAO auditLogDAO(); @@ -11397,6 +11409,677 @@ public EntityStats map(ResultSet rs, StatementContext ctx) throws SQLException { } } + /** DAO for distributed RDF index jobs. */ + interface RdfIndexJobDAO { + + @ConnectionAwareSqlUpdate( + value = + "INSERT INTO rdf_index_job (id, status, jobConfiguration, totalRecords, processedRecords, " + + "successRecords, failedRecords, stats, createdBy, createdAt, updatedAt) " + + "VALUES (:id, :status, :jobConfiguration, :totalRecords, :processedRecords, " + + ":successRecords, :failedRecords, :stats, :createdBy, :createdAt, :updatedAt)", + connectionType = MYSQL) + @ConnectionAwareSqlUpdate( + value = + "INSERT INTO rdf_index_job (id, status, jobConfiguration, totalRecords, processedRecords, " + + "successRecords, failedRecords, stats, createdBy, createdAt, updatedAt) " + + "VALUES (:id, :status, :jobConfiguration::jsonb, :totalRecords, :processedRecords, " + + ":successRecords, :failedRecords, :stats::jsonb, :createdBy, :createdAt, :updatedAt)", + connectionType = POSTGRES) + void insert( + @Bind("id") String id, + @Bind("status") String status, + @Bind("jobConfiguration") String jobConfiguration, + @Bind("totalRecords") long totalRecords, + @Bind("processedRecords") long processedRecords, + @Bind("successRecords") long successRecords, + @Bind("failedRecords") long failedRecords, + @Bind("stats") String stats, + @Bind("createdBy") String createdBy, + @Bind("createdAt") long createdAt, + @Bind("updatedAt") long updatedAt); + + @ConnectionAwareSqlUpdate( + value = + "UPDATE rdf_index_job SET status = :status, processedRecords = :processedRecords, " + + "successRecords = :successRecords, failedRecords = :failedRecords, stats = :stats, " + + "startedAt = :startedAt, completedAt = :completedAt, updatedAt = :updatedAt, " + + "errorMessage = :errorMessage WHERE id = :id", + connectionType = MYSQL) + @ConnectionAwareSqlUpdate( + value = + "UPDATE rdf_index_job SET status = :status, processedRecords = :processedRecords, " + + "successRecords = :successRecords, failedRecords = :failedRecords, stats = :stats::jsonb, " + + "startedAt = :startedAt, completedAt = :completedAt, updatedAt = :updatedAt, " + + "errorMessage = :errorMessage WHERE id = :id", + connectionType = POSTGRES) + void update( + @Bind("id") String id, + @Bind("status") String status, + @Bind("processedRecords") long processedRecords, + @Bind("successRecords") long successRecords, + @Bind("failedRecords") long failedRecords, + @Bind("stats") String stats, + @Bind("startedAt") Long startedAt, + @Bind("completedAt") Long completedAt, + @Bind("updatedAt") long updatedAt, + @Bind("errorMessage") String errorMessage); + + @SqlUpdate("UPDATE rdf_index_job SET updatedAt = :updatedAt WHERE id = :id") + void touchJob(@Bind("id") String id, @Bind("updatedAt") long updatedAt); + + @SqlQuery("SELECT * FROM rdf_index_job WHERE id = :id") + @RegisterRowMapper(RdfIndexJobMapper.class) + RdfIndexJobRecord findById(@Bind("id") String id); + + @SqlQuery("SELECT * FROM rdf_index_job WHERE status IN () ORDER BY createdAt DESC") + @RegisterRowMapper(RdfIndexJobMapper.class) + List findByStatuses(@BindList("statuses") List statuses); + + @SqlQuery( + "SELECT * FROM rdf_index_job WHERE status IN () ORDER BY createdAt DESC LIMIT :limit") + @RegisterRowMapper(RdfIndexJobMapper.class) + List findByStatusesWithLimit( + @BindList("statuses") List statuses, @Bind("limit") int limit); + + @SqlQuery("SELECT id FROM rdf_index_job WHERE status IN ('READY', 'RUNNING', 'STOPPING')") + List getRunningJobIds(); + + @SqlUpdate("DELETE FROM rdf_index_job") + void deleteAll(); + + class RdfIndexJobMapper implements RowMapper { + @Override + public RdfIndexJobRecord map(ResultSet rs, StatementContext ctx) throws SQLException { + return new RdfIndexJobRecord( + rs.getString("id"), + rs.getString("status"), + rs.getString("jobConfiguration"), + rs.getLong("totalRecords"), + rs.getLong("processedRecords"), + rs.getLong("successRecords"), + rs.getLong("failedRecords"), + rs.getString("stats"), + rs.getString("createdBy"), + rs.getLong("createdAt"), + (Long) rs.getObject("startedAt"), + (Long) rs.getObject("completedAt"), + rs.getLong("updatedAt"), + rs.getString("errorMessage")); + } + } + + record RdfIndexJobRecord( + String id, + String status, + String jobConfiguration, + long totalRecords, + long processedRecords, + long successRecords, + long failedRecords, + String stats, + String createdBy, + long createdAt, + Long startedAt, + Long completedAt, + long updatedAt, + String errorMessage) {} + } + + /** DAO for distributed RDF partitions. */ + interface RdfIndexPartitionDAO { + + @ConnectionAwareSqlUpdate( + value = + "INSERT INTO rdf_index_partition (id, jobId, entityType, partitionIndex, rangeStart, rangeEnd, " + + "estimatedCount, workUnits, priority, status, processingCursor, claimableAt) " + + "VALUES (:id, :jobId, :entityType, :partitionIndex, :rangeStart, :rangeEnd, " + + ":estimatedCount, :workUnits, :priority, :status, :cursor, :claimableAt)", + connectionType = MYSQL) + @ConnectionAwareSqlUpdate( + value = + "INSERT INTO rdf_index_partition (id, jobId, entityType, partitionIndex, rangeStart, rangeEnd, " + + "estimatedCount, workUnits, priority, status, processingCursor, claimableAt) " + + "VALUES (:id, :jobId, :entityType, :partitionIndex, :rangeStart, :rangeEnd, " + + ":estimatedCount, :workUnits, :priority, :status, :cursor, :claimableAt)", + connectionType = POSTGRES) + void insert( + @Bind("id") String id, + @Bind("jobId") String jobId, + @Bind("entityType") String entityType, + @Bind("partitionIndex") int partitionIndex, + @Bind("rangeStart") long rangeStart, + @Bind("rangeEnd") long rangeEnd, + @Bind("estimatedCount") long estimatedCount, + @Bind("workUnits") long workUnits, + @Bind("priority") int priority, + @Bind("status") String status, + @Bind("cursor") long cursor, + @Bind("claimableAt") long claimableAt); + + @SqlUpdate( + "UPDATE rdf_index_partition SET status = :status, processingCursor = :cursor, " + + "processedCount = :processedCount, successCount = :successCount, failedCount = :failedCount, " + + "assignedServer = :assignedServer, claimedAt = :claimedAt, startedAt = :startedAt, " + + "completedAt = :completedAt, lastUpdateAt = :lastUpdateAt, lastError = :lastError, " + + "retryCount = :retryCount WHERE id = :id") + void update( + @Bind("id") String id, + @Bind("status") String status, + @Bind("cursor") long cursor, + @Bind("processedCount") long processedCount, + @Bind("successCount") long successCount, + @Bind("failedCount") long failedCount, + @Bind("assignedServer") String assignedServer, + @Bind("claimedAt") Long claimedAt, + @Bind("startedAt") Long startedAt, + @Bind("completedAt") Long completedAt, + @Bind("lastUpdateAt") Long lastUpdateAt, + @Bind("lastError") String lastError, + @Bind("retryCount") int retryCount); + + @SqlUpdate( + "UPDATE rdf_index_partition SET processingCursor = :cursor, processedCount = :processedCount, " + + "successCount = :successCount, failedCount = :failedCount, lastUpdateAt = :lastUpdateAt " + + "WHERE id = :id") + void updateProgress( + @Bind("id") String id, + @Bind("cursor") long cursor, + @Bind("processedCount") long processedCount, + @Bind("successCount") long successCount, + @Bind("failedCount") long failedCount, + @Bind("lastUpdateAt") long lastUpdateAt); + + @SqlUpdate("UPDATE rdf_index_partition SET lastUpdateAt = :lastUpdateAt WHERE id = :id") + void updateHeartbeat(@Bind("id") String id, @Bind("lastUpdateAt") long lastUpdateAt); + + @SqlQuery("SELECT * FROM rdf_index_partition WHERE id = :id") + @RegisterRowMapper(RdfIndexPartitionMapper.class) + RdfIndexPartitionRecord findById(@Bind("id") String id); + + @SqlQuery( + "SELECT * FROM rdf_index_partition WHERE jobId = :jobId ORDER BY priority DESC, entityType, partitionIndex") + @RegisterRowMapper(RdfIndexPartitionMapper.class) + List findByJobId(@Bind("jobId") String jobId); + + @SqlQuery( + "SELECT COUNT(*) FROM rdf_index_partition WHERE jobId = :jobId AND status = 'PENDING'") + int countPendingPartitions(@Bind("jobId") String jobId); + + @SqlQuery( + "SELECT COUNT(*) FROM rdf_index_partition WHERE jobId = :jobId AND status = 'PROCESSING'") + int countInFlightPartitions(@Bind("jobId") String jobId); + + @ConnectionAwareSqlUpdate( + value = + "UPDATE rdf_index_partition p " + + "JOIN (SELECT id FROM rdf_index_partition WHERE jobId = :jobId AND status = 'PENDING' " + + "AND claimableAt <= :now " + + "ORDER BY priority DESC, entityType, partitionIndex LIMIT 1 FOR UPDATE SKIP LOCKED) t ON p.id = t.id " + + "SET p.status = 'PROCESSING', p.assignedServer = :serverId, p.claimedAt = :now, " + + "p.startedAt = :now, p.lastUpdateAt = :now", + connectionType = MYSQL) + @ConnectionAwareSqlUpdate( + value = + "UPDATE rdf_index_partition SET status = 'PROCESSING', " + + "assignedServer = :serverId, claimedAt = :now, startedAt = :now, lastUpdateAt = :now " + + "WHERE id = (SELECT id FROM rdf_index_partition WHERE jobId = :jobId AND status = 'PENDING' " + + "AND claimableAt <= :now " + + "ORDER BY priority DESC, entityType, partitionIndex LIMIT 1 FOR UPDATE SKIP LOCKED)", + connectionType = POSTGRES) + int claimNextPartitionAtomic( + @Bind("jobId") String jobId, @Bind("serverId") String serverId, @Bind("now") long now); + + @SqlQuery( + "SELECT * FROM rdf_index_partition WHERE jobId = :jobId AND status = 'PROCESSING' " + + "AND assignedServer = :serverId AND claimedAt = :claimedAt " + + "ORDER BY priority DESC, entityType, partitionIndex LIMIT 1") + @RegisterRowMapper(RdfIndexPartitionMapper.class) + RdfIndexPartitionRecord findLatestClaimedPartition( + @Bind("jobId") String jobId, + @Bind("serverId") String serverId, + @Bind("claimedAt") long claimedAt); + + @SqlUpdate( + "UPDATE rdf_index_partition SET status = 'PENDING', assignedServer = NULL, claimedAt = NULL, " + + "retryCount = retryCount + 1, lastError = 'Reclaimed due to stale heartbeat' " + + "WHERE jobId = :jobId AND status = 'PROCESSING' AND lastUpdateAt < :staleThreshold " + + "AND retryCount < :maxRetries") + int reclaimStalePartitionsForRetry( + @Bind("jobId") String jobId, + @Bind("staleThreshold") long staleThreshold, + @Bind("maxRetries") int maxRetries); + + @SqlUpdate( + "UPDATE rdf_index_partition SET status = 'FAILED', " + + "lastError = 'Exceeded max retries after stale heartbeat', completedAt = :now " + + "WHERE jobId = :jobId AND status = 'PROCESSING' AND lastUpdateAt < :staleThreshold " + + "AND retryCount >= :maxRetries") + int failStalePartitionsExceedingRetries( + @Bind("jobId") String jobId, + @Bind("staleThreshold") long staleThreshold, + @Bind("maxRetries") int maxRetries, + @Bind("now") long now); + + @SqlUpdate( + "UPDATE rdf_index_partition SET status = 'CANCELLED' WHERE jobId = :jobId AND status = 'PENDING'") + int cancelPendingPartitions(@Bind("jobId") String jobId); + + @SqlUpdate( + "UPDATE rdf_index_partition SET status = :status, assignedServer = NULL, claimedAt = NULL, " + + "lastError = :reason, lastUpdateAt = :updatedAt, completedAt = :completedAt " + + "WHERE jobId = :jobId AND status = 'PROCESSING' AND assignedServer = :serverId") + int releaseProcessingPartitions( + @Bind("jobId") String jobId, + @Bind("serverId") String serverId, + @Bind("status") String status, + @Bind("reason") String reason, + @Bind("updatedAt") long updatedAt, + @Bind("completedAt") Long completedAt); + + @SqlQuery( + "SELECT entityType, " + + "SUM(estimatedCount) as totalRecords, " + + "SUM(processedCount) as processedRecords, " + + "SUM(successCount) as successRecords, " + + "SUM(failedCount) as failedRecords, " + + "COUNT(*) as totalPartitions, " + + "SUM(CASE WHEN status = 'COMPLETED' THEN 1 ELSE 0 END) as completedPartitions, " + + "SUM(CASE WHEN status = 'FAILED' THEN 1 ELSE 0 END) as failedPartitions " + + "FROM rdf_index_partition WHERE jobId = :jobId GROUP BY entityType") + @RegisterRowMapper(RdfEntityStatsMapper.class) + List getEntityStats(@Bind("jobId") String jobId); + + @SqlQuery( + "SELECT " + + "SUM(estimatedCount) as totalRecords, " + + "SUM(processedCount) as processedRecords, " + + "SUM(successCount) as successRecords, " + + "SUM(failedCount) as failedRecords, " + + "COUNT(*) as totalPartitions, " + + "SUM(CASE WHEN status = 'COMPLETED' THEN 1 ELSE 0 END) as completedPartitions, " + + "SUM(CASE WHEN status = 'FAILED' THEN 1 ELSE 0 END) as failedPartitions, " + + "SUM(CASE WHEN status = 'PENDING' THEN 1 ELSE 0 END) as pendingPartitions, " + + "SUM(CASE WHEN status = 'PROCESSING' THEN 1 ELSE 0 END) as processingPartitions " + + "FROM rdf_index_partition WHERE jobId = :jobId") + @RegisterRowMapper(RdfAggregatedStatsMapper.class) + RdfAggregatedStatsRecord getAggregatedStats(@Bind("jobId") String jobId); + + @SqlQuery( + "SELECT assignedServer, " + + "SUM(processedCount) as processedRecords, " + + "SUM(successCount) as successRecords, " + + "SUM(failedCount) as failedRecords, " + + "COUNT(*) as totalPartitions, " + + "SUM(CASE WHEN status = 'COMPLETED' THEN 1 ELSE 0 END) as completedPartitions, " + + "SUM(CASE WHEN status = 'PROCESSING' THEN 1 ELSE 0 END) as processingPartitions " + + "FROM rdf_index_partition WHERE jobId = :jobId AND assignedServer IS NOT NULL " + + "GROUP BY assignedServer") + @RegisterRowMapper(RdfServerStatsMapper.class) + List getServerStats(@Bind("jobId") String jobId); + + @SqlQuery( + "SELECT DISTINCT assignedServer FROM rdf_index_partition " + + "WHERE jobId = :jobId AND assignedServer IS NOT NULL") + List getAssignedServers(@Bind("jobId") String jobId); + + @SqlUpdate("DELETE FROM rdf_index_partition") + void deleteAll(); + + class RdfIndexPartitionMapper implements RowMapper { + @Override + public RdfIndexPartitionRecord map(ResultSet rs, StatementContext ctx) throws SQLException { + return new RdfIndexPartitionRecord( + rs.getString("id"), + rs.getString("jobId"), + rs.getString("entityType"), + rs.getInt("partitionIndex"), + rs.getLong("rangeStart"), + rs.getLong("rangeEnd"), + rs.getLong("estimatedCount"), + rs.getLong("workUnits"), + rs.getInt("priority"), + rs.getString("status"), + rs.getLong("processingCursor"), + rs.getLong("processedCount"), + rs.getLong("successCount"), + rs.getLong("failedCount"), + rs.getString("assignedServer"), + (Long) rs.getObject("claimedAt"), + (Long) rs.getObject("startedAt"), + (Long) rs.getObject("completedAt"), + (Long) rs.getObject("lastUpdateAt"), + rs.getString("lastError"), + rs.getInt("retryCount"), + rs.getLong("claimableAt")); + } + } + + class RdfEntityStatsMapper implements RowMapper { + @Override + public RdfEntityStatsRecord map(ResultSet rs, StatementContext ctx) throws SQLException { + return new RdfEntityStatsRecord( + rs.getString("entityType"), + rs.getLong("totalRecords"), + rs.getLong("processedRecords"), + rs.getLong("successRecords"), + rs.getLong("failedRecords"), + rs.getInt("totalPartitions"), + rs.getInt("completedPartitions"), + rs.getInt("failedPartitions")); + } + } + + class RdfAggregatedStatsMapper implements RowMapper { + @Override + public RdfAggregatedStatsRecord map(ResultSet rs, StatementContext ctx) throws SQLException { + return new RdfAggregatedStatsRecord( + rs.getLong("totalRecords"), + rs.getLong("processedRecords"), + rs.getLong("successRecords"), + rs.getLong("failedRecords"), + rs.getInt("totalPartitions"), + rs.getInt("completedPartitions"), + rs.getInt("failedPartitions"), + rs.getInt("pendingPartitions"), + rs.getInt("processingPartitions")); + } + } + + class RdfServerStatsMapper implements RowMapper { + @Override + public RdfServerPartitionStatsRecord map(ResultSet rs, StatementContext ctx) + throws SQLException { + return new RdfServerPartitionStatsRecord( + rs.getString("assignedServer"), + rs.getLong("processedRecords"), + rs.getLong("successRecords"), + rs.getLong("failedRecords"), + rs.getInt("totalPartitions"), + rs.getInt("completedPartitions"), + rs.getInt("processingPartitions")); + } + } + + record RdfIndexPartitionRecord( + String id, + String jobId, + String entityType, + int partitionIndex, + long rangeStart, + long rangeEnd, + long estimatedCount, + long workUnits, + int priority, + String status, + long cursor, + long processedCount, + long successCount, + long failedCount, + String assignedServer, + Long claimedAt, + Long startedAt, + Long completedAt, + Long lastUpdateAt, + String lastError, + int retryCount, + long claimableAt) {} + + record RdfEntityStatsRecord( + String entityType, + long totalRecords, + long processedRecords, + long successRecords, + long failedRecords, + int totalPartitions, + int completedPartitions, + int failedPartitions) {} + + record RdfAggregatedStatsRecord( + long totalRecords, + long processedRecords, + long successRecords, + long failedRecords, + int totalPartitions, + int completedPartitions, + int failedPartitions, + int pendingPartitions, + int processingPartitions) {} + + record RdfServerPartitionStatsRecord( + String serverId, + long processedRecords, + long successRecords, + long failedRecords, + int totalPartitions, + int completedPartitions, + int processingPartitions) {} + } + + /** DAO for RDF distributed reindex lock. */ + interface RdfReindexLockDAO { + + @ConnectionAwareSqlUpdate( + value = + "INSERT IGNORE INTO rdf_reindex_lock (lockKey, jobId, serverId, acquiredAt, lastHeartbeat, expiresAt) " + + "VALUES (:lockKey, :jobId, :serverId, :acquiredAt, :lastHeartbeat, :expiresAt)", + connectionType = MYSQL) + @ConnectionAwareSqlUpdate( + value = + "INSERT INTO rdf_reindex_lock (lockKey, jobId, serverId, acquiredAt, lastHeartbeat, expiresAt) " + + "VALUES (:lockKey, :jobId, :serverId, :acquiredAt, :lastHeartbeat, :expiresAt) " + + "ON CONFLICT (lockKey) DO NOTHING", + connectionType = POSTGRES) + int insertIfNotExists( + @Bind("lockKey") String lockKey, + @Bind("jobId") String jobId, + @Bind("serverId") String serverId, + @Bind("acquiredAt") long acquiredAt, + @Bind("lastHeartbeat") long lastHeartbeat, + @Bind("expiresAt") long expiresAt); + + @SqlUpdate( + "UPDATE rdf_reindex_lock SET lastHeartbeat = :lastHeartbeat, expiresAt = :expiresAt " + + "WHERE lockKey = :lockKey AND jobId = :jobId") + int updateHeartbeat( + @Bind("lockKey") String lockKey, + @Bind("jobId") String jobId, + @Bind("lastHeartbeat") long lastHeartbeat, + @Bind("expiresAt") long expiresAt); + + @SqlQuery("SELECT * FROM rdf_reindex_lock WHERE lockKey = :lockKey") + @RegisterRowMapper(RdfReindexLockMapper.class) + RdfReindexLockRecord findByKey(@Bind("lockKey") String lockKey); + + @SqlUpdate("DELETE FROM rdf_reindex_lock WHERE lockKey = :lockKey") + void delete(@Bind("lockKey") String lockKey); + + @SqlUpdate("DELETE FROM rdf_reindex_lock WHERE lockKey = :lockKey AND jobId = :jobId") + int deleteByKeyAndJob(@Bind("lockKey") String lockKey, @Bind("jobId") String jobId); + + @SqlUpdate("DELETE FROM rdf_reindex_lock WHERE expiresAt < :now") + int deleteExpiredLocks(@Bind("now") long now); + + @SqlUpdate( + "UPDATE rdf_reindex_lock SET jobId = :toJobId, serverId = :serverId, " + + "lastHeartbeat = :heartbeat, expiresAt = :expiresAt " + + "WHERE lockKey = :lockKey AND jobId = :fromJobId") + int updateLockOwner( + @Bind("lockKey") String lockKey, + @Bind("fromJobId") String fromJobId, + @Bind("toJobId") String toJobId, + @Bind("serverId") String serverId, + @Bind("heartbeat") long heartbeat, + @Bind("expiresAt") long expiresAt); + + default boolean tryAcquireLock( + String lockKey, String jobId, String serverId, long acquiredAt, long expiresAt) { + deleteExpiredLocks(System.currentTimeMillis()); + int inserted = insertIfNotExists(lockKey, jobId, serverId, acquiredAt, acquiredAt, expiresAt); + if (inserted > 0) { + return true; + } + + RdfReindexLockRecord existing = findByKey(lockKey); + if (existing != null && existing.isExpired()) { + delete(lockKey); + inserted = insertIfNotExists(lockKey, jobId, serverId, acquiredAt, acquiredAt, expiresAt); + return inserted > 0; + } + return false; + } + + default void releaseLock(String lockKey, String jobId) { + deleteByKeyAndJob(lockKey, jobId); + } + + default boolean transferLock( + String lockKey, + String fromJobId, + String toJobId, + String serverId, + long heartbeat, + long expiresAt) { + return updateLockOwner(lockKey, fromJobId, toJobId, serverId, heartbeat, expiresAt) > 0; + } + + class RdfReindexLockMapper implements RowMapper { + @Override + public RdfReindexLockRecord map(ResultSet rs, StatementContext ctx) throws SQLException { + return new RdfReindexLockRecord( + rs.getString("lockKey"), + rs.getString("jobId"), + rs.getString("serverId"), + rs.getLong("acquiredAt"), + rs.getLong("lastHeartbeat"), + rs.getLong("expiresAt")); + } + } + + record RdfReindexLockRecord( + String lockKey, + String jobId, + String serverId, + long acquiredAt, + long lastHeartbeat, + long expiresAt) { + + public boolean isExpired() { + return System.currentTimeMillis() > expiresAt; + } + } + } + + /** DAO for RDF per-server distributed stats. */ + interface RdfIndexServerStatsDAO { + + record ServerStatsRecord( + String id, + String jobId, + String serverId, + String entityType, + long processedRecords, + long successRecords, + long failedRecords, + int partitionsCompleted, + int partitionsFailed, + long lastUpdatedAt) {} + + record AggregatedServerStats( + long processedRecords, + long successRecords, + long failedRecords, + int partitionsCompleted, + int partitionsFailed) {} + + @ConnectionAwareSqlUpdate( + value = + "INSERT INTO rdf_index_server_stats (id, jobId, serverId, entityType, processedRecords, " + + "successRecords, failedRecords, partitionsCompleted, partitionsFailed, lastUpdatedAt) " + + "VALUES (:id, :jobId, :serverId, :entityType, :processedRecords, :successRecords, " + + ":failedRecords, :partitionsCompleted, :partitionsFailed, :lastUpdatedAt) " + + "ON DUPLICATE KEY UPDATE " + + "processedRecords = processedRecords + VALUES(processedRecords), " + + "successRecords = successRecords + VALUES(successRecords), " + + "failedRecords = failedRecords + VALUES(failedRecords), " + + "partitionsCompleted = partitionsCompleted + VALUES(partitionsCompleted), " + + "partitionsFailed = partitionsFailed + VALUES(partitionsFailed), " + + "lastUpdatedAt = VALUES(lastUpdatedAt)", + connectionType = MYSQL) + @ConnectionAwareSqlUpdate( + value = + "INSERT INTO rdf_index_server_stats (id, jobId, serverId, entityType, processedRecords, " + + "successRecords, failedRecords, partitionsCompleted, partitionsFailed, lastUpdatedAt) " + + "VALUES (:id, :jobId, :serverId, :entityType, :processedRecords, :successRecords, " + + ":failedRecords, :partitionsCompleted, :partitionsFailed, :lastUpdatedAt) " + + "ON CONFLICT (jobId, serverId, entityType) DO UPDATE SET " + + "processedRecords = rdf_index_server_stats.processedRecords + EXCLUDED.processedRecords, " + + "successRecords = rdf_index_server_stats.successRecords + EXCLUDED.successRecords, " + + "failedRecords = rdf_index_server_stats.failedRecords + EXCLUDED.failedRecords, " + + "partitionsCompleted = rdf_index_server_stats.partitionsCompleted + EXCLUDED.partitionsCompleted, " + + "partitionsFailed = rdf_index_server_stats.partitionsFailed + EXCLUDED.partitionsFailed, " + + "lastUpdatedAt = EXCLUDED.lastUpdatedAt", + connectionType = POSTGRES) + void incrementStats( + @Bind("id") String id, + @Bind("jobId") String jobId, + @Bind("serverId") String serverId, + @Bind("entityType") String entityType, + @Bind("processedRecords") long processedRecords, + @Bind("successRecords") long successRecords, + @Bind("failedRecords") long failedRecords, + @Bind("partitionsCompleted") int partitionsCompleted, + @Bind("partitionsFailed") int partitionsFailed, + @Bind("lastUpdatedAt") long lastUpdatedAt); + + @SqlQuery("SELECT * FROM rdf_index_server_stats WHERE jobId = :jobId") + @RegisterRowMapper(RdfServerStatsRecordMapper.class) + List findByJobId(@Bind("jobId") String jobId); + + @SqlQuery( + "SELECT " + + "COALESCE(SUM(processedRecords), 0) as processedRecords, " + + "COALESCE(SUM(successRecords), 0) as successRecords, " + + "COALESCE(SUM(failedRecords), 0) as failedRecords, " + + "COALESCE(SUM(partitionsCompleted), 0) as partitionsCompleted, " + + "COALESCE(SUM(partitionsFailed), 0) as partitionsFailed " + + "FROM rdf_index_server_stats WHERE jobId = :jobId") + @RegisterRowMapper(RdfAggregatedServerStatsMapper.class) + AggregatedServerStats getAggregatedStats(@Bind("jobId") String jobId); + + @SqlUpdate("DELETE FROM rdf_index_server_stats") + void deleteAll(); + + class RdfServerStatsRecordMapper implements RowMapper { + @Override + public ServerStatsRecord map(ResultSet rs, StatementContext ctx) throws SQLException { + return new ServerStatsRecord( + rs.getString("id"), + rs.getString("jobId"), + rs.getString("serverId"), + rs.getString("entityType"), + rs.getLong("processedRecords"), + rs.getLong("successRecords"), + rs.getLong("failedRecords"), + rs.getInt("partitionsCompleted"), + rs.getInt("partitionsFailed"), + rs.getLong("lastUpdatedAt")); + } + } + + class RdfAggregatedServerStatsMapper implements RowMapper { + @Override + public AggregatedServerStats map(ResultSet rs, StatementContext ctx) throws SQLException { + return new AggregatedServerStats( + rs.getLong("processedRecords"), + rs.getLong("successRecords"), + rs.getLong("failedRecords"), + rs.getInt("partitionsCompleted"), + rs.getInt("partitionsFailed")); + } + } + } + @RegisterRowMapper(AuditLogRecordMapper.class) interface AuditLogDAO { @ConnectionAwareSqlUpdate( diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java index a776df3030b5..2ff1d3fbf18e 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java @@ -24,6 +24,7 @@ public class ListFilter extends Filter { public static final String NULL_PARAM = "null"; + private static final String MCP_EXECUTION_TABLE_NAME = "mcp_execution_entity"; public ListFilter() { this(Include.NON_DELETED); @@ -71,7 +72,7 @@ public String getCondition(String tableName) { conditions.add(getAgentTypeCondition()); conditions.add(getProviderCondition(tableName)); conditions.add(getEntityStatusCondition(tableName)); - conditions.add(getServerIdCondition()); + conditions.add(getServerIdCondition(tableName)); String condition = addCondition(conditions); return condition.isEmpty() ? "WHERE TRUE" : "WHERE " + condition; } @@ -399,9 +400,11 @@ public String getApiCollectionCondition(String apiEndpoint) { : getFqnPrefixCondition(apiEndpoint, apiCollection, "apiCollection"); } - private String getServerIdCondition() { + private String getServerIdCondition(String tableName) { String serverId = queryParams.get("serverId"); - return serverId == null ? "" : "serverId = :serverId"; + return serverId == null || !MCP_EXECUTION_TABLE_NAME.equals(tableName) + ? "" + : "serverId = :serverId"; } private String getEntityFQNHashCondition() { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfRepository.java index 308645ca3130..04f8d1c459a5 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfRepository.java @@ -4,9 +4,12 @@ import java.io.IOException; import java.io.InputStream; import java.util.ArrayList; +import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Set; import java.util.UUID; @@ -754,49 +757,96 @@ private String getJenaFormat(String mimeType) { return "TURTLE"; // default } - public String getEntityGraph(UUID entityId, String entityType, int depth) throws IOException { + public String getEntityGraph( + UUID entityId, + String entityType, + int depth, + Set entityTypes, + Set relationshipTypes) + throws IOException { if (!isEnabled()) { throw new IllegalStateException("RDF Repository is not enabled"); } - String entityUri = config.getBaseUri().toString() + "entity/" + entityType + "/" + entityId; + String validatedEntityType = requireKnownEntityType(entityType); + String entityUri = + config.getBaseUri().toString() + "entity/" + validatedEntityType + "/" + entityId; try { - Set visitedNodes = new HashSet<>(); - Set currentLevelNodes = new HashSet<>(); - List allEdges = new ArrayList<>(); - - currentLevelNodes.add(entityUri); - visitedNodes.add(entityUri); - - for (int currentDepth = 0; - currentDepth < depth && !currentLevelNodes.isEmpty(); - currentDepth++) { - Set nextLevelNodes = new HashSet<>(); - - // For each node at current level, get its relationships - for (String nodeUri : currentLevelNodes) { - String sparql = buildSingleNodeQuery(nodeUri); - String results = - storageService.executeSparqlQuery(sparql, "application/sparql-results+json"); - - if (results != null && !results.trim().isEmpty()) { - List edges = parseEdgesFromResults(results, visitedNodes, nextLevelNodes); - allEdges.addAll(edges); - } - } - - currentLevelNodes = nextLevelNodes; - visitedNodes.addAll(nextLevelNodes); - } - - return convertEdgesToGraphData(allEdges); + EntityGraphTraversalResult traversalResult = traverseEntityGraph(entityUri, depth); + FilteredEntityGraph filteredGraph = + applyGraphFilters( + entityUri, + traversalResult.nodeUris(), + traversalResult.edges(), + entityTypes, + relationshipTypes); + + return convertEdgesToGraphData( + entityUri, + filteredGraph.nodeUris(), + filteredGraph.edges(), + buildEntityTypeFilterOptions(traversalResult.nodeUris()), + buildRelationshipFilterOptions(traversalResult.edges())); } catch (Exception e) { LOG.error("Error getting entity graph for {}", entityUri, e); throw new IOException("Failed to get entity graph", e); } } + public String exportEntityGraph( + UUID entityId, + String entityType, + int depth, + Set entityTypes, + Set relationshipTypes, + String format) + throws IOException { + if (!isEnabled()) { + throw new IllegalStateException("RDF Repository is not enabled"); + } + + String validatedEntityType = requireKnownEntityType(entityType); + String normalizedFormat = normalizeEntityGraphExportFormat(format); + String entityUri = + config.getBaseUri().toString() + "entity/" + validatedEntityType + "/" + entityId; + + try { + EntityGraphTraversalResult traversalResult = traverseEntityGraph(entityUri, depth); + FilteredEntityGraph filteredGraph = + applyGraphFilters( + entityUri, + traversalResult.nodeUris(), + traversalResult.edges(), + entityTypes, + relationshipTypes); + + Model model = buildEntityGraphExportModel(filteredGraph.nodeUris(), filteredGraph.edges()); + + java.io.StringWriter writer = new java.io.StringWriter(); + model.write(writer, normalizedFormat); + + return writer.toString(); + } catch (Exception e) { + LOG.error("Error exporting entity graph for {}", entityUri, e); + throw new IOException("Failed to export entity graph", e); + } + } + + private String requireKnownEntityType(String entityType) { + if (entityType == null || entityType.isBlank()) { + throw new IllegalArgumentException("Entity type is required"); + } + + String trimmedEntityType = entityType.trim(); + if (!trimmedEntityType.matches("[A-Za-z][A-Za-z0-9]*") + || !Entity.hasEntityRepository(trimmedEntityType)) { + throw new IllegalArgumentException("Invalid entity type"); + } + + return trimmedEntityType; + } + /** * Get glossary term relationship graph with pagination support. * This method queries the RDF store for glossary terms and their relationships, @@ -1494,56 +1544,157 @@ private String formatRelationshipLabel(String relationship) { }; } - private String buildSingleNodeQuery(String nodeUri) { + private EntityGraphTraversalResult traverseEntityGraph(String rootUri, int depth) { + Set visitedNodes = new HashSet<>(); + Set currentLevelNodes = new HashSet<>(); + Set discoveredNodes = new HashSet<>(); + Set edgeKeys = new HashSet<>(); + List allEdges = new ArrayList<>(); + + currentLevelNodes.add(rootUri); + visitedNodes.add(rootUri); + discoveredNodes.add(rootUri); + + for (int currentDepth = 0; + currentDepth < depth && !currentLevelNodes.isEmpty(); + currentDepth++) { + String sparql = buildEntityGraphBatchQuery(currentLevelNodes); + String results = storageService.executeSparqlQuery(sparql, "application/sparql-results+json"); + + Set nextLevelNodes = new HashSet<>(); + if (results != null && !results.trim().isEmpty()) { + allEdges.addAll( + parseEntityGraphEdgesFromResults( + results, visitedNodes, nextLevelNodes, discoveredNodes, edgeKeys)); + } + + nextLevelNodes.removeAll(visitedNodes); + visitedNodes.addAll(nextLevelNodes); + currentLevelNodes = nextLevelNodes; + } + + return new EntityGraphTraversalResult(discoveredNodes, allEdges); + } + + private String buildEntityGraphBatchQuery(Set nodeUris) { + String entityPrefix = escapeSparqlStringLiteral(config.getBaseUri().toString() + "entity/"); + String valuesClause = + nodeUris.stream() + .sorted() + .map(this::escapeSparqlUri) + .map(uri -> "<" + uri + ">") + .collect(java.util.stream.Collectors.joining(" ")); + return "PREFIX om: " + "PREFIX rdfs: " + "PREFIX rdf: " + "SELECT DISTINCT ?subject ?predicate ?object WHERE { " + " { " - + " GRAPH ?g { <" - + nodeUri - + "> ?predicate ?object . " - + " FILTER(isIRI(?object) && " - + " ?predicate != rdf:type && " - + " ?predicate != rdfs:label) } " - + " BIND(<" - + nodeUri - + "> AS ?subject) " + + " VALUES ?frontier { " + + valuesClause + + " } " + + " GRAPH ?g { " + + " ?frontier ?predicate ?object . " + + " FILTER(isIRI(?object) && " + + " STRSTARTS(STR(?object), \"" + + entityPrefix + + "\") && " + + " ?predicate != rdf:type && " + + " ?predicate != rdfs:label) " + + " } " + + " BIND(?frontier AS ?subject) " + " } UNION { " - + " GRAPH ?g { ?subject ?predicate <" - + nodeUri - + "> . " - + " FILTER(isIRI(?subject) && " - + " ?predicate != rdf:type && " - + " ?predicate != rdfs:label) } " - + " BIND(<" - + nodeUri - + "> AS ?object) " + + " VALUES ?frontier { " + + valuesClause + + " } " + + " GRAPH ?g { " + + " ?subject ?predicate ?frontier . " + + " FILTER(isIRI(?subject) && " + + " STRSTARTS(STR(?subject), \"" + + entityPrefix + + "\") && " + + " ?predicate != rdf:type && " + + " ?predicate != rdfs:label) " + + " } " + + " BIND(?frontier AS ?object) " + " } " - + "} LIMIT 200"; + + "} LIMIT 5000"; } - private List parseEdgesFromResults( - String sparqlResults, Set visitedNodes, Set nextLevelNodes) { + private String escapeSparqlUri(String uri) { + if (uri == null || uri.isBlank()) { + throw new IllegalArgumentException("Invalid URI for SPARQL: " + uri); + } + + if (uri.contains("<") || uri.contains(">") || uri.chars().anyMatch(Character::isWhitespace)) { + throw new IllegalArgumentException("Invalid URI for SPARQL: " + uri); + } + + try { + java.net.URI parsedUri = java.net.URI.create(uri); + if (!parsedUri.isAbsolute()) { + throw new IllegalArgumentException("Invalid URI for SPARQL: " + uri); + } + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Invalid URI for SPARQL: " + uri, e); + } + + return uri; + } + + private String escapeSparqlStringLiteral(String value) { + return value + .replace("\\", "\\\\") + .replace("\"", "\\\"") + .replace("\n", "\\n") + .replace("\r", "\\r") + .replace("\t", "\\t"); + } + + private List parseEntityGraphEdgesFromResults( + String sparqlResults, + Set visitedNodes, + Set nextLevelNodes, + Set discoveredNodes, + Set edgeKeys) { List edges = new ArrayList<>(); com.fasterxml.jackson.databind.JsonNode resultsJson = JsonUtils.readTree(sparqlResults); if (resultsJson.has("results") && resultsJson.get("results").has("bindings")) { for (com.fasterxml.jackson.databind.JsonNode binding : resultsJson.get("results").get("bindings")) { - String subjectUri = binding.get("subject").get("value").asText(); - String objectUri = binding.get("object").get("value").asText(); - String predicate = binding.get("predicate").get("value").asText(); + String subjectUri = + binding.has("subject") ? binding.get("subject").get("value").asText() : null; + String objectUri = + binding.has("object") ? binding.get("object").get("value").asText() : null; + String predicate = + binding.has("predicate") ? binding.get("predicate").get("value").asText() : null; + + if (!isEntityUri(subjectUri) || !isEntityUri(objectUri)) { + continue; + } - EdgeInfo edge = new EdgeInfo(subjectUri, objectUri, extractPredicateName(predicate)); - edges.add(edge); + String relationType = extractEntityRelationType(predicate); + if (relationType == null || relationType.isBlank()) { + continue; + } - if (!visitedNodes.contains(objectUri)) { - nextLevelNodes.add(objectUri); + String edgeKey = subjectUri + "|" + relationType + "|" + objectUri; + if (!edgeKeys.add(edgeKey)) { + continue; } + + EdgeInfo edge = new EdgeInfo(subjectUri, objectUri, relationType, predicate); + edges.add(edge); + discoveredNodes.add(subjectUri); + discoveredNodes.add(objectUri); + if (!visitedNodes.contains(subjectUri)) { nextLevelNodes.add(subjectUri); } + if (!visitedNodes.contains(objectUri)) { + nextLevelNodes.add(objectUri); + } } } @@ -1554,60 +1705,412 @@ private static class EdgeInfo { final String fromUri; final String toUri; final String relation; + final String predicateUri; - EdgeInfo(String fromUri, String toUri, String relation) { + EdgeInfo(String fromUri, String toUri, String relation, String predicateUri) { this.fromUri = fromUri; this.toUri = toUri; this.relation = relation; + this.predicateUri = predicateUri; + } + } + + private FilteredEntityGraph applyGraphFilters( + String rootUri, + Set nodeUris, + List edges, + Set entityTypeFilters, + Set relationshipTypeFilters) { + if ((entityTypeFilters == null || entityTypeFilters.isEmpty()) + && (relationshipTypeFilters == null || relationshipTypeFilters.isEmpty())) { + return new FilteredEntityGraph(new HashSet<>(nodeUris), edges); + } + + Set normalizedEntityFilters = new HashSet<>(); + if (entityTypeFilters != null) { + entityTypeFilters.stream() + .map(this::normalizeEntityTypeFilter) + .filter(value -> !value.isBlank()) + .forEach(normalizedEntityFilters::add); + } + + Set normalizedRelationshipFilters = new HashSet<>(); + if (relationshipTypeFilters != null) { + relationshipTypeFilters.stream() + .map(this::normalizeRelationTypeFilter) + .filter(value -> !value.isBlank()) + .forEach(normalizedRelationshipFilters::add); + } + + Set allowedNodes = new HashSet<>(); + for (String nodeUri : nodeUris) { + if (rootUri.equals(nodeUri) + || normalizedEntityFilters.isEmpty() + || normalizedEntityFilters.contains( + normalizeEntityTypeFilter(extractEntityTypeFromUri(nodeUri)))) { + allowedNodes.add(nodeUri); + } + } + + List filteredEdges = new ArrayList<>(); + Set connectedNodes = new HashSet<>(); + connectedNodes.add(rootUri); + + for (EdgeInfo edge : edges) { + boolean relationshipAllowed = + normalizedRelationshipFilters.isEmpty() + || normalizedRelationshipFilters.contains(normalizeRelationTypeFilter(edge.relation)); + if (!relationshipAllowed) { + continue; + } + + if (!allowedNodes.contains(edge.fromUri) || !allowedNodes.contains(edge.toUri)) { + continue; + } + + filteredEdges.add(edge); + connectedNodes.add(edge.fromUri); + connectedNodes.add(edge.toUri); } + + Set filteredNodes = new HashSet<>(); + for (String nodeUri : allowedNodes) { + if (rootUri.equals(nodeUri) || connectedNodes.contains(nodeUri)) { + filteredNodes.add(nodeUri); + } + } + filteredNodes.add(rootUri); + + return new FilteredEntityGraph(filteredNodes, filteredEdges); + } + + private List buildEntityTypeFilterOptions(Set nodeUris) { + Map counts = new LinkedHashMap<>(); + for (String nodeUri : nodeUris) { + String entityType = extractEntityTypeFromUri(nodeUri); + counts.merge(entityType, 1, Integer::sum); + } + return buildFilterOptions(counts); } - private String convertEdgesToGraphData(List edges) { + private List buildRelationshipFilterOptions(List edges) { + Map counts = new LinkedHashMap<>(); + for (EdgeInfo edge : edges) { + counts.merge(edge.relation, 1, Integer::sum); + } + return buildFilterOptions(counts); + } + + private List buildFilterOptions(Map counts) { + return counts.entrySet().stream() + .map( + entry -> + new FilterOptionInfo( + entry.getKey(), formatRelationshipLabel(entry.getKey()), entry.getValue())) + .sorted( + Comparator.comparingInt(FilterOptionInfo::count) + .reversed() + .thenComparing(FilterOptionInfo::label)) + .toList(); + } + + private Model buildEntityGraphExportModel(Set nodeUris, List edges) { + Model model = ModelFactory.createDefaultModel(); + configureEntityGraphPrefixes(model); + + Map nodeMap = new HashMap<>(); + List orderedNodeUris = nodeUris.stream().sorted().toList(); + + for (String nodeUri : orderedNodeUris) { + com.fasterxml.jackson.databind.node.ObjectNode node = createNodeFromUri(nodeUri); + nodeMap.put(nodeUri, node); + } + + enhanceNodesWithEntityDetails(nodeMap); + + for (String nodeUri : orderedNodeUris) { + addEntityGraphNodeToModel(model, nodeUri, nodeMap.get(nodeUri)); + } + + for (EdgeInfo edge : edges) { + if (edge.predicateUri == null || edge.predicateUri.isBlank()) { + continue; + } + + Resource fromResource = model.createResource(edge.fromUri); + Resource toResource = model.createResource(edge.toUri); + Property predicate = model.createProperty(edge.predicateUri); + fromResource.addProperty(predicate, toResource); + } + + return model; + } + + private void configureEntityGraphPrefixes(Model model) { + model.setNsPrefix("om", "https://open-metadata.org/ontology/"); + model.setNsPrefix("rdf", "http://www.w3.org/1999/02/22-rdf-syntax-ns#"); + model.setNsPrefix("rdfs", "http://www.w3.org/2000/01/rdf-schema#"); + model.setNsPrefix("dcat", "http://www.w3.org/ns/dcat#"); + model.setNsPrefix("prov", "http://www.w3.org/ns/prov#"); + model.setNsPrefix("foaf", "http://xmlns.com/foaf/0.1/"); + model.setNsPrefix("skos", "http://www.w3.org/2004/02/skos/core#"); + model.setNsPrefix("dprod", "https://ekgf.github.io/dprod/"); + } + + private void addEntityGraphNodeToModel( + Model model, String nodeUri, com.fasterxml.jackson.databind.node.ObjectNode node) { + Resource resource = model.createResource(nodeUri); + String entityType = + node != null + ? node.path("type").asText(extractEntityTypeFromUri(nodeUri)) + : extractEntityTypeFromUri(nodeUri); + String rdfTypeUri = resolvePrefixedUri(RdfUtils.getRdfType(entityType)); + + if (rdfTypeUri != null) { + resource.addProperty( + model.createProperty("http://www.w3.org/1999/02/22-rdf-syntax-ns#", "type"), + model.createResource(rdfTypeUri)); + } + + if (node == null) { + return; + } + + addLiteralIfPresent( + resource, + model.createProperty("http://www.w3.org/2000/01/rdf-schema#", "label"), + node.path("label").asText(null)); + addLiteralIfPresent( + resource, + model.createProperty("https://open-metadata.org/ontology/", "name"), + node.path("name").asText(null)); + addLiteralIfPresent( + resource, + model.createProperty("https://open-metadata.org/ontology/", "fullyQualifiedName"), + node.path("fullyQualifiedName").asText(null)); + addLiteralIfPresent( + resource, + model.createProperty("https://open-metadata.org/ontology/", "description"), + node.path("description").asText(null)); + } + + private void addLiteralIfPresent(Resource resource, Property property, String value) { + if (value != null && !value.isBlank()) { + resource.addProperty(property, value); + } + } + + private String resolvePrefixedUri(String value) { + if (value == null || value.isBlank()) { + return null; + } + + if (value.startsWith("http://") || value.startsWith("https://")) { + return value; + } + + int separatorIndex = value.indexOf(':'); + if (separatorIndex <= 0 || separatorIndex == value.length() - 1) { + return value; + } + + String prefix = value.substring(0, separatorIndex); + String localName = value.substring(separatorIndex + 1); + + return switch (prefix) { + case "om" -> "https://open-metadata.org/ontology/" + localName; + case "rdf" -> "http://www.w3.org/1999/02/22-rdf-syntax-ns#" + localName; + case "rdfs" -> "http://www.w3.org/2000/01/rdf-schema#" + localName; + case "dcat" -> "http://www.w3.org/ns/dcat#" + localName; + case "prov" -> "http://www.w3.org/ns/prov#" + localName; + case "foaf" -> "http://xmlns.com/foaf/0.1/" + localName; + case "skos" -> "http://www.w3.org/2004/02/skos/core#" + localName; + case "dprod" -> "https://ekgf.github.io/dprod/" + localName; + default -> value; + }; + } + + public static String normalizeEntityGraphExportFormat(String format) { + if (format == null || format.isBlank()) { + return "TURTLE"; + } + + return switch (format.trim().toLowerCase(Locale.ROOT)) { + case "jsonld", "json-ld" -> "JSON-LD"; + case "turtle", "ttl" -> "TURTLE"; + default -> throw new IllegalArgumentException("Unsupported export format"); + }; + } + + private String convertEdgesToGraphData( + String rootUri, + Set nodeUris, + List edges, + List entityTypeOptions, + List relationshipTypeOptions) { com.fasterxml.jackson.databind.node.ObjectNode graphData = JsonUtils.getObjectMapper().createObjectNode(); com.fasterxml.jackson.databind.node.ArrayNode nodes = JsonUtils.getObjectMapper().createArrayNode(); com.fasterxml.jackson.databind.node.ArrayNode graphEdges = JsonUtils.getObjectMapper().createArrayNode(); + com.fasterxml.jackson.databind.node.ObjectNode filterOptions = + JsonUtils.getObjectMapper().createObjectNode(); + com.fasterxml.jackson.databind.node.ArrayNode entityTypeFilterOptions = + JsonUtils.getObjectMapper().createArrayNode(); + com.fasterxml.jackson.databind.node.ArrayNode relationshipTypeFilterOptions = + JsonUtils.getObjectMapper().createArrayNode(); - Set addedNodes = new HashSet<>(); Map nodeMap = new HashMap<>(); - for (EdgeInfo edge : edges) { - String fromUri = edge.fromUri; - String toUri = edge.toUri; - - if (!addedNodes.contains(fromUri)) { - com.fasterxml.jackson.databind.node.ObjectNode fromNode = createNodeFromUri(fromUri); - nodes.add(fromNode); - nodeMap.put(fromUri, fromNode); - addedNodes.add(fromUri); - } + List orderedNodeUris = + nodeUris.stream() + .sorted( + Comparator.comparing((String uri) -> !rootUri.equals(uri)) + .thenComparing(this::extractEntityTypeFromUri) + .thenComparing(uri -> uri)) + .toList(); - if (!addedNodes.contains(toUri)) { - com.fasterxml.jackson.databind.node.ObjectNode toNode = createNodeFromUri(toUri); - nodes.add(toNode); - nodeMap.put(toUri, toNode); - addedNodes.add(toUri); - } + for (String nodeUri : orderedNodeUris) { + com.fasterxml.jackson.databind.node.ObjectNode node = createNodeFromUri(nodeUri); + nodes.add(node); + nodeMap.put(nodeUri, node); + } + for (EdgeInfo edge : edges) { com.fasterxml.jackson.databind.node.ObjectNode graphEdge = JsonUtils.getObjectMapper().createObjectNode(); - graphEdge.put("from", fromUri); - graphEdge.put("to", toUri); + graphEdge.put("from", edge.fromUri); + graphEdge.put("to", edge.toUri); graphEdge.put("label", formatRelationshipLabel(edge.relation)); + graphEdge.put("relationType", edge.relation); graphEdge.put("arrows", "to"); graphEdges.add(graphEdge); } + for (FilterOptionInfo filterOption : entityTypeOptions) { + entityTypeFilterOptions.add(createFilterOptionNode(filterOption)); + } + + for (FilterOptionInfo filterOption : relationshipTypeOptions) { + relationshipTypeFilterOptions.add(createFilterOptionNode(filterOption)); + } + enhanceNodesWithEntityDetails(nodeMap); graphData.set("nodes", nodes); graphData.set("edges", graphEdges); + graphData.put("totalNodes", nodes.size()); + graphData.put("totalEdges", graphEdges.size()); + graphData.put("source", "rdf"); + + filterOptions.set("entityTypes", entityTypeFilterOptions); + filterOptions.set("relationshipTypes", relationshipTypeFilterOptions); + graphData.set("filterOptions", filterOptions); return JsonUtils.pojoToJson(graphData); } + private com.fasterxml.jackson.databind.node.ObjectNode createFilterOptionNode( + FilterOptionInfo filterOption) { + com.fasterxml.jackson.databind.node.ObjectNode option = + JsonUtils.getObjectMapper().createObjectNode(); + option.put("id", filterOption.id()); + option.put("label", filterOption.label()); + option.put("count", filterOption.count()); + return option; + } + + private boolean isEntityUri(String uri) { + if (uri == null || !uri.startsWith(config.getBaseUri().toString() + "entity/")) { + return false; + } + String[] parts = uri.split("/entity/")[1].split("/"); + return parts.length >= 2 && !parts[0].isBlank() && !parts[1].isBlank(); + } + + private String extractEntityRelationType(String predicateUri) { + if (predicateUri == null || predicateUri.isBlank()) { + return null; + } + + String localName = extractUriLocalName(predicateUri); + if (localName == null || localName.isBlank()) { + return null; + } + + String normalized = localName.replaceAll("[^A-Za-z0-9]", "").toLowerCase(Locale.ROOT); + return switch (normalized) { + case "used" -> "uses"; + case "wasderivedfrom", "upstream" -> "upstream"; + case "wasinfluencedby", "downstream" -> "downstream"; + case "wasgeneratedby" -> "processedBy"; + default -> toCanonicalIdentifier(localName); + }; + } + + private String normalizeEntityTypeFilter(String entityType) { + return entityType == null ? "" : entityType.trim().toLowerCase(Locale.ROOT); + } + + private String normalizeRelationTypeFilter(String relationType) { + String canonical = toCanonicalIdentifier(relationType); + return canonical == null ? "" : canonical.toLowerCase(Locale.ROOT); + } + + private String extractUriLocalName(String uri) { + if (uri == null || uri.isBlank()) { + return null; + } + if (uri.contains("#")) { + return uri.substring(uri.lastIndexOf('#') + 1); + } + if (uri.contains("/")) { + return uri.substring(uri.lastIndexOf('/') + 1); + } + return uri; + } + + private String toCanonicalIdentifier(String value) { + String localName = extractUriLocalName(value); + if (localName == null || localName.isBlank()) { + return null; + } + + if (localName.equals(localName.toUpperCase(Locale.ROOT))) { + return localName.toLowerCase(Locale.ROOT); + } + + if (localName.matches("[a-z]+([A-Z][a-z0-9]+)+")) { + return Character.toLowerCase(localName.charAt(0)) + localName.substring(1); + } + + String spaced = + localName.replaceAll("([a-z0-9])([A-Z])", "$1 $2").replace('_', ' ').replace('-', ' '); + String[] parts = spaced.trim().split("\\s+"); + if (parts.length == 0) { + return localName; + } + + StringBuilder builder = new StringBuilder(parts[0].toLowerCase(Locale.ROOT)); + for (int i = 1; i < parts.length; i++) { + if (parts[i].isBlank()) { + continue; + } + String normalizedPart = parts[i].toLowerCase(Locale.ROOT); + builder + .append(Character.toUpperCase(normalizedPart.charAt(0))) + .append(normalizedPart.substring(1)); + } + return builder.toString(); + } + + private record EntityGraphTraversalResult(Set nodeUris, List edges) {} + + private record FilteredEntityGraph(Set nodeUris, List edges) {} + + private record FilterOptionInfo(String id, String label, int count) {} + public void executeSparqlUpdate(String update) { if (!isEnabled()) { throw new IllegalStateException("RDF not enabled"); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfPropertyMapper.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfPropertyMapper.java index b69dd1a3d363..810b6cb5d3e5 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfPropertyMapper.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfPropertyMapper.java @@ -43,7 +43,10 @@ public class RdfPropertyMapper { // Properties that should be mapped to structured RDF instead of JSON literals private static final Set STRUCTURED_PROPERTIES = - Set.of("changeDescription", "votes", "lifeCycle", "customProperties", "extension"); + Set.of("votes", "lifeCycle", "customProperties", "extension"); + + // Properties that should be omitted from RDF because they are audit/helper data. + private static final Set IGNORED_PROPERTIES = Set.of("changeDescription"); // Lineage properties that need special handling private static final Set LINEAGE_PROPERTIES = @@ -119,7 +122,7 @@ private void processContextMappings( // Skip fields that are handled separately with typed predicates // (e.g., relatedTerms which use typed relations like broader, synonym, etc.) - if (TYPED_RELATION_FIELDS.contains(fieldName)) { + if (TYPED_RELATION_FIELDS.contains(fieldName) || IGNORED_PROPERTIES.contains(fieldName)) { continue; } @@ -400,7 +403,6 @@ private UUID tryResolveGlossaryTermIdFromHref(JsonNode tagLabel) { private void addStructuredProperty( String fieldName, JsonNode value, Resource entityResource, Model model) { switch (fieldName) { - case "changeDescription" -> addChangeDescription(value, entityResource, model); case "votes" -> addVotes(value, entityResource, model); case "lifeCycle" -> addLifeCycle(value, entityResource, model); case "extension" -> addExtension(value, entityResource, model); @@ -421,102 +423,9 @@ private void addStructuredArrayProperty( } } - /** - * Converts ChangeDescription to structured RDF triples. Enables SPARQL queries like: "Find all - * entities where description was changed by user X after date Y" - * - *

Structure: entity -> om:hasChangeDescription -> _:changeNode _:changeNode a - * om:ChangeDescription _:changeNode om:previousVersion "1.0" _:changeNode om:fieldsAdded -> - * _:fieldChange1 - */ - private void addChangeDescription(JsonNode changeDesc, Resource entityResource, Model model) { - if (changeDesc == null || changeDesc.isNull()) { - return; - } - - // Create a blank node for the change description - String changeNodeUri = - baseUri + "change/" + entityResource.getLocalName() + "/" + UUID.randomUUID(); - Resource changeNode = model.createResource(changeNodeUri); - - // Link entity to change description - Property hasChangeDesc = model.createProperty(OM_NS, "hasChangeDescription"); - entityResource.addProperty(hasChangeDesc, changeNode); - - // Add type - changeNode.addProperty(RDF.type, model.createResource(OM_NS + "ChangeDescription")); - - // Add previous version - if (changeDesc.has("previousVersion")) { - changeNode.addProperty( - model.createProperty(OM_NS, "previousVersion"), - model.createTypedLiteral(changeDesc.get("previousVersion").asDouble())); - } - - // Add fields added - if (changeDesc.has("fieldsAdded") && changeDesc.get("fieldsAdded").isArray()) { - addFieldChanges( - changeDesc.get("fieldsAdded"), changeNode, "fieldsAdded", entityResource, model); - } - - // Add fields updated - if (changeDesc.has("fieldsUpdated") && changeDesc.get("fieldsUpdated").isArray()) { - addFieldChanges( - changeDesc.get("fieldsUpdated"), changeNode, "fieldsUpdated", entityResource, model); - } - - // Add fields deleted - if (changeDesc.has("fieldsDeleted") && changeDesc.get("fieldsDeleted").isArray()) { - addFieldChanges( - changeDesc.get("fieldsDeleted"), changeNode, "fieldsDeleted", entityResource, model); - } - } - - /** - * Adds field change details as structured RDF - */ - private void addFieldChanges( - JsonNode fieldsArray, - Resource changeNode, - String changeType, - Resource entityResource, - Model model) { - Property changeProp = model.createProperty(OM_NS, changeType); - - for (JsonNode fieldChange : fieldsArray) { - // Create a blank node for each field change - String fieldChangeUri = - baseUri + "fieldChange/" + entityResource.getLocalName() + "/" + UUID.randomUUID(); - Resource fieldChangeNode = model.createResource(fieldChangeUri); - - changeNode.addProperty(changeProp, fieldChangeNode); - fieldChangeNode.addProperty(RDF.type, model.createResource(OM_NS + "FieldChange")); - - // Add field name - if (fieldChange.has("name")) { - fieldChangeNode.addProperty( - model.createProperty(OM_NS, "fieldName"), fieldChange.get("name").asText()); - } - - // Add old value (as string representation for queryability) - if (fieldChange.has("oldValue") && !fieldChange.get("oldValue").isNull()) { - JsonNode oldVal = fieldChange.get("oldValue"); - String oldValueStr = oldVal.isTextual() ? oldVal.asText() : oldVal.toString(); - fieldChangeNode.addProperty(model.createProperty(OM_NS, "oldValue"), oldValueStr); - } - - // Add new value (as string representation for queryability) - if (fieldChange.has("newValue") && !fieldChange.get("newValue").isNull()) { - JsonNode newVal = fieldChange.get("newValue"); - String newValueStr = newVal.isTextual() ? newVal.asText() : newVal.toString(); - fieldChangeNode.addProperty(model.createProperty(OM_NS, "newValue"), newValueStr); - } - } - } - /** * Converts Votes to structured RDF triples. Enables SPARQL queries like: "Find all entities with - * more than 10 upvotes" or "Find entities upvoted by user X" + * more than 10 upvotes" without exposing individual voter identities as graph edges. */ private void addVotes(JsonNode votes, Resource entityResource, Model model) { if (votes == null || votes.isNull()) { @@ -547,30 +456,6 @@ private void addVotes(JsonNode votes, Resource entityResource, Model model) { model.createProperty(OM_NS, "downVotes"), model.createTypedLiteral(votes.get("downVotes").asInt())); } - - // Add upVoters as entity references - if (votes.has("upVoters") && votes.get("upVoters").isArray()) { - Property upVotersProp = model.createProperty(OM_NS, "upVoters"); - for (JsonNode voter : votes.get("upVoters")) { - if (voter.has("id") && voter.has("type")) { - String voterUri = - baseUri + "entity/" + voter.get("type").asText() + "/" + voter.get("id").asText(); - votesNode.addProperty(upVotersProp, model.createResource(voterUri)); - } - } - } - - // Add downVoters as entity references - if (votes.has("downVoters") && votes.get("downVoters").isArray()) { - Property downVotersProp = model.createProperty(OM_NS, "downVoters"); - for (JsonNode voter : votes.get("downVoters")) { - if (voter.has("id") && voter.has("type")) { - String voterUri = - baseUri + "entity/" + voter.get("type").asText() + "/" + voter.get("id").asText(); - votesNode.addProperty(downVotersProp, model.createResource(voterUri)); - } - } - } } /** diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfResource.java index 56e7bbb379df..a07e26176622 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfResource.java @@ -20,10 +20,15 @@ import jakarta.ws.rs.core.SecurityContext; import jakarta.ws.rs.core.UriInfo; import java.io.IOException; +import java.util.Arrays; +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.Set; import java.util.UUID; import javax.validation.constraints.NotEmpty; import lombok.extern.slf4j.Slf4j; import org.openmetadata.schema.api.rdf.SparqlQuery; +import org.openmetadata.schema.utils.JsonUtils; import org.openmetadata.service.Entity; import org.openmetadata.service.OpenMetadataApplicationConfig; import org.openmetadata.service.rdf.RdfRepository; @@ -39,6 +44,8 @@ @Slf4j public class RdfResource { public static final String COLLECTION_PATH = "/v1/rdf"; + private static final int MIN_GRAPH_DEPTH = 1; + private static final int MAX_GRAPH_DEPTH = 5; private volatile RdfRepository rdfRepository; private final Authorizer authorizer; private volatile SemanticSearchEngine semanticSearchEngine; @@ -240,24 +247,148 @@ public Response exploreEntityGraph( @Parameter(description = "Depth of relationships to explore") @QueryParam("depth") @DefaultValue("2") - int depth) { + int depth, + @Parameter(description = "Comma-separated entity types to keep in the graph") + @QueryParam("entityTypes") + String entityTypes, + @Parameter(description = "Comma-separated relationship types to keep in the graph") + @QueryParam("relationshipTypes") + String relationshipTypes) { authorizer.authorizeAdmin(securityContext); try { + String validatedEntityType = validateEntityType(entityType); + int clampedDepth = clampGraphDepth(depth); if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { return Response.status(Response.Status.SERVICE_UNAVAILABLE) - .entity("{\"error\": \"RDF service not enabled\"}") + .entity(buildErrorResponse("RDF service not enabled")) .build(); } - String graphData = getRdfRepository().getEntityGraph(entityId, entityType, depth); + String graphData = + getRdfRepository() + .getEntityGraph( + entityId, + validatedEntityType, + clampedDepth, + parseCsvFilter(entityTypes), + parseCsvFilter(relationshipTypes)); return Response.ok(graphData, MediaType.APPLICATION_JSON).build(); - + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .build(); } catch (Exception e) { LOG.error("Error exploring entity graph", e); - return Response.serverError().entity("{\"error\": \"An internal error occurred\"}").build(); + return Response.serverError() + .entity(buildErrorResponse("An internal error occurred")) + .build(); + } + } + + @GET + @Path("/graph/explore/export") + @Produces({JSON_LD, TURTLE, MediaType.APPLICATION_JSON}) + @Operation( + operationId = "exportEntityGraph", + summary = "Export explored entity graph", + description = "Export the currently explored entity graph in Turtle or JSON-LD format", + responses = { + @ApiResponse(responseCode = "200", description = "Entity graph exported successfully"), + @ApiResponse(responseCode = "400", description = "Invalid request"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response exportEntityGraph( + @Context SecurityContext securityContext, + @Parameter(description = "Entity ID", required = true) @QueryParam("entityId") UUID entityId, + @Parameter(description = "Entity type", required = true) @QueryParam("entityType") + String entityType, + @Parameter(description = "Depth of relationships to explore") + @QueryParam("depth") + @DefaultValue("2") + int depth, + @Parameter(description = "Comma-separated entity types to keep in the graph") + @QueryParam("entityTypes") + String entityTypes, + @Parameter(description = "Comma-separated relationship types to keep in the graph") + @QueryParam("relationshipTypes") + String relationshipTypes, + @Parameter(description = "Export format: turtle or jsonld") + @QueryParam("format") + @DefaultValue("turtle") + String format) { + authorizer.authorizeAdmin(securityContext); + try { + String validatedEntityType = validateEntityType(entityType); + int clampedDepth = clampGraphDepth(depth); + String normalizedFormat = RdfRepository.normalizeEntityGraphExportFormat(format); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF service not enabled")) + .build(); + } + + String result = + getRdfRepository() + .exportEntityGraph( + entityId, + validatedEntityType, + clampedDepth, + parseCsvFilter(entityTypes), + parseCsvFilter(relationshipTypes), + normalizedFormat); + + MediaType mediaType = + switch (normalizedFormat) { + case "JSON-LD" -> MediaType.valueOf(JSON_LD); + default -> MediaType.valueOf(TURTLE); + }; + + return Response.ok(result, mediaType).build(); + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .build(); + } catch (Exception e) { + LOG.error("Error exporting entity graph", e); + return Response.serverError() + .entity(buildErrorResponse("An internal error occurred")) + .build(); } } + private String validateEntityType(String entityType) { + if (entityType == null || entityType.isBlank()) { + throw new IllegalArgumentException("Entity type is required"); + } + + String trimmedEntityType = entityType.trim(); + if (!trimmedEntityType.matches("[A-Za-z][A-Za-z0-9]*") + || !Entity.hasEntityRepository(trimmedEntityType)) { + throw new IllegalArgumentException("Invalid entity type"); + } + + return trimmedEntityType; + } + + private String buildErrorResponse(String message) { + return JsonUtils.pojoToJson(Map.of("error", message)); + } + + private int clampGraphDepth(int depth) { + return Math.min(Math.max(depth, MIN_GRAPH_DEPTH), MAX_GRAPH_DEPTH); + } + + private Set parseCsvFilter(String values) { + if (values == null || values.isBlank()) { + return Set.of(); + } + + return Arrays.stream(values.split(",")) + .map(String::trim) + .filter(value -> !value.isEmpty()) + .collect(java.util.stream.Collectors.toCollection(LinkedHashSet::new)); + } + @GET @Path("/sparql") @Operation( @@ -430,25 +561,35 @@ public Response getFullLineage( String direction) { authorizer.authorizeAdmin(securityContext); try { - String query = buildLineageQuery(entityId, entityType, direction); - String results = - getRdfRepository() != null - ? getRdfRepository().executeSparqlQueryWithInference(query, SPARQL_JSON, "custom") - : null; - if (results == null) { + String validatedEntityType = validateEntityType(entityType); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { return Response.status(Response.Status.SERVICE_UNAVAILABLE) - .entity("{\"error\": \"RDF service not enabled\"}") + .entity(buildErrorResponse("RDF service not enabled")) .build(); } + + String query = + buildLineageQuery( + entityId, validatedEntityType, direction, getRdfRepository().getBaseUri()); + String results = + getRdfRepository().executeSparqlQueryWithInference(query, SPARQL_JSON, "custom"); return Response.ok(results).build(); + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .build(); } catch (Exception e) { LOG.error("Error getting lineage with inference", e); - return Response.serverError().entity("{\"error\": \"An internal error occurred\"}").build(); + return Response.serverError() + .entity(buildErrorResponse("An internal error occurred")) + .build(); } } - private String buildLineageQuery(UUID entityId, String entityType, String direction) { - String entityUri = "https://open-metadata.org/entity/" + entityType + "/" + entityId; + private String buildLineageQuery( + UUID entityId, String entityType, String direction, String baseUri) { + String normalizedBaseUri = baseUri.endsWith("/") ? baseUri : baseUri + "/"; + String entityUri = normalizedBaseUri + "entity/" + entityType + "/" + entityId; return switch (direction.toLowerCase()) { case "upstream" -> String.format( diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/ODCSConverter.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/ODCSConverter.java index 3be0a07736f2..a8439f264e7f 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/ODCSConverter.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/util/ODCSConverter.java @@ -330,8 +330,8 @@ private static ODCSDataContract.OdcsStatus mapContractStatusToODCS(EntityStatus if (status == null) return ODCSDataContract.OdcsStatus.DRAFT; return switch (status) { case APPROVED -> ODCSDataContract.OdcsStatus.ACTIVE; - case DEPRECATED -> ODCSDataContract.OdcsStatus.DEPRECATED; case ARCHIVED -> ODCSDataContract.OdcsStatus.RETIRED; + case DEPRECATED -> ODCSDataContract.OdcsStatus.DEPRECATED; case DRAFT, IN_REVIEW, REJECTED, UNPROCESSED -> ODCSDataContract.OdcsStatus.DRAFT; }; } diff --git a/openmetadata-service/src/main/resources/json/data/app/RdfIndexApp.json b/openmetadata-service/src/main/resources/json/data/app/RdfIndexApp.json index 3695ec8e493a..c351ed8e8a98 100644 --- a/openmetadata-service/src/main/resources/json/data/app/RdfIndexApp.json +++ b/openmetadata-service/src/main/resources/json/data/app/RdfIndexApp.json @@ -2,15 +2,18 @@ "name": "RdfIndexApp", "displayName": "RDF Knowledge Graph Indexing", "appConfiguration": { - "entities": [ - "all" - ], + "entities": [], "recreateIndex": false, - "batchSize": "100" + "batchSize": 100, + "producerThreads": 2, + "consumerThreads": 3, + "queueSize": 5000, + "useDistributedIndexing": true, + "partitionSize": 10000 }, "appSchedule": { "scheduleTimeline": "Custom", "cronExpression": "0 0 * * *" }, "supportsInterrupt": true -} \ No newline at end of file +} diff --git a/openmetadata-service/src/main/resources/json/data/appMarketPlaceDefinition/RdfIndexApp.json b/openmetadata-service/src/main/resources/json/data/appMarketPlaceDefinition/RdfIndexApp.json index 74ee1963e36e..d825db2b47cf 100644 --- a/openmetadata-service/src/main/resources/json/data/appMarketPlaceDefinition/RdfIndexApp.json +++ b/openmetadata-service/src/main/resources/json/data/appMarketPlaceDefinition/RdfIndexApp.json @@ -12,15 +12,19 @@ "scheduleType": "ScheduledOrManual", "permission": "All", "className": "org.openmetadata.service.apps.bundles.rdf.RdfIndexApp", + "allowConfiguration": true, "runtime": { "enabled": true }, "supportsInterrupt": true, "appConfiguration": { - "entities": [ - "all" - ], + "entities": [], "recreateIndex": false, - "batchSize": "100" + "batchSize": 100, + "producerThreads": 2, + "consumerThreads": 3, + "queueSize": 5000, + "useDistributedIndexing": true, + "partitionSize": 10000 } -} \ No newline at end of file +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/RdfIndexAppTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/RdfIndexAppTest.java index 3476da14b28f..8e6d21ef4e56 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/RdfIndexAppTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/RdfIndexAppTest.java @@ -3,10 +3,13 @@ import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; +import com.fasterxml.jackson.core.type.TypeReference; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.UUID; +import java.util.concurrent.CompletableFuture; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; @@ -18,14 +21,28 @@ import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.entity.app.AppRunRecord; import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; import org.openmetadata.schema.type.Include; import org.openmetadata.schema.type.Relationship; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.rdf.distributed.RdfIndexJob; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; import org.openmetadata.service.jdbi3.CollectionDAO; import org.openmetadata.service.jdbi3.CollectionDAO.EntityRelationshipDAO; import org.openmetadata.service.jdbi3.CollectionDAO.EntityRelationshipObject; +import org.openmetadata.service.jdbi3.EntityDAO; +import org.openmetadata.service.jdbi3.EntityRepository; import org.openmetadata.service.rdf.RdfRepository; import org.openmetadata.service.search.SearchRepository; +import org.quartz.JobDataMap; +import org.quartz.JobDetail; +import org.quartz.JobExecutionContext; +import org.quartz.JobKey; @ExtendWith(MockitoExtension.class) @DisplayName("RdfIndexApp Tests") @@ -39,6 +56,30 @@ class RdfIndexAppTest { private static RdfRepository mockRdfRepository; private RdfIndexApp rdfIndexApp; + private static class TestableRdfIndexApp extends RdfIndexApp { + private AppRunRecord appRunRecord; + private JobExecutionContext pushedContext; + private AppRunRecord pushedRecord; + private boolean pushedUpdate; + + TestableRdfIndexApp(CollectionDAO collectionDAO, SearchRepository searchRepository) { + super(collectionDAO, searchRepository); + } + + @Override + protected AppRunRecord getJobRecord(JobExecutionContext jobExecutionContext) { + return appRunRecord; + } + + @Override + protected void pushAppStatusUpdates( + JobExecutionContext jobExecutionContext, AppRunRecord appRecord, boolean update) { + this.pushedContext = jobExecutionContext; + this.pushedRecord = appRecord; + this.pushedUpdate = update; + } + } + @BeforeAll static void setUpClass() { mockRdfRepository = mock(RdfRepository.class); @@ -56,6 +97,8 @@ static void tearDownClass() { @BeforeEach void setUp() { lenient().when(collectionDAO.relationshipDAO()).thenReturn(relationshipDAO); + lenient().when(mockRdfRepository.isEnabled()).thenReturn(true); + clearInvocations(mockRdfRepository); rdfIndexApp = new RdfIndexApp(collectionDAO, searchRepository); } @@ -383,6 +426,240 @@ void testValidateValidConfig() { } } + @Nested + @DisplayName("Distributed Job Serialization Tests") + class DistributedJobSerializationTests { + + @Test + @DisplayName("Should deserialize distributed RDF entity stats from JSON") + void testDistributedEntityStatsJsonRoundTrip() { + Map original = + Map.of( + "app", + RdfIndexJob.EntityTypeStats.builder() + .entityType("app") + .totalRecords(12) + .processedRecords(4) + .successRecords(3) + .failedRecords(1) + .totalPartitions(2) + .completedPartitions(1) + .failedPartitions(0) + .build()); + + String json = JsonUtils.pojoToJson(original); + Map roundTrip = + JsonUtils.readValue( + json, new TypeReference>() {}); + + assertEquals(12, roundTrip.get("app").getTotalRecords()); + assertEquals(3, roundTrip.get("app").getSuccessRecords()); + assertEquals(2, roundTrip.get("app").getTotalPartitions()); + } + } + + @Nested + @DisplayName("Entity Selection Tests") + class EntitySelectionTests { + + @Test + @DisplayName("Should skip entity types without repositories") + void testResolveEntityTypesSkipsUnsupportedEntities() throws Exception { + try (MockedStatic entityMock = mockStatic(Entity.class)) { + EntityRepository mockRepository = mock(EntityRepository.class); + entityMock.when(() -> Entity.getEntityRepository("table")).thenReturn(mockRepository); + entityMock + .when(() -> Entity.getEntityRepository("queryCostRecord")) + .thenThrow(new IllegalStateException("Unsupported entity")); + + var method = RdfIndexApp.class.getDeclaredMethod("resolveEntityTypes", Set.class); + method.setAccessible(true); + + @SuppressWarnings("unchecked") + Set result = + (Set) method.invoke(rdfIndexApp, Set.of("table", "queryCostRecord")); + + assertEquals(Set.of("table"), result); + } + } + + @Test + @DisplayName( + "Should treat null or empty entity selection as all supported repository-backed entities") + void testResolveEntityTypesDefaultsEmptySelectionToAllSupportedEntities() throws Exception { + try (MockedStatic entityMock = mockStatic(Entity.class)) { + EntityRepository mockRepository = mock(EntityRepository.class); + entityMock.when(Entity::getEntityList).thenReturn(Set.of("table", "queryCostRecord")); + entityMock.when(() -> Entity.getEntityRepository("table")).thenReturn(mockRepository); + entityMock + .when(() -> Entity.getEntityRepository("queryCostRecord")) + .thenThrow(new IllegalStateException("Unsupported entity")); + + var method = RdfIndexApp.class.getDeclaredMethod("resolveEntityTypes", Set.class); + method.setAccessible(true); + + @SuppressWarnings("unchecked") + Set nullSelection = (Set) method.invoke(rdfIndexApp, new Object[] {null}); + @SuppressWarnings("unchecked") + Set emptySelection = (Set) method.invoke(rdfIndexApp, Set.of()); + + assertEquals(Set.of("table"), nullSelection); + assertEquals(Set.of("table"), emptySelection); + } + } + } + + @Nested + @DisplayName("Run Record Update Tests") + class RunRecordUpdateTests { + + @Test + @DisplayName("Should persist updated app run record state during execution") + void testUpdateRecordToDbAndNotifyPersistsRunRecord() throws Exception { + TestableRdfIndexApp testApp = new TestableRdfIndexApp(collectionDAO, searchRepository); + + EventPublisherJob jobConfig = new EventPublisherJob(); + jobConfig.setStatus(EventPublisherJob.Status.FAILED); + jobConfig.setFailure( + new IndexingError().withMessage("distributed rdf job initialization failed")); + Stats stats = new Stats(); + stats.setJobStats(new StepStats().withTotalRecords(5).withFailedRecords(5)); + jobConfig.setStats(stats); + + var jobDataField = RdfIndexApp.class.getDeclaredField("jobData"); + jobDataField.setAccessible(true); + jobDataField.set(testApp, jobConfig); + + JobExecutionContext context = mock(JobExecutionContext.class); + testApp.appRunRecord = new AppRunRecord().withStatus(AppRunRecord.Status.RUNNING); + + testApp.updateRecordToDbAndNotify(context); + + assertEquals(AppRunRecord.Status.FAILED, testApp.appRunRecord.getStatus()); + assertNotNull(testApp.appRunRecord.getFailureContext()); + assertNotNull(testApp.appRunRecord.getSuccessContext()); + assertSame(context, testApp.pushedContext); + assertSame(testApp.appRunRecord, testApp.pushedRecord); + assertTrue(testApp.pushedUpdate); + } + + @Test + @DisplayName("Should publish distributed RDF progress before job completion") + void testMonitorDistributedJobPublishesRunningProgress() throws Exception { + TestableRdfIndexApp testApp = new TestableRdfIndexApp(collectionDAO, searchRepository); + + EventPublisherJob jobConfig = new EventPublisherJob(); + jobConfig.setStatus(EventPublisherJob.Status.RUNNING); + + var jobDataField = RdfIndexApp.class.getDeclaredField("jobData"); + jobDataField.setAccessible(true); + jobDataField.set(testApp, jobConfig); + + JobExecutionContext context = mock(JobExecutionContext.class); + JobDetail jobDetail = mock(JobDetail.class); + JobDataMap jobDataMap = new JobDataMap(); + when(context.getJobDetail()).thenReturn(jobDetail); + when(jobDetail.getJobDataMap()).thenReturn(jobDataMap); + + var jobExecutionContextField = RdfIndexApp.class.getDeclaredField("jobExecutionContext"); + jobExecutionContextField.setAccessible(true); + jobExecutionContextField.set(testApp, context); + + testApp.appRunRecord = new AppRunRecord().withStatus(AppRunRecord.Status.RUNNING); + + var distributedExecutorField = RdfIndexApp.class.getDeclaredField("distributedExecutor"); + distributedExecutorField.setAccessible(true); + var mockDistributedExecutor = + mock( + org.openmetadata.service.apps.bundles.rdf.distributed.DistributedRdfIndexExecutor + .class); + when(mockDistributedExecutor.getJobWithFreshStats()) + .thenReturn( + RdfIndexJob.builder() + .id(UUID.randomUUID()) + .status( + org.openmetadata + .service + .apps + .bundles + .searchIndex + .distributed + .IndexJobStatus + .RUNNING) + .totalRecords(10) + .processedRecords(7) + .successRecords(7) + .failedRecords(0) + .build()); + distributedExecutorField.set(testApp, mockDistributedExecutor); + + var method = + RdfIndexApp.class.getDeclaredMethod( + "monitorDistributedJob", UUID.class, java.util.concurrent.Future.class); + method.setAccessible(true); + method.invoke(testApp, UUID.randomUUID(), CompletableFuture.completedFuture(null)); + + assertNotNull(testApp.getJobData().getStats()); + assertEquals(10, testApp.getJobData().getStats().getJobStats().getTotalRecords()); + assertEquals(7, testApp.getJobData().getStats().getJobStats().getSuccessRecords()); + assertSame(context, testApp.pushedContext); + assertNotNull(testApp.appRunRecord.getSuccessContext()); + assertEquals(AppRunRecord.Status.RUNNING, testApp.appRunRecord.getStatus()); + } + + @Test + @DisplayName("Should clear RDF data when recreateIndex is enabled before distributed indexing") + void testExecuteClearsRdfDataWhenRecreateIndexEnabled() throws Exception { + TestableRdfIndexApp testApp = new TestableRdfIndexApp(collectionDAO, searchRepository); + testApp.appRunRecord = new AppRunRecord().withStatus(AppRunRecord.Status.RUNNING); + + EventPublisherJob jobConfig = new EventPublisherJob(); + jobConfig.setEntities(Set.of("table")); + jobConfig.setRecreateIndex(true); + jobConfig.setUseDistributedIndexing(true); + jobConfig.setStatus(EventPublisherJob.Status.STARTED); + + var jobDataField = RdfIndexApp.class.getDeclaredField("jobData"); + jobDataField.setAccessible(true); + jobDataField.set(testApp, jobConfig); + + @SuppressWarnings("unchecked") + EntityRepository repository = mock(EntityRepository.class); + @SuppressWarnings("unchecked") + EntityDAO entityDAO = mock(EntityDAO.class); + when(repository.getDao()).thenReturn(entityDAO); + when(entityDAO.listTotalCount()).thenReturn(0); + + JobExecutionContext context = mock(JobExecutionContext.class); + JobDetail jobDetail = mock(JobDetail.class); + JobDataMap jobDataMap = new JobDataMap(); + when(context.getJobDetail()).thenReturn(jobDetail); + when(jobDetail.getJobDataMap()).thenReturn(jobDataMap); + when(jobDetail.getKey()).thenReturn(JobKey.jobKey("rdf-index-test")); + + RdfIndexJob completedJob = + RdfIndexJob.builder().id(UUID.randomUUID()).status(IndexJobStatus.COMPLETED).build(); + + try (MockedStatic entityMock = mockStatic(Entity.class); + var ignored = + mockConstruction( + org.openmetadata.service.apps.bundles.rdf.distributed.DistributedRdfIndexExecutor + .class, + (mock, mockContext) -> { + when(mock.createJob(eq(Set.of("table")), eq(jobConfig), anyString())) + .thenReturn(completedJob); + when(mock.getJobWithFreshStats()).thenReturn(completedJob); + })) { + entityMock.when(() -> Entity.getEntityRepository("table")).thenReturn(repository); + + testApp.execute(context); + } + + verify(mockRdfRepository).clearAll(); + assertEquals(EventPublisherJob.Status.COMPLETED, jobConfig.getStatus()); + } + } + @Nested @DisplayName("IndexingTask Record Tests") class IndexingTaskTests { @@ -573,6 +850,72 @@ void testProcessBatchRelationshipsHandlesLineageWithDetails() throws Exception { verify(mockRdfRepository) .addLineageWithDetails(eq("table"), eq(fromId), eq("table"), eq(toId), any()); } + + @Test + @DisplayName("Should skip event subscription relationships using canonical entity type") + void testProcessBatchRelationshipsSkipsEventSubscriptionEntities() throws Exception { + List mockEntities = new ArrayList<>(); + EntityInterface mockEntity = mock(EntityInterface.class); + UUID entityId = UUID.randomUUID(); + when(mockEntity.getId()).thenReturn(entityId); + mockEntities.add(mockEntity); + + List mockRelationships = new ArrayList<>(); + mockRelationships.add( + EntityRelationshipObject.builder() + .fromId(entityId.toString()) + .toId(UUID.randomUUID().toString()) + .fromEntity("table") + .toEntity(Entity.EVENT_SUBSCRIPTION) + .relation(Relationship.HAS.ordinal()) + .build()); + + when(relationshipDAO.findToBatchWithRelations(anyList(), anyString(), anyList())) + .thenReturn(mockRelationships); + when(relationshipDAO.findFromBatch(anyList(), anyInt(), any(Include.class))) + .thenReturn(new ArrayList<>()); + + var method = + RdfIndexApp.class.getDeclaredMethod( + "processBatchRelationships", String.class, List.class); + method.setAccessible(true); + method.invoke(rdfIndexApp, "table", mockEntities); + + verifyNoInteractions(mockRdfRepository); + } + + @Test + @DisplayName("Should skip event subscription relationships using legacy camelCase entity type") + void testProcessBatchRelationshipsSkipsLegacyEventSubscriptionEntities() throws Exception { + List mockEntities = new ArrayList<>(); + EntityInterface mockEntity = mock(EntityInterface.class); + UUID entityId = UUID.randomUUID(); + when(mockEntity.getId()).thenReturn(entityId); + mockEntities.add(mockEntity); + + List mockRelationships = new ArrayList<>(); + mockRelationships.add( + EntityRelationshipObject.builder() + .fromId(entityId.toString()) + .toId(UUID.randomUUID().toString()) + .fromEntity("table") + .toEntity("eventSubscription") + .relation(Relationship.HAS.ordinal()) + .build()); + + when(relationshipDAO.findToBatchWithRelations(anyList(), anyString(), anyList())) + .thenReturn(mockRelationships); + when(relationshipDAO.findFromBatch(anyList(), anyInt(), any(Include.class))) + .thenReturn(new ArrayList<>()); + + var method = + RdfIndexApp.class.getDeclaredMethod( + "processBatchRelationships", String.class, List.class); + method.setAccessible(true); + method.invoke(rdfIndexApp, "table", mockEntities); + + verifyNoInteractions(mockRdfRepository); + } } @Nested diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexCoordinatorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexCoordinatorTest.java new file mode 100644 index 000000000000..243b085893a4 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexCoordinatorTest.java @@ -0,0 +1,416 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.junit.jupiter.MockitoExtension; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.ServerIdentityResolver; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexJobDAO; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexJobDAO.RdfIndexJobRecord; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexPartitionDAO; +import org.openmetadata.service.jdbi3.CollectionDAO.RdfIndexPartitionDAO.RdfAggregatedStatsRecord; + +@ExtendWith(MockitoExtension.class) +class DistributedRdfIndexCoordinatorTest { + + private static final String TEST_SERVER_ID = "rdf-test-server"; + + @Mock private CollectionDAO collectionDAO; + @Mock private RdfIndexJobDAO jobDAO; + @Mock private RdfIndexPartitionDAO partitionDAO; + @Mock private RdfPartitionCalculator partitionCalculator; + + private DistributedRdfIndexCoordinator coordinator; + private MockedStatic serverIdentityMock; + + @BeforeEach + void setUp() { + ServerIdentityResolver resolver = mock(ServerIdentityResolver.class); + when(resolver.getServerId()).thenReturn(TEST_SERVER_ID); + + serverIdentityMock = mockStatic(ServerIdentityResolver.class); + serverIdentityMock.when(ServerIdentityResolver::getInstance).thenReturn(resolver); + + lenient().when(collectionDAO.rdfIndexJobDAO()).thenReturn(jobDAO); + lenient().when(collectionDAO.rdfIndexPartitionDAO()).thenReturn(partitionDAO); + + coordinator = new DistributedRdfIndexCoordinator(collectionDAO, partitionCalculator); + } + + @AfterEach + void tearDown() { + if (serverIdentityMock != null) { + serverIdentityMock.close(); + } + } + + @Test + void getJobWithAggregatedStatsKeepsCompletedAtNullForNonTerminalJob() { + UUID jobId = UUID.randomUUID(); + EventPublisherJob jobConfiguration = new EventPublisherJob().withEntities(Set.of("table")); + RdfIndexJobRecord jobRecord = + new RdfIndexJobRecord( + jobId.toString(), + IndexJobStatus.READY.name(), + JsonUtils.pojoToJson(jobConfiguration), + 25L, + 0L, + 0L, + 0L, + JsonUtils.pojoToJson( + Map.of( + "table", + RdfIndexJob.EntityTypeStats.builder() + .entityType("table") + .totalRecords(25) + .build())), + "admin", + System.currentTimeMillis(), + null, + null, + System.currentTimeMillis(), + null); + + when(jobDAO.findById(jobId.toString())).thenReturn(jobRecord); + when(partitionDAO.getAggregatedStats(jobId.toString())) + .thenReturn(new RdfAggregatedStatsRecord(25L, 0L, 0L, 0L, 1, 0, 0, 1, 0)); + when(partitionDAO.getEntityStats(jobId.toString())) + .thenReturn( + List.of( + new CollectionDAO.RdfIndexPartitionDAO.RdfEntityStatsRecord( + "table", 25L, 0L, 0L, 0L, 1, 0, 0))); + when(partitionDAO.getServerStats(jobId.toString())).thenReturn(List.of()); + + RdfIndexJob refreshed = coordinator.getJobWithAggregatedStats(jobId); + + assertEquals(IndexJobStatus.RUNNING, refreshed.getStatus()); + assertNull(refreshed.getCompletedAt()); + + verify(jobDAO) + .update( + eq(jobId.toString()), + eq(IndexJobStatus.RUNNING.name()), + eq(0L), + eq(0L), + eq(0L), + anyString(), + isNull(), + isNull(), + anyLong(), + isNull()); + } + + @Test + void getJobWithAggregatedStatsPreservesCompletedAtForTerminalJob() { + UUID jobId = UUID.randomUUID(); + long completedAt = System.currentTimeMillis() - 5000; + EventPublisherJob jobConfiguration = new EventPublisherJob().withEntities(Set.of("table")); + RdfIndexJobRecord jobRecord = + new RdfIndexJobRecord( + jobId.toString(), + IndexJobStatus.RUNNING.name(), + JsonUtils.pojoToJson(jobConfiguration), + 25L, + 25L, + 25L, + 0L, + JsonUtils.pojoToJson( + Map.of( + "table", + RdfIndexJob.EntityTypeStats.builder() + .entityType("table") + .totalRecords(25) + .build())), + "admin", + System.currentTimeMillis(), + System.currentTimeMillis() - 10000, + completedAt, + System.currentTimeMillis(), + null); + + when(jobDAO.findById(jobId.toString())).thenReturn(jobRecord); + when(partitionDAO.getAggregatedStats(jobId.toString())) + .thenReturn(new RdfAggregatedStatsRecord(25L, 25L, 25L, 0L, 1, 1, 0, 0, 0)); + when(partitionDAO.getEntityStats(jobId.toString())) + .thenReturn( + List.of( + new CollectionDAO.RdfIndexPartitionDAO.RdfEntityStatsRecord( + "table", 25L, 25L, 25L, 0L, 1, 1, 0))); + when(partitionDAO.getServerStats(jobId.toString())).thenReturn(List.of()); + + RdfIndexJob refreshed = coordinator.getJobWithAggregatedStats(jobId); + + assertEquals(IndexJobStatus.COMPLETED, refreshed.getStatus()); + assertEquals(completedAt, refreshed.getCompletedAt()); + + verify(jobDAO) + .update( + eq(jobId.toString()), + eq(IndexJobStatus.COMPLETED.name()), + eq(25L), + eq(25L), + eq(0L), + anyString(), + eq(jobRecord.startedAt()), + eq(completedAt), + anyLong(), + isNull()); + } + + @Test + void getJobWithAggregatedStatsSetsCompletedAtWhenTerminalJobWasPreviouslyUnset() { + UUID jobId = UUID.randomUUID(); + EventPublisherJob jobConfiguration = new EventPublisherJob().withEntities(Set.of("table")); + RdfIndexJobRecord jobRecord = + new RdfIndexJobRecord( + jobId.toString(), + IndexJobStatus.RUNNING.name(), + JsonUtils.pojoToJson(jobConfiguration), + 25L, + 25L, + 25L, + 0L, + JsonUtils.pojoToJson( + Map.of( + "table", + RdfIndexJob.EntityTypeStats.builder() + .entityType("table") + .totalRecords(25) + .build())), + "admin", + System.currentTimeMillis(), + System.currentTimeMillis() - 10000, + null, + System.currentTimeMillis(), + null); + + when(jobDAO.findById(jobId.toString())).thenReturn(jobRecord); + when(partitionDAO.getAggregatedStats(jobId.toString())) + .thenReturn(new RdfAggregatedStatsRecord(25L, 25L, 25L, 0L, 1, 1, 0, 0, 0)); + when(partitionDAO.getEntityStats(jobId.toString())) + .thenReturn( + List.of( + new CollectionDAO.RdfIndexPartitionDAO.RdfEntityStatsRecord( + "table", 25L, 25L, 25L, 0L, 1, 1, 0))); + when(partitionDAO.getServerStats(jobId.toString())).thenReturn(List.of()); + + RdfIndexJob refreshed = coordinator.getJobWithAggregatedStats(jobId); + + assertEquals(IndexJobStatus.COMPLETED, refreshed.getStatus()); + assertNotNull(refreshed.getCompletedAt()); + + ArgumentCaptor completedAtCaptor = ArgumentCaptor.forClass(Long.class); + verify(jobDAO) + .update( + eq(jobId.toString()), + eq(IndexJobStatus.COMPLETED.name()), + eq(25L), + eq(25L), + eq(0L), + anyString(), + eq(jobRecord.startedAt()), + completedAtCaptor.capture(), + anyLong(), + isNull()); + + assertNotNull(completedAtCaptor.getValue()); + assertEquals(completedAtCaptor.getValue(), refreshed.getCompletedAt()); + } + + @Test + void claimNextPartitionUsesUniqueMillisecondTimestamps() { + UUID jobId = UUID.randomUUID(); + when(partitionDAO.claimNextPartitionAtomic(eq(jobId.toString()), eq(TEST_SERVER_ID), anyLong())) + .thenReturn(1); + when(partitionDAO.findLatestClaimedPartition( + eq(jobId.toString()), eq(TEST_SERVER_ID), anyLong())) + .thenAnswer( + invocation -> + new CollectionDAO.RdfIndexPartitionDAO.RdfIndexPartitionRecord( + UUID.randomUUID().toString(), + jobId.toString(), + "table", + 0, + 0L, + 100L, + 100L, + 100L, + 1, + "PROCESSING", + 0L, + 0L, + 0L, + 0L, + TEST_SERVER_ID, + invocation.getArgument(2, Long.class), + invocation.getArgument(2, Long.class), + null, + invocation.getArgument(2, Long.class), + null, + 0, + 0L)); + + coordinator.claimNextPartition(jobId); + coordinator.claimNextPartition(jobId); + + ArgumentCaptor claimTimes = ArgumentCaptor.forClass(Long.class); + verify(partitionDAO, times(2)) + .claimNextPartitionAtomic(eq(jobId.toString()), eq(TEST_SERVER_ID), claimTimes.capture()); + + List capturedTimes = claimTimes.getAllValues(); + assertEquals(2, capturedTimes.size()); + assertNotEquals(capturedTimes.get(0), capturedTimes.get(1)); + assertTrue(capturedTimes.get(1) > capturedTimes.get(0)); + } + + @Test + void updateJobStatusPreservesExistingCompletedAt() { + UUID jobId = UUID.randomUUID(); + long completedAt = System.currentTimeMillis() - 5000; + EventPublisherJob jobConfiguration = new EventPublisherJob().withEntities(Set.of("table")); + RdfIndexJobRecord jobRecord = + new RdfIndexJobRecord( + jobId.toString(), + IndexJobStatus.RUNNING.name(), + JsonUtils.pojoToJson(jobConfiguration), + 25L, + 25L, + 25L, + 0L, + JsonUtils.pojoToJson(Map.of()), + "admin", + System.currentTimeMillis(), + System.currentTimeMillis() - 10000, + completedAt, + System.currentTimeMillis(), + null); + + when(jobDAO.findById(jobId.toString())).thenReturn(jobRecord); + + coordinator.updateJobStatus(jobId, IndexJobStatus.COMPLETED, null); + + verify(jobDAO) + .update( + eq(jobId.toString()), + eq(IndexJobStatus.COMPLETED.name()), + eq(25L), + eq(25L), + eq(0L), + anyString(), + eq(jobRecord.startedAt()), + eq(completedAt), + anyLong(), + isNull()); + } + + @Test + void hasClaimableWorkUsesCountQueries() { + UUID jobId = UUID.randomUUID(); + EventPublisherJob jobConfiguration = new EventPublisherJob().withEntities(Set.of("table")); + RdfIndexJobRecord jobRecord = + new RdfIndexJobRecord( + jobId.toString(), + IndexJobStatus.READY.name(), + JsonUtils.pojoToJson(jobConfiguration), + 25L, + 0L, + 0L, + 0L, + JsonUtils.pojoToJson(Map.of()), + "admin", + System.currentTimeMillis(), + null, + null, + System.currentTimeMillis(), + null); + + when(jobDAO.findById(jobId.toString())).thenReturn(jobRecord); + when(partitionDAO.getAggregatedStats(jobId.toString())) + .thenReturn(new RdfAggregatedStatsRecord(25L, 0L, 0L, 0L, 1, 0, 0, 1, 0)); + when(partitionDAO.getEntityStats(jobId.toString())).thenReturn(List.of()); + when(partitionDAO.getServerStats(jobId.toString())).thenReturn(List.of()); + when(partitionDAO.countPendingPartitions(jobId.toString())).thenReturn(0); + when(partitionDAO.countInFlightPartitions(jobId.toString())).thenReturn(1); + + assertTrue(coordinator.hasClaimableWork(jobId)); + + verify(partitionDAO).countPendingPartitions(jobId.toString()); + verify(partitionDAO).countInFlightPartitions(jobId.toString()); + verify(partitionDAO, never()).findByJobId(jobId.toString()); + } + + @Test + void hasClaimableWorkReturnsFalseWhenNoClaimableOrInflightPartitionsExist() { + UUID jobId = UUID.randomUUID(); + EventPublisherJob jobConfiguration = new EventPublisherJob().withEntities(Set.of("table")); + RdfIndexJobRecord jobRecord = + new RdfIndexJobRecord( + jobId.toString(), + IndexJobStatus.READY.name(), + JsonUtils.pojoToJson(jobConfiguration), + 25L, + 0L, + 0L, + 0L, + JsonUtils.pojoToJson(Map.of()), + "admin", + System.currentTimeMillis(), + null, + null, + System.currentTimeMillis(), + null); + + when(jobDAO.findById(jobId.toString())).thenReturn(jobRecord); + when(partitionDAO.getAggregatedStats(jobId.toString())) + .thenReturn(new RdfAggregatedStatsRecord(25L, 0L, 0L, 0L, 1, 0, 0, 1, 0)); + when(partitionDAO.getEntityStats(jobId.toString())).thenReturn(List.of()); + when(partitionDAO.getServerStats(jobId.toString())).thenReturn(List.of()); + when(partitionDAO.countPendingPartitions(jobId.toString())).thenReturn(0); + when(partitionDAO.countInFlightPartitions(jobId.toString())).thenReturn(0); + + assertFalse(coordinator.hasClaimableWork(jobId)); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexExecutorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexExecutorTest.java new file mode 100644 index 000000000000..12effe2ba9b4 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/DistributedRdfIndexExecutorTest.java @@ -0,0 +1,147 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.rdf.RdfRepository; + +class DistributedRdfIndexExecutorTest { + + @Test + void executeReleasesCoordinatorStateWhenWorkerStartupFails() { + CollectionDAO collectionDAO = org.mockito.Mockito.mock(CollectionDAO.class); + DistributedRdfIndexCoordinator coordinator = + org.mockito.Mockito.mock(DistributedRdfIndexCoordinator.class); + + UUID jobId = UUID.randomUUID(); + EventPublisherJob jobConfiguration = + new EventPublisherJob().withEntities(Set.of("table")).withConsumerThreads(1); + RdfIndexJob job = + RdfIndexJob.builder() + .id(jobId) + .status( + org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus.READY) + .jobConfiguration(jobConfiguration) + .build(); + + when(coordinator.getBlockingJob()).thenReturn(Optional.empty()); + when(coordinator.tryAcquireReindexLock(any(UUID.class))).thenReturn(true); + when(coordinator.createJob(Set.of("table"), jobConfiguration, "admin")).thenReturn(job); + when(coordinator.initializePartitions(jobId)).thenReturn(job); + when(coordinator.transferReindexLock(any(UUID.class), any(UUID.class))).thenReturn(true); + when(coordinator.getJobWithAggregatedStats(jobId)).thenReturn(job); + + DistributedRdfIndexExecutor executor = + new DistributedRdfIndexExecutor(collectionDAO, coordinator, "rdf-test-server"); + executor.createJob(Set.of("table"), jobConfiguration, "admin"); + + try (MockedStatic rdfRepository = mockStatic(RdfRepository.class)) { + rdfRepository.when(RdfRepository::getInstance).thenThrow(new IllegalStateException("boom")); + + assertThrows(IllegalStateException.class, () -> executor.execute(jobConfiguration)); + } + + verify(coordinator).releaseReindexLock(jobId); + assertFalse(DistributedRdfIndexExecutor.isCoordinatingJob(jobId)); + verify(coordinator) + .updateJobStatus( + jobId, + org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus.RUNNING, + null); + } + + @Test + void stopAndCoordinatorCleanupOnlyTearDownLocalExecutionOnce() throws Exception { + CollectionDAO collectionDAO = org.mockito.Mockito.mock(CollectionDAO.class); + DistributedRdfIndexCoordinator coordinator = + org.mockito.Mockito.mock(DistributedRdfIndexCoordinator.class); + ExecutorService workerExecutor = org.mockito.Mockito.mock(ExecutorService.class); + + when(workerExecutor.awaitTermination(anyLong(), any(TimeUnit.class))).thenReturn(true); + + DistributedRdfIndexExecutor executor = + new DistributedRdfIndexExecutor(collectionDAO, coordinator, "rdf-test-server"); + + Thread lockRefreshThread = startSleepingThread("rdf-lock-refresh-test"); + Thread staleReclaimerThread = startSleepingThread("rdf-stale-reclaimer-test"); + + setField(executor, "workerExecutor", workerExecutor); + setField(executor, "lockRefreshThread", lockRefreshThread); + setField(executor, "staleReclaimerThread", staleReclaimerThread); + getField(executor, "localExecutionCleaned", AtomicBoolean.class).set(false); + + executor.stop(); + invokeCleanupCoordinatorExecution(executor); + + verify(workerExecutor).shutdownNow(); + verify(workerExecutor).awaitTermination(30, TimeUnit.SECONDS); + assertFalse(lockRefreshThread.isAlive()); + assertFalse(staleReclaimerThread.isAlive()); + assertTrue(getField(executor, "localExecutionCleaned", AtomicBoolean.class).get()); + } + + private static Thread startSleepingThread(String name) { + return Thread.ofPlatform() + .name(name) + .start( + () -> { + try { + Thread.sleep(Long.MAX_VALUE); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + } + + private static void invokeCleanupCoordinatorExecution(DistributedRdfIndexExecutor executor) + throws Exception { + Method method = + DistributedRdfIndexExecutor.class.getDeclaredMethod("cleanupCoordinatorExecution"); + method.setAccessible(true); + method.invoke(executor); + } + + private static void setField(Object target, String fieldName, Object value) throws Exception { + Field field = target.getClass().getDeclaredField(fieldName); + field.setAccessible(true); + field.set(target, value); + } + + @SuppressWarnings("unchecked") + private static T getField(Object target, String fieldName, Class type) throws Exception { + Field field = target.getClass().getDeclaredField(fieldName); + field.setAccessible(true); + return (T) field.get(target); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobParticipantTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobParticipantTest.java new file mode 100644 index 000000000000..9e0a8bbf8b35 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfDistributedJobParticipantTest.java @@ -0,0 +1,59 @@ +/* + * Copyright 2026 Collate + * 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. + */ + +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.ServerIdentityResolver; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.rdf.RdfRepository; + +class RdfDistributedJobParticipantTest { + + private MockedStatic serverIdentityMock; + + @BeforeEach + void setUp() { + ServerIdentityResolver resolver = mock(ServerIdentityResolver.class); + when(resolver.getServerId()).thenReturn("rdf-test-server"); + + serverIdentityMock = mockStatic(ServerIdentityResolver.class); + serverIdentityMock.when(ServerIdentityResolver::getInstance).thenReturn(resolver); + RdfRepository.reset(); + } + + @AfterEach + void tearDown() { + RdfRepository.reset(); + if (serverIdentityMock != null) { + serverIdentityMock.close(); + } + } + + @Test + void startDoesNotFailWhenRdfRepositoryIsNotInitialized() { + RdfDistributedJobParticipant participant = + new RdfDistributedJobParticipant(mock(CollectionDAO.class)); + + assertDoesNotThrow(participant::start); + assertDoesNotThrow(participant::stop); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionWorkerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionWorkerTest.java new file mode 100644 index 000000000000..aa8a3d2904f7 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/rdf/distributed/RdfPartitionWorkerTest.java @@ -0,0 +1,100 @@ +package org.openmetadata.service.apps.bundles.rdf.distributed; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.junit.jupiter.MockitoExtension; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.rdf.RdfBatchProcessor; +import org.openmetadata.service.jdbi3.EntityRepository; +import org.openmetadata.service.jdbi3.ListFilter; + +@ExtendWith(MockitoExtension.class) +class RdfPartitionWorkerTest { + + @Mock private DistributedRdfIndexCoordinator coordinator; + @Mock private RdfBatchProcessor batchProcessor; + + private RdfPartitionWorker worker; + + @BeforeEach + void setUp() { + worker = new RdfPartitionWorker(coordinator, batchProcessor, 100); + } + + @Test + void initializeKeysetCursorHandlesRepositoryBackedEntities() throws Exception { + @SuppressWarnings("unchecked") + EntityRepository repository = mock(EntityRepository.class); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock.when(() -> Entity.getEntityRepository("table")).thenReturn(repository); + when(repository.getCursorAtOffset(any(ListFilter.class), eq(4))).thenReturn("cursor-4"); + + assertNull( + invokePrivate( + worker, + "initializeKeysetCursor", + new Class[] {String.class, long.class}, + "table", + 0L)); + assertEquals( + "cursor-4", + invokePrivate( + worker, + "initializeKeysetCursor", + new Class[] {String.class, long.class}, + "table", + 5L)); + } + } + + @Test + void initializeKeysetCursorRejectsOffsetsBeyondSupportedRange() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + invokePrivate( + worker, + "initializeKeysetCursor", + new Class[] {String.class, long.class}, + "table", + (long) Integer.MAX_VALUE + 2L)); + + assertTrue(exception.getMessage().contains("does not support offsets above")); + } + + private Object invokePrivate( + RdfPartitionWorker target, String methodName, Class[] parameterTypes, Object... args) + throws Exception { + Method method = RdfPartitionWorker.class.getDeclaredMethod(methodName, parameterTypes); + method.setAccessible(true); + try { + return method.invoke(target, args); + } catch (InvocationTargetException e) { + if (e.getCause() instanceof Exception exception) { + throw exception; + } + if (e.getCause() instanceof Error error) { + throw error; + } + throw e; + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorkerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorkerTest.java index e992fa0b1785..32d5563b3337 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorkerTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorkerTest.java @@ -783,6 +783,22 @@ void initializeKeysetCursorReturnsNullWhenRepositoryCursorMissing() throws Excep } } + @Test + void initializeKeysetCursorRejectsOffsetsBeyondSupportedRange() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + invokePrivate( + worker, + "initializeKeysetCursor", + new Class[] {String.class, long.class}, + "table", + (long) Integer.MAX_VALUE + 2L)); + + assertTrue(exception.getMessage().contains("does not support offsets above")); + } + @Test void testPartitionResult_RecordWithReaderFailuresDefaultsWarningsToZero() { PartitionWorker.PartitionResult result = new PartitionWorker.PartitionResult(10, 2, false, 3); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java index 5243f8e6b401..07a98a6d86d0 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java @@ -119,4 +119,15 @@ void test_getAgentTypeCondition_singleAgentTypeWithComma() { assertEquals("CollateAI", filter.getQueryParams().get("agentType_0")); assertNull(filter.getQueryParams().get("agentType_1")); } + + @Test + void test_serverIdConditionOnlyAppliesToMcpExecutionTable() { + ListFilter filter = new ListFilter(); + filter.addQueryParam("serverId", "mcp-server-1"); + + assertFalse(filter.getCondition("table_entity").contains("serverId = :serverId")); + assertEquals( + "WHERE mcp_execution_entity.deleted = FALSE AND serverId = :serverId", + filter.getCondition("mcp_execution_entity")); + } } diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfPropertyMapperTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfPropertyMapperTest.java index bae9ceebe335..b59adc5e5ca0 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfPropertyMapperTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfPropertyMapperTest.java @@ -61,85 +61,42 @@ void setUp() { class ChangeDescriptionTests { @Test - @DisplayName("ChangeDescription should be stored as structured RDF, not JSON literal") - void testChangeDescriptionStructured() throws Exception { + @DisplayName("ChangeDescription should be ignored during RDF field processing") + void testChangeDescriptionIsIgnored() throws Exception { ObjectNode changeDesc = objectMapper.createObjectNode(); changeDesc.put("previousVersion", 1.0); - ArrayNode fieldsAdded = objectMapper.createArrayNode(); - ObjectNode addedField = objectMapper.createObjectNode(); - addedField.put("name", "description"); - addedField.put("newValue", "New description value"); - fieldsAdded.add(addedField); - changeDesc.set("fieldsAdded", fieldsAdded); - - ArrayNode fieldsUpdated = objectMapper.createArrayNode(); - ObjectNode updatedField = objectMapper.createObjectNode(); - updatedField.put("name", "tags"); - updatedField.put("oldValue", "[]"); - updatedField.put("newValue", "[\"PII\"]"); - fieldsUpdated.add(updatedField); - changeDesc.set("fieldsUpdated", fieldsUpdated); - - // Use reflection to call the private method - java.lang.reflect.Method method = - RdfPropertyMapper.class.getDeclaredMethod( - "addChangeDescription", JsonNode.class, Resource.class, Model.class); - method.setAccessible(true); - method.invoke(propertyMapper, changeDesc, entityResource, model); - - // Verify structured RDF was created - Property hasChangeDesc = model.createProperty(OM_NS, "hasChangeDescription"); - assertTrue( - model.contains(entityResource, hasChangeDesc), - "Entity should have hasChangeDescription property"); - - // Find the change description resource - Resource changeDescResource = - model.listObjectsOfProperty(entityResource, hasChangeDesc).next().asResource(); - - // Verify type - assertTrue( - model.contains( - changeDescResource, RDF.type, model.createResource(OM_NS + "ChangeDescription")), - "ChangeDescription should have correct type"); + ObjectNode entityJson = objectMapper.createObjectNode(); + entityJson.set("changeDescription", changeDesc); - // Verify previousVersion is stored as a typed literal, not JSON - Property prevVersion = model.createProperty(OM_NS, "previousVersion"); - assertTrue( - model.contains(changeDescResource, prevVersion), - "ChangeDescription should have previousVersion"); - - // Verify fieldsAdded are stored as structured nodes - Property fieldsAddedProp = model.createProperty(OM_NS, "fieldsAdded"); - assertTrue( - model.contains(changeDescResource, fieldsAddedProp), - "ChangeDescription should have fieldsAdded"); + invokePrivate( + "processContextMappings", + new Class[] {Map.class, JsonNode.class, Resource.class, Model.class}, + Map.of("changeDescription", Map.of("@id", "om:hasChangeDescription", "@type", "@json")), + entityJson, + entityResource, + model); - // Verify the field change has a name property (not stored as JSON blob) - Resource fieldChangeResource = - model.listObjectsOfProperty(changeDescResource, fieldsAddedProp).next().asResource(); - Property fieldNameProp = model.createProperty(OM_NS, "fieldName"); - assertTrue( - model.contains(fieldChangeResource, fieldNameProp), - "FieldChange should have fieldName property"); + assertFalse( + model.contains(entityResource, model.createProperty(OM_NS, "hasChangeDescription")), + "ChangeDescription helper nodes should not be emitted into RDF"); } @Test - @DisplayName("Empty ChangeDescription should not create any triples") - void testEmptyChangeDescription() throws Exception { + @DisplayName("Structured property dispatch should ignore changeDescription") + void testStructuredDispatchIgnoresChangeDescription() throws Exception { ObjectNode changeDesc = objectMapper.createObjectNode(); - java.lang.reflect.Method method = - RdfPropertyMapper.class.getDeclaredMethod( - "addChangeDescription", JsonNode.class, Resource.class, Model.class); - method.setAccessible(true); - method.invoke(propertyMapper, changeDesc, entityResource, model); + invokePrivate( + "addStructuredProperty", + new Class[] {String.class, JsonNode.class, Resource.class, Model.class}, + "changeDescription", + changeDesc, + entityResource, + model); - Property hasChangeDesc = model.createProperty(OM_NS, "hasChangeDescription"); - assertTrue( - model.contains(entityResource, hasChangeDesc), - "Entity should still have hasChangeDescription for empty change"); + assertFalse( + model.contains(entityResource, model.createProperty(OM_NS, "hasChangeDescription"))); } } @@ -148,7 +105,7 @@ void testEmptyChangeDescription() throws Exception { class VotesTests { @Test - @DisplayName("Votes should be stored as structured RDF with upVotes/downVotes as integers") + @DisplayName("Votes should keep counts but omit voter relationship edges") void testVotesStructured() throws Exception { ObjectNode votes = objectMapper.createObjectNode(); votes.put("upVotes", 10); @@ -192,9 +149,12 @@ void testVotesStructured() throws Exception { stmt = model.getProperty(votesResource, downVotesProp); assertEquals(2, stmt.getInt(), "downVotes should be 2"); - // Verify upVoters are stored as entity references + // Verify individual voter references are not stored as graph edges Property upVotersProp = model.createProperty(OM_NS, "upVoters"); - assertTrue(model.contains(votesResource, upVotersProp), "Votes should have upVoters"); + assertFalse(model.contains(votesResource, upVotersProp), "Votes should not expose upVoters"); + assertFalse( + model.contains(votesResource, model.createProperty(OM_NS, "downVoters")), + "Votes should not expose downVoters"); } } @@ -823,11 +783,12 @@ void testContainerVotesAndExtensionHelpersCoverRemainingBranches() throws Except .listObjectsOfProperty(entityResource, model.createProperty(OM_NS, "hasVotes")) .next() .asResource(); - assertTrue( + assertFalse( model.contains( votesResource, model.createProperty(OM_NS, "downVoters"), - model.createResource(BASE_URI + "entity/user/" + reviewerId))); + model.createResource(BASE_URI + "entity/user/" + reviewerId)), + "Vote helpers should not emit voter references"); ObjectNode extension = objectMapper.createObjectNode(); extension.put("threshold", 2.5); @@ -868,7 +829,7 @@ void testStructuredPropertyDispatchAndCustomProperties() throws Exception { changeDescription, entityResource, model); - assertTrue( + assertFalse( model.contains(entityResource, model.createProperty(OM_NS, "hasChangeDescription"))); ObjectNode votes = objectMapper.createObjectNode(); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/ai/McpExecutionResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/ai/McpExecutionResourceTest.java index 152adedb349a..971e60e80f13 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/ai/McpExecutionResourceTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/ai/McpExecutionResourceTest.java @@ -2,6 +2,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyInt; @@ -11,6 +12,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockStatic; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; import jakarta.ws.rs.core.Response; @@ -128,6 +130,27 @@ void testListWithServerId() { } } + @Test + void testListRejectsPartialTimeRange() { + try (MockedStatic entityMock = mockStatic(Entity.class)) { + McpExecutionRepository mockRepo = mock(McpExecutionRepository.class); + Authorizer mockAuth = mock(Authorizer.class); + McpExecutionResource resource = createResource(entityMock, mockRepo, mockAuth); + + SecurityContext securityContext = mock(SecurityContext.class); + doNothing().when(mockAuth).authorize(any(), any(), any()); + + assertThrows( + IllegalArgumentException.class, + () -> resource.list(securityContext, null, 1000L, null, 10)); + assertThrows( + IllegalArgumentException.class, + () -> resource.list(securityContext, null, null, 2000L, 10)); + + verifyNoInteractions(mockRepo); + } + } + @Test void testGet() { try (MockedStatic entityMock = mockStatic(Entity.class)) { diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/RdfResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/RdfResourceTest.java new file mode 100644 index 000000000000..0e398c2332be --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/RdfResourceTest.java @@ -0,0 +1,133 @@ +/* + * Copyright 2024 Collate + * 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. + */ + +package org.openmetadata.service.resources.rdf; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import jakarta.ws.rs.core.Response; +import jakarta.ws.rs.core.SecurityContext; +import java.lang.reflect.Field; +import java.util.Set; +import java.util.UUID; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; +import org.openmetadata.service.Entity; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.security.Authorizer; + +class RdfResourceTest { + + private Authorizer authorizer; + private SecurityContext securityContext; + private RdfResource rdfResource; + + @BeforeEach + void setUp() { + authorizer = Mockito.mock(Authorizer.class); + securityContext = Mockito.mock(SecurityContext.class); + doNothing().when(authorizer).authorizeAdmin(securityContext); + rdfResource = new RdfResource(authorizer); + } + + private void setRdfRepository(RdfRepository repository) throws Exception { + Field field = RdfResource.class.getDeclaredField("rdfRepository"); + field.setAccessible(true); + field.set(rdfResource, repository); + } + + @Test + void exploreEntityGraphRejectsInvalidEntityType() { + Response response = + rdfResource.exploreEntityGraph( + securityContext, UUID.randomUUID(), "table } UNION { ?s ?p ?o", 2, null, null); + + assertEquals(Response.Status.BAD_REQUEST.getStatusCode(), response.getStatus()); + assertTrue(String.valueOf(response.getEntity()).contains("Invalid entity type")); + } + + @Test + void getFullLineageRejectsInvalidEntityType() { + Response response = + rdfResource.getFullLineage( + securityContext, UUID.randomUUID(), "table } UNION { ?s ?p ?o", "both"); + + assertEquals(Response.Status.BAD_REQUEST.getStatusCode(), response.getStatus()); + assertTrue(String.valueOf(response.getEntity()).contains("Invalid entity type")); + } + + @Test + void exportEntityGraphRejectsUnsupportedFormat() { + Response response = + rdfResource.exportEntityGraph( + securityContext, UUID.randomUUID(), "table", 2, null, null, "rdfxml"); + + assertEquals(Response.Status.BAD_REQUEST.getStatusCode(), response.getStatus()); + } + + @Test + void exploreEntityGraphClampsDepthToMaximum() throws Exception { + RdfRepository repository = Mockito.mock(RdfRepository.class); + UUID entityId = UUID.randomUUID(); + when(repository.isEnabled()).thenReturn(true); + when(repository.getEntityGraph(entityId, "table", 5, Set.of(), Set.of())).thenReturn("{}"); + setRdfRepository(repository); + + Response response; + try (MockedStatic entityMock = Mockito.mockStatic(Entity.class)) { + entityMock.when(() -> Entity.hasEntityRepository("table")).thenReturn(true); + response = rdfResource.exploreEntityGraph(securityContext, entityId, "table", 99, null, null); + } + + assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(repository).getEntityGraph(eq(entityId), eq("table"), eq(5), eq(Set.of()), eq(Set.of())); + } + + @Test + void exportEntityGraphClampsDepthToMaximum() throws Exception { + RdfRepository repository = Mockito.mock(RdfRepository.class); + UUID entityId = UUID.randomUUID(); + when(repository.isEnabled()).thenReturn(true); + when(repository.exportEntityGraph(entityId, "table", 5, Set.of(), Set.of(), "TURTLE")) + .thenReturn("@prefix ex: ."); + setRdfRepository(repository); + + Response response; + try (MockedStatic entityMock = Mockito.mockStatic(Entity.class)) { + entityMock.when(() -> Entity.hasEntityRepository("table")).thenReturn(true); + response = + rdfResource.exportEntityGraph( + securityContext, entityId, "table", 99, null, null, "turtle"); + } + + assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(repository) + .exportEntityGraph( + eq(entityId), eq("table"), eq(5), eq(Set.of()), eq(Set.of()), eq("TURTLE")); + } + + @Test + void normalizeEntityGraphExportFormatAcceptsAliases() { + assertEquals("TURTLE", RdfRepository.normalizeEntityGraphExportFormat("ttl")); + assertEquals("TURTLE", RdfRepository.normalizeEntityGraphExportFormat("turtle")); + assertEquals("JSON-LD", RdfRepository.normalizeEntityGraphExportFormat("jsonld")); + assertEquals("JSON-LD", RdfRepository.normalizeEntityGraphExportFormat("json-ld")); + } +} diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/rdfIndexingAppConfig.json b/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/rdfIndexingAppConfig.json new file mode 100644 index 000000000000..e5f2f4d28015 --- /dev/null +++ b/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/rdfIndexingAppConfig.json @@ -0,0 +1,153 @@ +{ + "$id": "https://open-metadata.org/schema/entity/applications/configuration/rdfIndexingApp.json", + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "RdfIndexingApp", + "type": "object", + "description": "RDF indexing application configuration.", + "definitions": { + "rdfIndexingType": { + "description": "Application type.", + "type": "string", + "enum": ["RdfIndexing"], + "default": "RdfIndexing" + } + }, + "properties": { + "type": { + "title": "Application Type", + "description": "Application Type", + "$ref": "#/definitions/rdfIndexingType", + "default": "RdfIndexing" + }, + "entities": { + "title": "Entities", + "description": "List of entities that you need to reindex. Leave empty to index all supported entities.", + "type": "array", + "items": { + "type": "string", + "enum": [ + "aiApplication", + "aiGovernancePolicy", + "apiCollection", + "apiEndpoint", + "apiService", + "app", + "appMarketPlaceDefinition", + "bot", + "chart", + "classification", + "container", + "dashboard", + "dashboardDataModel", + "dashboardService", + "dataContract", + "dataInsightChart", + "dataInsightCustomChart", + "dataProduct", + "database", + "databaseSchema", + "databaseService", + "directory", + "document", + "domain", + "driveService", + "eventsubscription", + "file", + "glossary", + "glossaryTerm", + "ingestionPipeline", + "kpi", + "learningResource", + "llmModel", + "llmService", + "messagingService", + "metadataService", + "metric", + "mlmodel", + "mlmodelService", + "notificationTemplate", + "persona", + "pipeline", + "pipelineService", + "policy", + "promptTemplate", + "query", + "report", + "role", + "searchIndex", + "searchService", + "securityService", + "spreadsheet", + "storageService", + "storedProcedure", + "table", + "tag", + "team", + "testCase", + "testConnectionDefinition", + "testDefinition", + "testSuite", + "topic", + "type", + "user", + "webAnalyticEvent", + "workflow", + "workflowDefinition", + "worksheet" + ] + }, + "default": [], + "uiFieldType": "treeSelect", + "uniqueItems": true + }, + "recreateIndex": { + "title": "Recreate RDF Store", + "description": "Recreate the RDF store before indexing.", + "type": "boolean", + "default": false + }, + "batchSize": { + "title": "Batch Size", + "description": "Maximum number of entities processed in a batch.", + "type": "integer", + "default": 100, + "minimum": 1 + }, + "producerThreads": { + "title": "Number of Producer Threads", + "description": "Number of producer threads to use for non-distributed RDF reindexing", + "type": "integer", + "default": 2, + "minimum": 1 + }, + "consumerThreads": { + "title": "Number of Consumer Threads", + "description": "Number of consumer threads to use for non-distributed RDF reindexing", + "type": "integer", + "default": 3, + "minimum": 1 + }, + "queueSize": { + "title": "Queue Size", + "description": "Queue size to use internally for non-distributed RDF reindexing.", + "type": "integer", + "default": 5000, + "minimum": 1 + }, + "useDistributedIndexing": { + "title": "Use Distributed Indexing", + "description": "Enable distributed RDF indexing across multiple servers with partition coordination and recovery.", + "type": "boolean", + "default": true + }, + "partitionSize": { + "title": "Partition Size", + "description": "Number of entities per partition for distributed RDF indexing. Smaller values create more partitions for better distribution across servers.", + "type": "integer", + "default": 10000, + "minimum": 1000, + "maximum": 50000 + } + }, + "additionalProperties": false +} diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/applications/jobStatus.json b/openmetadata-spec/src/main/resources/json/schema/entity/applications/jobStatus.json index f8ba929466fd..8a0f40b96ba0 100644 --- a/openmetadata-spec/src/main/resources/json/schema/entity/applications/jobStatus.json +++ b/openmetadata-spec/src/main/resources/json/schema/entity/applications/jobStatus.json @@ -11,8 +11,11 @@ "oneOf": [ { "$ref": "configuration/internal/searchIndexingAppConfig.json" + }, + { + "$ref": "configuration/internal/rdfIndexingAppConfig.json" } ] } } -} \ No newline at end of file +} diff --git a/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/card/card.tsx b/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/card/card.tsx index 55997976708f..e12de5518566 100644 --- a/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/card/card.tsx +++ b/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/card/card.tsx @@ -32,43 +32,43 @@ const sizes = sortCx({ export const cardStyles = sortCx({ common: { - root: 'tw:outline-focus-ring tw:focus-visible:outline-2 tw:focus-visible:outline-offset-2 tw:relative tw:overflow-hidden tw:rounded-xl tw:transition tw:duration-100', + root: 'tw:focus-visible:outline-2 tw:focus-visible:outline-offset-2 tw:relative tw:overflow-hidden tw:rounded-xl tw:transition tw:duration-100', }, variants: { default: { - root: 'tw:ring-1 tw:ring-inset tw:ring-secondary tw:bg-primary', + root: 'tw:border-1 tw:border-secondary tw:bg-primary', }, elevated: { - root: 'tw:ring-1 tw:ring-inset tw:ring-secondary tw:bg-primary tw:shadow-md', + root: 'tw:border-1 tw:border-secondary tw:bg-primary tw:shadow-md', }, - outlined: { root: 'tw:ring-2 tw:ring-inset tw:ring-primary tw:bg-primary' }, + outlined: { root: 'tw:border-2 tw:border-primary tw:bg-primary' }, ghost: { root: 'tw:bg-transparent' }, }, colors: { default: { root: '' }, brand: { - root: 'tw:bg-utility-brand-50 tw:ring-1 tw:ring-inset tw:ring-utility-brand-200', + root: 'tw:bg-utility-brand-50 tw:border-1 tw:border-utility-brand-200', }, error: { - root: 'tw:bg-utility-error-50 tw:ring-1 tw:ring-inset tw:ring-utility-error-200', + root: 'tw:bg-utility-error-50 tw:border-1 tw:border-utility-error-200', }, warning: { - root: 'tw:bg-utility-warning-50 tw:ring-1 tw:ring-inset tw:ring-utility-warning-200', + root: 'tw:bg-utility-warning-50 tw:border-1 tw:border-utility-warning-200', }, success: { - root: 'tw:bg-utility-success-50 tw:ring-1 tw:ring-inset tw:ring-utility-success-200', + root: 'tw:bg-utility-success-50 tw:border-1 tw:border-utility-success-200', }, }, interactive: { root: 'tw:cursor-pointer' }, interactiveVariants: { - default: { root: 'tw:hover:bg-primary_hover tw:hover:ring-primary' }, + default: { root: 'tw:hover:bg-primary_hover tw:hover:border-primary' }, elevated: { root: 'tw:hover:bg-primary_hover tw:hover:shadow-lg' }, - outlined: { root: 'tw:hover:bg-secondary tw:hover:ring-primary' }, + outlined: { root: 'tw:hover:bg-secondary tw:hover:border-primary' }, ghost: { - root: 'tw:hover:bg-secondary tw:hover:ring-1 tw:hover:ring-inset tw:hover:ring-secondary', + root: 'tw:hover:bg-secondary tw:hover:border-1 tw:hover:border-inset tw:hover:border-secondary', }, }, - selected: { root: 'tw:ring-2 tw:ring-inset tw:ring-brand' }, + selected: { root: 'tw:border-2 tw:border-brand' }, }); // ─── Sub-component interfaces ────────────────────────────────────────────────── @@ -103,7 +103,7 @@ const CardHeader = ({

diff --git a/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/dropdown/dropdown.tsx b/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/dropdown/dropdown.tsx index 51144adfb8c5..829db4711c97 100644 --- a/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/dropdown/dropdown.tsx +++ b/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/dropdown/dropdown.tsx @@ -17,6 +17,7 @@ import { Popover as AriaPopover, Separator as AriaSeparator, } from 'react-aria-components'; +import { CheckboxBase } from '@/components/base/checkbox/checkbox'; import { cx } from '@/utils/cx'; interface DropdownItemProps extends AriaMenuItemProps { @@ -28,6 +29,8 @@ interface DropdownItemProps extends AriaMenuItemProps { unstyled?: boolean; /** An icon to be displayed on the left side of the item. */ icon?: FC<{ className?: string }>; + /** If true, shows a checkbox on the left to indicate selection state. */ + showCheckbox?: boolean; } const DropdownItem = ({ @@ -36,6 +39,7 @@ const DropdownItem = ({ addon, icon: Icon, unstyled, + showCheckbox, ...props }: DropdownItemProps) => { if (unstyled) { @@ -57,16 +61,26 @@ const DropdownItem = ({ {(state) => (
+ {showCheckbox && ( + + )} + {Icon && (