diff --git a/.github/workflows/sql-test-and-build-workflow.yml b/.github/workflows/sql-test-and-build-workflow.yml index 964cc63eb20..e9a280ee803 100644 --- a/.github/workflows/sql-test-and-build-workflow.yml +++ b/.github/workflows/sql-test-and-build-workflow.yml @@ -4,6 +4,7 @@ on: pull_request: push: branches-ignore: + - 'backport/**' - 'dependabot/**' paths: - '**/*.java' @@ -11,7 +12,6 @@ on: - '!sql-jdbc/**' - '**gradle*' - '**lombok*' - - '**spotless*' - 'integ-test/**' - '**/*.jar' - '**/*.pom' @@ -26,111 +26,204 @@ jobs: build-linux: needs: Get-CI-Image-Tag strategy: - # Run all jobs fail-fast: false matrix: - java: - - 11 - - 17 - - 21 + java: [11, 17, 21] + test-type: ['unit', 'integration', 'doc'] runs-on: ubuntu-latest container: - # using the same image which is used by opensearch-build team to build the OpenSearch Distribution - # this image tag is subject to change as more dependencies and updates will arrive over time image: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-version-linux }} options: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-options }} steps: - - name: Run start commands - run: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-command }} - - - uses: actions/checkout@v4 - - - name: Set up JDK ${{ matrix.java }} - uses: actions/setup-java@v4 - with: - distribution: 'temurin' - java-version: ${{ matrix.java }} - - - name: Build with Gradle - run: | - chown -R 1000:1000 `pwd` - su `id -un 1000` -c "./gradlew --continue build" - - - name: Run backward compatibility tests - run: | - chown -R 1000:1000 `pwd` - su `id -un 1000` -c "./scripts/bwctest.sh" - - - name: Create Artifact Path - run: | - mkdir -p opensearch-sql-builds - cp -r ./plugin/build/distributions/*.zip opensearch-sql-builds/ - - # This step uses the codecov-action Github action: https://github.com/codecov/codecov-action - - name: Upload SQL Coverage Report - if: always() - uses: codecov/codecov-action@v4 - with: - flags: sql-engine - token: ${{ secrets.CODECOV_TOKEN }} - - - name: Upload Artifacts - uses: actions/upload-artifact@v4 - with: - name: opensearch-sql-ubuntu-latest-${{ matrix.java }} - path: opensearch-sql-builds - - - name: Upload test reports - if: always() - uses: actions/upload-artifact@v4 - continue-on-error: true - with: - name: test-reports-ubuntu-latest-${{ matrix.java }} - path: | - sql/build/reports/** - ppl/build/reports/** - core/build/reports/** - common/build/reports/** - opensearch/build/reports/** - integ-test/build/reports/** - protocol/build/reports/** - legacy/build/reports/** - plugin/build/reports/** + - name: Run start commands + run: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-command }} + + - uses: actions/checkout@v4 + + - name: Set up JDK ${{ matrix.java }} + uses: actions/setup-java@v4 + with: + distribution: 'temurin' + java-version: ${{ matrix.java }} + + - name: Build and Test + run: | + chown -R 1000:1000 `pwd` + if [ "${{ matrix.test-type }}" = "unit" ]; then + su `id -un 1000` -c "./gradlew --continue build -x integTest -x doctest" + elif [ "${{ matrix.test-type }}" = "integration" ]; then + su `id -un 1000` -c "./gradlew --continue integTest" + else + su `id -un 1000` -c "./gradlew --continue doctest" + fi + + - name: Create Artifact Path + if: ${{ matrix.test-type == 'unit' }} + run: | + mkdir -p opensearch-sql-builds + cp -r ./plugin/build/distributions/*.zip opensearch-sql-builds/ + + - name: Upload SQL Coverage Report + if: ${{ always() }} + uses: codecov/codecov-action@v4 + continue-on-error: true + with: + flags: sql-engine + token: ${{ secrets.CODECOV_TOKEN }} + + - name: Upload Artifacts + if: ${{ matrix.test-type == 'unit' }} + uses: actions/upload-artifact@v4 + continue-on-error: true + with: + name: opensearch-sql-ubuntu-latest-${{ matrix.java }} + path: opensearch-sql-builds + + - name: Upload test reports + if: ${{ always() }} + uses: actions/upload-artifact@v4 + continue-on-error: true + with: + name: test-reports-ubuntu-latest-${{ matrix.java }}-${{ matrix.test-type }} + path: | + sql/build/reports/** + ppl/build/reports/** + core/build/reports/** + common/build/reports/** + opensearch/build/reports/** + integ-test/build/reports/** + protocol/build/reports/** + legacy/build/reports/** + plugin/build/reports/** + doctest/build/testclusters/docTestCluster-0/logs/* + integ-test/build/testclusters/*/logs/* build-windows-macos: strategy: - # Run all jobs fail-fast: false matrix: entry: - - { os: windows-latest, java: 11, os_build_args: -x doctest -PbuildPlatform=windows } - - { os: macos-13, java: 11} - - { os: windows-latest, java: 17, os_build_args: -x doctest -PbuildPlatform=windows } - - { os: macos-13, java: 17 } - - { os: windows-latest, java: 21, os_build_args: -x doctest -PbuildPlatform=windows } - - { os: macos-13, java: 21 } + - { os: windows-latest, java: 11, os_build_args: -PbuildPlatform=windows } + - { os: windows-latest, java: 17, os_build_args: -PbuildPlatform=windows } + - { os: windows-latest, java: 21, os_build_args: -PbuildPlatform=windows } + - { os: macos-13, java: 11, os_build_args: '' } + - { os: macos-13, java: 17, os_build_args: '' } + - { os: macos-13, java: 21, os_build_args: '' } + test-type: ['unit', 'integration', 'doc'] + exclude: + # Exclude doctest for Windows + - test-type: doc + entry: { os: windows-latest, java: 11, os_build_args: -PbuildPlatform=windows } + - test-type: doc + entry: { os: windows-latest, java: 17, os_build_args: -PbuildPlatform=windows } + - test-type: doc + entry: { os: windows-latest, java: 21, os_build_args: -PbuildPlatform=windows } + runs-on: ${{ matrix.entry.os }} + steps: + - uses: actions/checkout@v4 + + - name: Set up JDK ${{ matrix.entry.java }} + uses: actions/setup-java@v4 + with: + distribution: 'temurin' + java-version: ${{ matrix.entry.java }} + + - name: Build and Test + run: | + if [ "${{ matrix.test-type }}" = "unit" ]; then + ./gradlew --continue build -x integTest -x doctest ${{ matrix.entry.os_build_args }} + elif [ "${{ matrix.test-type }}" = "integration" ]; then + ./gradlew --continue integTest ${{ matrix.entry.os_build_args }} + else + ./gradlew --continue doctest ${{ matrix.entry.os_build_args }} + fi + shell: bash + + - name: Create Artifact Path + if: ${{ matrix.test-type == 'unit' }} + run: | + mkdir -p opensearch-sql-builds + cp -r ./plugin/build/distributions/*.zip opensearch-sql-builds/ + + - name: Upload SQL Coverage Report + if: ${{ always() && matrix.entry.os == 'ubuntu-latest' }} + uses: codecov/codecov-action@v4 + continue-on-error: true + with: + flags: sql-engine + token: ${{ secrets.CODECOV_TOKEN }} + + - name: Upload Artifacts + if: ${{ matrix.test-type == 'unit' }} + uses: actions/upload-artifact@v4 + continue-on-error: true + with: + name: opensearch-sql-${{ matrix.entry.os }}-${{ matrix.entry.java }} + path: opensearch-sql-builds + + - name: Upload test reports + if: ${{ always() }} + uses: actions/upload-artifact@v4 + continue-on-error: true + with: + name: test-reports-${{ matrix.entry.os }}-${{ matrix.entry.java }}-${{ matrix.test-type }} + path: | + sql/build/reports/** + ppl/build/reports/** + core/build/reports/** + common/build/reports/** + opensearch/build/reports/** + integ-test/build/reports/** + protocol/build/reports/** + legacy/build/reports/** + plugin/build/reports/** + doctest/build/testclusters/docTestCluster-0/logs/* + integ-test/build/testclusters/*/logs/* + + bwc-tests: + needs: Get-CI-Image-Tag + runs-on: ubuntu-latest + strategy: + matrix: + java: [11, 17, 21] + container: + image: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-version-linux }} + options: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-options }} steps: - - uses: actions/checkout@v4 - - - name: Set up JDK ${{ matrix.java }} - uses: actions/setup-java@v4 - with: - distribution: 'temurin' - java-version: ${{ matrix.entry.java }} - - - name: Build with Gradle - run: ./gradlew --continue build ${{ matrix.entry.os_build_args }} - - - name: Create Artifact Path - run: | - mkdir -p opensearch-sql-builds - cp -r ./plugin/build/distributions/*.zip opensearch-sql-builds/ - - - name: Upload Artifacts - uses: actions/upload-artifact@v4 - with: - name: opensearch-sql-${{ matrix.entry.os }}-${{ matrix.entry.java }} - path: opensearch-sql-builds + - name: Run start commands + run: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-command }} + + - uses: actions/checkout@v4 + + - name: Set up JDK ${{ matrix.java }} + uses: actions/setup-java@v4 + with: + distribution: 'temurin' + java-version: ${{ matrix.java }} + + - name: Run backward compatibility tests + run: | + chown -R 1000:1000 `pwd` + su `id -un 1000` -c "./scripts/bwctest.sh" + + - name: Upload test reports + if: ${{ always() }} + uses: actions/upload-artifact@v4 + continue-on-error: true + with: + name: test-reports-ubuntu-latest-${{ matrix.java }}-bwc + path: | + sql/build/reports/** + ppl/build/reports/** + core/build/reports/** + common/build/reports/** + opensearch/build/reports/** + integ-test/build/reports/** + protocol/build/reports/** + legacy/build/reports/** + plugin/build/reports/** + doctest/build/testclusters/docTestCluster-0/logs/* + integ-test/build/testclusters/*/logs/* diff --git a/DEVELOPER_GUIDE.rst b/DEVELOPER_GUIDE.rst index 353627996c6..8767eb7c0c1 100644 --- a/DEVELOPER_GUIDE.rst +++ b/DEVELOPER_GUIDE.rst @@ -219,7 +219,9 @@ Building and Running Tests Gradle Build ------------ -Most of the time you just need to run ./gradlew build which will make sure you pass all checks and testing. While you’re developing, you may want to run specific Gradle task only. In this case, you can run ./gradlew with task name which only triggers the task along with those it depends on. Here is a list for common tasks: +Most of the time you just need to run ``./gradlew build`` which will make sure you pass all checks and testing. While you're developing, you may want to run specific Gradle task only. In this case, you can run ./gradlew with task name which only triggers the task along with those it depends on. Here is a list for common tasks: + +For faster local iterations, skip integration tests. ``./gradlew build -x integTest``. .. list-table:: :widths: 30 50 diff --git a/async-query-core/build.gradle b/async-query-core/build.gradle index 37bf6748c9d..147e59bd6a4 100644 --- a/async-query-core/build.gradle +++ b/async-query-core/build.gradle @@ -86,6 +86,7 @@ spotless { } test { + maxParallelForks = Runtime.runtime.availableProcessors() useJUnitPlatform() testLogging { events "skipped", "failed" diff --git a/async-query/build.gradle b/async-query/build.gradle index fba74aa216e..c616b5f37d2 100644 --- a/async-query/build.gradle +++ b/async-query/build.gradle @@ -53,6 +53,7 @@ dependencies { } test { + maxParallelForks = Runtime.runtime.availableProcessors() useJUnitPlatform { includeEngines("junit-jupiter") } diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java index 53b465aa6de..0319883c9b3 100644 --- a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java @@ -513,20 +513,20 @@ public String loadResultIndexMappings() { @RequiredArgsConstructor public class FlintDatasetMock { - final String query; - final String refreshQuery; + public final String query; + public final String refreshQuery; final FlintIndexType indexType; - final String indexName; + public final String indexName; boolean isLegacy = false; boolean isSpecialCharacter = false; - String latestId; + public String latestId; public FlintDatasetMock isLegacy(boolean isLegacy) { this.isLegacy = isLegacy; return this; } - FlintDatasetMock isSpecialCharacter(boolean isSpecialCharacter) { + public FlintDatasetMock isSpecialCharacter(boolean isSpecialCharacter) { this.isSpecialCharacter = isSpecialCharacter; return this; } diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/IndexQuerySpecAlterTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/IndexQuerySpecAlterTest.java deleted file mode 100644 index d69c7d48645..00000000000 --- a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/IndexQuerySpecAlterTest.java +++ /dev/null @@ -1,1124 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.spark.asyncquery; - -import com.amazonaws.services.emrserverless.model.CancelJobRunResult; -import com.amazonaws.services.emrserverless.model.GetJobRunResult; -import com.amazonaws.services.emrserverless.model.JobRun; -import com.amazonaws.services.emrserverless.model.ValidationException; -import com.google.common.collect.ImmutableList; -import java.util.HashMap; -import java.util.Map; -import org.junit.Test; -import org.junit.jupiter.api.Assertions; -import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; -import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; -import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; -import org.opensearch.sql.spark.client.EMRServerlessClientFactory; -import org.opensearch.sql.spark.client.StartJobRequest; -import org.opensearch.sql.spark.flint.FlintIndexState; -import org.opensearch.sql.spark.flint.FlintIndexType; -import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; -import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; -import org.opensearch.sql.spark.rest.model.LangType; - -public class IndexQuerySpecAlterTest extends AsyncQueryExecutorServiceSpec { - - @Test - public void testAlterIndexQueryConvertingToManualRefresh() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false)"); - MockFlintIndex ALTER_COVERING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false)"); - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," - + " incremental_refresh=false) "); - ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(1); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryConvertingToManualRefreshWithNoIncrementalRefresh() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false)"); - MockFlintIndex ALTER_COVERING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false)"); - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false)"); - ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - existingOptions.put("checkpoint_location", "s3://checkpoint/location"); - mockDS.updateIndexOptions(existingOptions, true); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(1); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryWithRedundantOperation() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false)"); - MockFlintIndex ALTER_COVERING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false)"); - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," - + " incremental_refresh=false) "); - ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public String startJobRun(StartJobRequest startJobRequest) { - return "jobId"; - } - - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - super.cancelJobRun(applicationId, jobId, allowExceptionPropagation); - throw new ValidationException("Job run is not in a cancellable state"); - } - }; - EMRServerlessClientFactory emrServerlessCientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessCientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "false"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(1); - emrsClient.getJobRunResultCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryConvertingToAutoRefresh() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=true," - + " incremental_refresh=false)"); - MockFlintIndex ALTER_COVERING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=true," - + " incremental_refresh=false)"); - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=true," - + " incremental_refresh=false) "); - ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = new LocalEMRSClient(); - EMRServerlessClientFactory clientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(clientFactory); - - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "false"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - assertEquals( - "RUNNING", - asyncQueryExecutorService - .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) - .getStatus()); - - flintIndexJob.assertState(FlintIndexState.ACTIVE); - emrsClient.startJobRunCalled(1); - emrsClient.getJobRunResultCalled(1); - emrsClient.cancelJobRunCalled(0); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryWithOutAnyAutoRefresh() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (" - + " incremental_refresh=false)"); - MockFlintIndex ALTER_COVERING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (" - + " incremental_refresh=false)"); - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (" + " incremental_refresh=false) "); - ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = new LocalEMRSClient(); - EMRServerlessClientFactory clientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(clientFactory); - - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "false"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - assertEquals( - "RUNNING", - asyncQueryExecutorService - .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) - .getStatus()); - - flintIndexJob.assertState(FlintIndexState.ACTIVE); - emrsClient.startJobRunCalled(1); - emrsClient.getJobRunResultCalled(1); - emrsClient.cancelJobRunCalled(0); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryOfFullRefreshWithInvalidOptions() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false, checkpoint_location=\"s3://ckp/skp\")"); - MockFlintIndex ALTER_COVERING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false, checkpoint_location=\"s3://ckp/skp\")"); - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," - + " incremental_refresh=false, checkpoint_location=\"s3://ckp/skp\") "); - ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals( - "Altering to full refresh only allows: [auto_refresh, incremental_refresh]" - + " options", - asyncQueryExecutionResponse.getError()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("true", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryOfIncrementalRefreshWithInvalidOptions() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=true, output_mode=\"complete\")"); - MockFlintIndex ALTER_COVERING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=true, output_mode=\"complete\")"); - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," - + " incremental_refresh=true, output_mode=\"complete\") "); - ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals( - "Altering to incremental refresh only allows: [auto_refresh, incremental_refresh," - + " watermark_delay, checkpoint_location] options", - asyncQueryExecutionResponse.getError()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("true", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryOfIncrementalRefreshWithInsufficientOptions() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=true)"); - MockFlintIndex ALTER_COVERING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=true)"); - ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - existingOptions.put("incremental_refresh", "false"); - mockDS.updateIndexOptions(existingOptions, true); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals( - "Conversion to incremental refresh index cannot proceed due to missing" - + " attributes: checkpoint_location.", - asyncQueryExecutionResponse.getError()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("true", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryOfIncrementalRefreshWithInsufficientOptionsForMV() { - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," - + " incremental_refresh=true) "); - ImmutableList.of(ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - existingOptions.put("incremental_refresh", "false"); - mockDS.updateIndexOptions(existingOptions, true); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals( - "Conversion to incremental refresh index cannot proceed due to missing" - + " attributes: checkpoint_location, watermark_delay.", - asyncQueryExecutionResponse.getError()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("true", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryOfIncrementalRefreshWithEmptyExistingOptionsForMV() { - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," - + " incremental_refresh=true) "); - ImmutableList.of(ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - existingOptions.put("incremental_refresh", "false"); - existingOptions.put("watermark_delay", ""); - existingOptions.put("checkpoint_location", ""); - mockDS.updateIndexOptions(existingOptions, true); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals( - "Conversion to incremental refresh index cannot proceed due to missing" - + " attributes: checkpoint_location, watermark_delay.", - asyncQueryExecutionResponse.getError()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("true", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryOfIncrementalRefresh() { - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," - + " incremental_refresh=true) "); - ImmutableList.of(ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - existingOptions.put("incremental_refresh", "false"); - existingOptions.put("watermark_delay", "watermark_delay"); - existingOptions.put("checkpoint_location", "s3://checkpoint/location"); - mockDS.updateIndexOptions(existingOptions, true); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); - emrsClient.startJobRunCalled(0); - emrsClient.getJobRunResultCalled(1); - emrsClient.cancelJobRunCalled(1); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - Assertions.assertEquals("true", options.get("incremental_refresh")); - }); - } - - @Test - public void testAlterIndexQueryWithIncrementalRefreshAlreadyExisting() { - MockFlintIndex ALTER_MV = - new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false) "); - ImmutableList.of(ALTER_MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - existingOptions.put("incremental_refresh", "true"); - existingOptions.put("watermark_delay", "watermark_delay"); - existingOptions.put("checkpoint_location", "s3://checkpoint/location"); - mockDS.updateIndexOptions(existingOptions, true); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); - emrsClient.startJobRunCalled(0); - emrsClient.getJobRunResultCalled(1); - emrsClient.cancelJobRunCalled(1); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - Assertions.assertEquals("true", options.get("incremental_refresh")); - }); - } - - @Test - public void testAlterIndexQueryWithInvalidInitialState() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false)"); - ImmutableList.of(ALTER_SKIPPING) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.updating(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals( - "Transaction failed as flint index is not in a valid state.", - asyncQueryExecutionResponse.getError()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - flintIndexJob.assertState(FlintIndexState.UPDATING); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("true", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryWithValidationExceptionWithSuccess() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false)"); - ImmutableList.of(ALTER_SKIPPING) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - super.cancelJobRun(applicationId, jobId, allowExceptionPropagation); - throw new ValidationException("Job run is not in a cancellable state"); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(1); - emrsClient.getJobRunResultCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryWithResourceNotFoundExceptionWithSuccess() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false)"); - ImmutableList.of(ALTER_SKIPPING) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - super.cancelJobRun(applicationId, jobId, allowExceptionPropagation); - throw new ValidationException("Random validation exception"); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals("Internal Server Error.", asyncQueryExecutionResponse.getError()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(1); - emrsClient.getJobRunResultCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - } - - @Test - public void testAlterIndexQueryWithUnknownError() { - MockFlintIndex ALTER_SKIPPING = - new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=false)"); - ImmutableList.of(ALTER_SKIPPING) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - super.cancelJobRun(applicationId, jobId, allowExceptionPropagation); - throw new IllegalArgumentException("Unknown Error"); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - // Mock flint index - mockDS.createIndex(); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - mockDS.updateIndexOptions(existingOptions, false); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. alter index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals("Internal Server Error.", asyncQueryExecutionResponse.getError()); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(1); - emrsClient.getJobRunResultCalled(0); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = mockDS.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - } -} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/IndexQuerySpecTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/IndexQuerySpecTest.java deleted file mode 100644 index 920981abf1e..00000000000 --- a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/IndexQuerySpecTest.java +++ /dev/null @@ -1,1065 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.spark.asyncquery; - -import com.amazonaws.services.emrserverless.model.CancelJobRunResult; -import com.amazonaws.services.emrserverless.model.GetJobRunResult; -import com.amazonaws.services.emrserverless.model.JobRun; -import com.amazonaws.services.emrserverless.model.ValidationException; -import com.google.common.collect.ImmutableList; -import org.junit.Assert; -import org.junit.Test; -import org.junit.jupiter.api.Assertions; -import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; -import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; -import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; -import org.opensearch.sql.spark.client.EMRServerlessClientFactory; -import org.opensearch.sql.spark.flint.FlintIndexState; -import org.opensearch.sql.spark.flint.FlintIndexType; -import org.opensearch.sql.spark.leasemanager.ConcurrencyLimitExceededException; -import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; -import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; -import org.opensearch.sql.spark.rest.model.LangType; - -public class IndexQuerySpecTest extends AsyncQueryExecutorServiceSpec { - private final String specialName = "`test ,:\"+/\\|?#><`"; - private final String encodedName = "test%20%2c%3a%22%2b%2f%5c%7c%3f%23%3e%3c"; - - public final String REFRESH_SI = "REFRESH SKIPPING INDEX on mys3.default.http_logs"; - public final String REFRESH_CI = "REFRESH INDEX covering ON mys3.default.http_logs"; - public final String REFRESH_MV = "REFRESH MATERIALIZED VIEW mys3.default.http_logs_metrics"; - public final String REFRESH_SCI = "REFRESH SKIPPING INDEX on mys3.default." + specialName; - - public final FlintDatasetMock LEGACY_SKIPPING = - new FlintDatasetMock( - "DROP SKIPPING INDEX ON mys3.default.http_logs", - REFRESH_SI, - FlintIndexType.SKIPPING, - "flint_mys3_default_http_logs_skipping_index") - .isLegacy(true); - public final FlintDatasetMock LEGACY_COVERING = - new FlintDatasetMock( - "DROP INDEX covering ON mys3.default.http_logs", - REFRESH_CI, - FlintIndexType.COVERING, - "flint_mys3_default_http_logs_covering_index") - .isLegacy(true); - public final FlintDatasetMock LEGACY_MV = - new FlintDatasetMock( - "DROP MATERIALIZED VIEW mys3.default.http_logs_metrics", - REFRESH_MV, - FlintIndexType.MATERIALIZED_VIEW, - "flint_mys3_default_http_logs_metrics") - .isLegacy(true); - - public final FlintDatasetMock LEGACY_SPECIAL_CHARACTERS = - new FlintDatasetMock( - "DROP SKIPPING INDEX ON mys3.default." + specialName, - REFRESH_SCI, - FlintIndexType.SKIPPING, - "flint_mys3_default_" + encodedName + "_skipping_index") - .isLegacy(true) - .isSpecialCharacter(true); - - public final FlintDatasetMock SKIPPING = - new FlintDatasetMock( - "DROP SKIPPING INDEX ON mys3.default.http_logs", - REFRESH_SI, - FlintIndexType.SKIPPING, - "flint_mys3_default_http_logs_skipping_index") - .latestId("ZmxpbnRfbXlzM19kZWZhdWx0X2h0dHBfbG9nc19za2lwcGluZ19pbmRleA=="); - public final FlintDatasetMock COVERING = - new FlintDatasetMock( - "DROP INDEX covering ON mys3.default.http_logs", - REFRESH_CI, - FlintIndexType.COVERING, - "flint_mys3_default_http_logs_covering_index") - .latestId("ZmxpbnRfbXlzM19kZWZhdWx0X2h0dHBfbG9nc19jb3ZlcmluZ19pbmRleA=="); - public final FlintDatasetMock MV = - new FlintDatasetMock( - "DROP MATERIALIZED VIEW mys3.default.http_logs_metrics", - REFRESH_MV, - FlintIndexType.MATERIALIZED_VIEW, - "flint_mys3_default_http_logs_metrics") - .latestId("ZmxpbnRfbXlzM19kZWZhdWx0X2h0dHBfbG9nc19tZXRyaWNz"); - public final FlintDatasetMock SPECIAL_CHARACTERS = - new FlintDatasetMock( - "DROP SKIPPING INDEX ON mys3.default." + specialName, - REFRESH_SCI, - FlintIndexType.SKIPPING, - "flint_mys3_default_" + encodedName + "_skipping_index") - .isSpecialCharacter(true) - .latestId( - "ZmxpbnRfbXlzM19kZWZhdWx0X3Rlc3QlMjAlMmMlM2ElMjIlMmIlMmYlNWMlN2MlM2YlMjMlM2UlM2Nfc2tpcHBpbmdfaW5kZXg="); - - public final String CREATE_SI_AUTO = - "CREATE SKIPPING INDEX ON mys3.default.http_logs" - + "(l_orderkey VALUE_SET) WITH (auto_refresh = true)"; - - public final String CREATE_CI_AUTO = - "CREATE INDEX covering ON mys3.default.http_logs " - + "(l_orderkey, l_quantity) WITH (auto_refresh = true)"; - - public final String CREATE_MV_AUTO = - "CREATE MATERIALIZED VIEW mys3.default.http_logs_metrics AS select * " - + "from mys3.default.https WITH (auto_refresh = true)"; - - /** - * Happy case. expectation is - * - *

(1) Drop Index response is SUCCESS - */ - @Test - public void legacyBasicDropAndFetchAndCancel() { - ImmutableList.of(LEGACY_SKIPPING, LEGACY_COVERING, LEGACY_SPECIAL_CHARACTERS) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - - // 1.drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - assertNotNull(response.getQueryId()); - assertTrue(clusterService.state().routingTable().hasIndex(mockDS.indexName)); - - // 2.fetch result - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryResults.getStatus()); - assertNull(asyncQueryResults.getError()); - emrsClient.cancelJobRunCalled(1); - - // 3.cancel - IllegalArgumentException exception = - assertThrows( - IllegalArgumentException.class, - () -> - asyncQueryExecutorService.cancelQuery( - response.getQueryId(), asyncQueryRequestContext)); - assertEquals("can't cancel index DML query", exception.getMessage()); - }); - } - - /** - * Legacy Test, without state index support. Not EMR-S job running. expectation is - * - *

(1) Drop Index response is SUCCESS - */ - @Test - public void legacyDropIndexNoJobRunning() { - ImmutableList.of(LEGACY_SKIPPING, LEGACY_COVERING, LEGACY_MV, LEGACY_SPECIAL_CHARACTERS) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - throw new ValidationException("Job run is not in a cancellable state"); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - - // 1.drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2.fetch result. - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryResults.getStatus()); - assertNull(asyncQueryResults.getError()); - }); - } - - /** - * Legacy Test, without state index support. Cancel EMR-S job call timeout. expectation is - * - *

(1) Drop Index response is FAILED - */ - @Test - public void legacyDropIndexCancelJobTimeout() { - ImmutableList.of(LEGACY_SKIPPING, LEGACY_COVERING, LEGACY_MV, LEGACY_SPECIAL_CHARACTERS) - .forEach( - mockDS -> { - // Mock EMR-S always return running. - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - return new GetJobRunResult().withJobRun(new JobRun().withState("Running")); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - - // 1. drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryResults.getStatus()); - assertEquals("Cancel job operation timed out.", asyncQueryResults.getError()); - }); - } - - /** - * Legacy Test, without state index support. Not EMR-S job running. expectation is - * - *

(1) Drop Index response is SUCCESS - */ - @Test - public void legacyDropIndexSpecialCharacter() { - FlintDatasetMock mockDS = LEGACY_SPECIAL_CHARACTERS; - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - throw new ValidationException("Job run is not in a cancellable state"); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - - // 1.drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest(mockDS.query, MYGLUE_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2.fetch result. - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryResults.getStatus()); - assertNull(asyncQueryResults.getError()); - } - - /** - * Happy case. expectation is - * - *

(1) Drop Index response is SUCCESS (2) change index state to: DELETED - */ - @Test - public void dropAndFetchAndCancel() { - ImmutableList.of(SKIPPING, COVERING, MV, SPECIAL_CHARACTERS) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // 1.drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - assertNotNull(response.getQueryId()); - assertTrue(clusterService.state().routingTable().hasIndex(mockDS.indexName)); - - // assert state is DELETED - flintIndexJob.assertState(FlintIndexState.DELETED); - - // 2.fetch result - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryResults.getStatus()); - assertNull(asyncQueryResults.getError()); - emrsClient.cancelJobRunCalled(1); - - // 3.cancel - IllegalArgumentException exception = - assertThrows( - IllegalArgumentException.class, - () -> - asyncQueryExecutorService.cancelQuery( - response.getQueryId(), asyncQueryRequestContext)); - assertEquals("can't cancel index DML query", exception.getMessage()); - }); - } - - /** - * Cancel EMR-S job, but not job running. expectation is - * - *

(1) Drop Index response is SUCCESS (2) change index state to: DELETED - */ - @Test - public void dropIndexNoJobRunning() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - // Mock EMR-S job is not running - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - throw new ValidationException("Job run is not in a cancellable state"); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state in refresh state. - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // 1.drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2.fetch result. - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryResults.getStatus()); - assertNull(asyncQueryResults.getError()); - - flintIndexJob.assertState(FlintIndexState.DELETED); - }); - } - - /** - * Cancel EMR-S job call timeout, expectation is - * - *

(1) Drop Index response is failed, (2) change index state to: CANCELLING - */ - @Test - public void dropIndexCancelJobTimeout() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - // Mock EMR-S always return running. - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - return new GetJobRunResult().withJobRun(new JobRun().withState("Running")); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // 1. drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryResults.getStatus()); - assertEquals("Cancel job operation timed out.", asyncQueryResults.getError()); - flintIndexJob.assertState(FlintIndexState.REFRESHING); - }); - } - - /** - * Drop Index operation is retryable, expectation is - * - *

(1) call EMR-S (2) change index state to: DELETED - */ - @Test - public void dropIndexWithIndexInRefreshingState() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // 1. drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - assertEquals( - "SUCCESS", - asyncQueryExecutorService - .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) - .getStatus()); - - flintIndexJob.assertState(FlintIndexState.DELETED); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(1); - emrsClient.getJobRunResultCalled(1); - }); - } - - /** - * Index state is stable, Drop Index operation is retryable, expectation is - * - *

(1) call EMR-S (2) change index state to: DELETED - */ - @Test - public void dropIndexWithIndexInActiveState() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - flintIndexJob.active(); - - // 1. drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); - flintIndexJob.assertState(FlintIndexState.DELETED); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(1); - emrsClient.getJobRunResultCalled(1); - }); - } - - /** - * Index state is stable, expectation is - * - *

(1) call EMR-S (2) change index state to: DELETED - */ - @Test - public void dropIndexWithIndexInCreatingState() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - flintIndexJob.creating(); - - // 1. drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - assertEquals( - "SUCCESS", - asyncQueryExecutorService - .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) - .getStatus()); - - flintIndexJob.assertState(FlintIndexState.DELETED); - }); - } - - /** - * Index state is stable, Drop Index operation is retryable, expectation is - * - *

(1) call EMR-S (2) change index state to: DELETED - */ - @Test - public void dropIndexWithIndexInEmptyState() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - - // 1. drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - assertEquals( - "SUCCESS", - asyncQueryExecutorService - .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) - .getStatus()); - - flintIndexJob.assertState(FlintIndexState.DELETED); - }); - } - - /** - * Couldn't acquire lock as the index is in transitioning state. Will result in error. - * - *

(1) not call EMR-S (2) change index state to: DELETED - */ - @Test - public void dropIndexWithIndexInDeletedState() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - Assert.fail("should not call cancelJobRun"); - return null; - } - - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - Assert.fail("should not call getJobRunResult"); - return null; - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - flintIndexJob.deleting(); - - // 1. drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - AsyncQueryExecutionResponse asyncQueryExecutionResponse = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - // 2. fetch result - assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); - assertEquals( - "Transaction failed as flint index is not in a valid state.", - asyncQueryExecutionResponse.getError()); - flintIndexJob.assertState(FlintIndexState.DELETING); - }); - } - - /** - * Cancel EMR-S job, but not job running. expectation is - * - *

(1) Drop Index response is SUCCESS (2) change index state to: DELETED - */ - @Test - public void dropIndexSpecialCharacter() { - FlintDatasetMock mockDS = SPECIAL_CHARACTERS; - // Mock EMR-S job is not running - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - throw new IllegalArgumentException("Job run is not in a cancellable state"); - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - // Mock index state in refresh state. - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob(flintIndexStateModelService, mockDS.latestId, MYGLUE_DATASOURCE); - flintIndexJob.refreshing(); - - // 1.drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest(mockDS.query, MYGLUE_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2.fetch result. - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryResults.getStatus()); - assertEquals("Internal Server Error.", asyncQueryResults.getError()); - - flintIndexJob.assertState(FlintIndexState.REFRESHING); - } - - /** - * No Job running, expectation is - * - *

(1) not call EMR-S (2) change index state to: DELETED - */ - @Test - public void edgeCaseNoIndexStateDoc() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - Assert.fail("should not call cancelJobRun"); - return null; - } - - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - Assert.fail("should not call getJobRunResult"); - return null; - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - mockDS.createIndex(); - - // 1. drop index - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. fetch result - AsyncQueryExecutionResponse asyncQueryResults = - asyncQueryExecutorService.getAsyncQueryResults( - response.getQueryId(), asyncQueryRequestContext); - assertEquals("FAILED", asyncQueryResults.getStatus()); - assertTrue(asyncQueryResults.getError().contains("no state found")); - }); - } - - @Test - public void concurrentRefreshJobLimitNotApplied() { - EMRServerlessClientFactory emrServerlessClientFactory = new LocalEMRServerlessClientFactory(); - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // Mock flint index - COVERING.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob(flintIndexStateModelService, COVERING.latestId, MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // query with auto refresh - String query = - "CREATE INDEX covering ON mys3.default.http_logs(l_orderkey, " - + "l_quantity) WITH (auto_refresh = true)"; - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - assertNull(response.getSessionId()); - } - - @Test - public void concurrentRefreshJobLimitAppliedToDDLWithAuthRefresh() { - EMRServerlessClientFactory emrServerlessClientFactory = new LocalEMRServerlessClientFactory(); - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - setConcurrentRefreshJob(1); - - // Mock flint index - COVERING.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob(flintIndexStateModelService, COVERING.latestId, MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // query with auto_refresh = true. - String query = - "CREATE INDEX covering ON mys3.default.http_logs(l_orderkey, " - + "l_quantity) WITH (auto_refresh = true)"; - ConcurrencyLimitExceededException exception = - assertThrows( - ConcurrencyLimitExceededException.class, - () -> - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext)); - assertEquals("domain concurrent refresh job can not exceed 1", exception.getMessage()); - } - - @Test - public void concurrentRefreshJobLimitAppliedToRefresh() { - EMRServerlessClientFactory emrServerlessClientFactory = new LocalEMRServerlessClientFactory(); - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - setConcurrentRefreshJob(1); - - // Mock flint index - COVERING.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob(flintIndexStateModelService, COVERING.latestId, MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - // query with auto_refresh = true. - String query = "REFRESH INDEX covering ON mys3.default.http_logs"; - ConcurrencyLimitExceededException exception = - assertThrows( - ConcurrencyLimitExceededException.class, - () -> - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext)); - assertEquals("domain concurrent refresh job can not exceed 1", exception.getMessage()); - } - - @Test - public void concurrentRefreshJobLimitNotAppliedToDDL() { - String query = "CREATE INDEX covering ON mys3.default.http_logs(l_orderkey, l_quantity)"; - EMRServerlessClientFactory emrServerlessClientFactory = new LocalEMRServerlessClientFactory(); - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - setConcurrentRefreshJob(1); - - // Mock flint index - COVERING.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob(flintIndexStateModelService, COVERING.latestId, MYS3_DATASOURCE); - flintIndexJob.refreshing(); - - CreateAsyncQueryResponse asyncQueryResponse = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - } - - /** Cancel create flint index statement with auto_refresh=true, should throw exception. */ - @Test - public void cancelAutoRefreshCreateFlintIndexShouldThrowException() { - ImmutableList.of(CREATE_SI_AUTO, CREATE_CI_AUTO, CREATE_MV_AUTO) - .forEach( - query -> { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public CancelJobRunResult cancelJobRun( - String applicationId, String jobId, boolean allowExceptionPropagation) { - Assert.fail("should not call cancelJobRun"); - return null; - } - - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - Assert.fail("should not call getJobRunResult"); - return null; - } - }; - EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService(emrServerlessClientFactory); - - // 1. submit create / refresh index query - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - - // 2. cancel query - IllegalArgumentException exception = - assertThrows( - IllegalArgumentException.class, - () -> - asyncQueryExecutorService.cancelQuery( - response.getQueryId(), asyncQueryRequestContext)); - assertEquals( - "can't cancel index DML query, using ALTER auto_refresh=off statement to stop" - + " job, using VACUUM statement to stop job and delete data", - exception.getMessage()); - }); - } - - /** Cancel REFRESH statement should success */ - @Test - public void cancelRefreshStatement() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService( - (accountId) -> - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult( - String applicationId, String jobId) { - return new GetJobRunResult() - .withJobRun(new JobRun().withState("Cancelled")); - } - }); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - - // 1. Submit REFRESH statement - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.refreshQuery, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - // mock index state. - flintIndexJob.refreshing(); - - // 2. Cancel query - String cancelResponse = - asyncQueryExecutorService.cancelQuery( - response.getQueryId(), asyncQueryRequestContext); - - assertNotNull(cancelResponse); - assertTrue(clusterService.state().routingTable().hasIndex(mockDS.indexName)); - - // assert state is active - flintIndexJob.assertState(FlintIndexState.ACTIVE); - }); - } - - /** Cancel REFRESH statement should success */ - @Test - public void cancelRefreshStatementWithActiveState() { - ImmutableList.of(SKIPPING, COVERING, MV) - .forEach( - mockDS -> { - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService( - (accountId) -> - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult( - String applicationId, String jobId) { - return new GetJobRunResult() - .withJobRun(new JobRun().withState("Cancelled")); - } - }); - - // Mock flint index - mockDS.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); - - // 1. Submit REFRESH statement - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - mockDS.refreshQuery, MYS3_DATASOURCE, LangType.SQL, null), - asyncQueryRequestContext); - // mock index state. - flintIndexJob.active(); - - // 2. Cancel query - IllegalStateException illegalStateException = - Assertions.assertThrows( - IllegalStateException.class, - () -> - asyncQueryExecutorService.cancelQuery( - response.getQueryId(), asyncQueryRequestContext)); - Assertions.assertEquals( - "Transaction failed as flint index is not in a valid state.", - illegalStateException.getMessage()); - - // assert state is active - flintIndexJob.assertState(FlintIndexState.ACTIVE); - }); - } - - @Test - public void cancelRefreshStatementWithFailureInFetchingIndexMetadata() { - String indexName = "flint_my_glue_mydb_http_logs_covering_corrupted_index"; - MockFlintIndex mockFlintIndex = - new MockFlintIndex(client(), indexName, FlintIndexType.COVERING, null); - AsyncQueryExecutorService asyncQueryExecutorService = - createAsyncQueryExecutorService( - (accountId) -> - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); - } - }); - - mockFlintIndex.createIndex(); - // Mock index state - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, indexName + "_latest_id", MYS3_DATASOURCE); - - // 1. Submit REFRESH statement - CreateAsyncQueryResponse response = - asyncQueryExecutorService.createAsyncQuery( - new CreateAsyncQueryRequest( - "REFRESH INDEX covering_corrupted ON my_glue.mydb.http_logs", - MYS3_DATASOURCE, - LangType.SQL, - null), - asyncQueryRequestContext); - // mock index state. - flintIndexJob.refreshing(); - - // 2. Cancel query - Assertions.assertThrows( - IllegalStateException.class, - () -> - asyncQueryExecutorService.cancelQuery(response.getQueryId(), asyncQueryRequestContext)); - } -} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterConversionTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterConversionTest.java new file mode 100644 index 00000000000..2c2dc0714f3 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterConversionTest.java @@ -0,0 +1,240 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import com.amazonaws.services.emrserverless.model.CancelJobRunResult; +import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.amazonaws.services.emrserverless.model.JobRun; +import com.amazonaws.services.emrserverless.model.ValidationException; +import com.google.common.collect.ImmutableList; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorServiceSpec; +import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.client.EMRServerlessClientFactory; +import org.opensearch.sql.spark.client.StartJobRequest; +import org.opensearch.sql.spark.flint.FlintIndexState; +import org.opensearch.sql.spark.flint.FlintIndexType; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; +import org.opensearch.sql.spark.rest.model.LangType; + +import java.util.HashMap; +import java.util.Map; + +public class IndexQuerySpecAlterConversionTest extends AsyncQueryExecutorServiceSpec { + @Test + public void testAlterIndexQueryConvertingToManualRefresh() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false)"); + MockFlintIndex ALTER_COVERING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false)"); + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," + + " incremental_refresh=false) "); + ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(1); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryConvertingToManualRefreshWithNoIncrementalRefresh() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false)"); + MockFlintIndex ALTER_COVERING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false)"); + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false)"); + ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + existingOptions.put("checkpoint_location", "s3://checkpoint/location"); + mockDS.updateIndexOptions(existingOptions, true); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(1); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryConvertingToAutoRefresh() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=true," + + " incremental_refresh=false)"); + MockFlintIndex ALTER_COVERING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=true," + + " incremental_refresh=false)"); + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=true," + + " incremental_refresh=false) "); + ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = new LocalEMRSClient(); + EMRServerlessClientFactory clientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(clientFactory); + + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "false"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + assertEquals( + "RUNNING", + asyncQueryExecutorService + .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) + .getStatus()); + + flintIndexJob.assertState(FlintIndexState.ACTIVE); + emrsClient.startJobRunCalled(1); + emrsClient.getJobRunResultCalled(1); + emrsClient.cancelJobRunCalled(0); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterIncrementalTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterIncrementalTest.java new file mode 100644 index 00000000000..47b6a239033 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterIncrementalTest.java @@ -0,0 +1,430 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import com.amazonaws.services.emrserverless.model.CancelJobRunResult; +import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.amazonaws.services.emrserverless.model.JobRun; +import com.amazonaws.services.emrserverless.model.ValidationException; +import com.google.common.collect.ImmutableList; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorServiceSpec; +import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.client.EMRServerlessClientFactory; +import org.opensearch.sql.spark.client.StartJobRequest; +import org.opensearch.sql.spark.flint.FlintIndexState; +import org.opensearch.sql.spark.flint.FlintIndexType; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; +import org.opensearch.sql.spark.rest.model.LangType; + +import java.util.HashMap; +import java.util.Map; + +public class IndexQuerySpecAlterIncrementalTest extends AsyncQueryExecutorServiceSpec { + @Test + public void testAlterIndexQueryOfIncrementalRefreshWithInvalidOptions() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=true, output_mode=\"complete\")"); + MockFlintIndex ALTER_COVERING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=true, output_mode=\"complete\")"); + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," + + " incremental_refresh=true, output_mode=\"complete\") "); + ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals( + "Altering to incremental refresh only allows: [auto_refresh, incremental_refresh," + + " watermark_delay, checkpoint_location] options", + asyncQueryExecutionResponse.getError()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("true", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryOfIncrementalRefreshWithInsufficientOptions() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=true)"); + MockFlintIndex ALTER_COVERING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=true)"); + ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + existingOptions.put("incremental_refresh", "false"); + mockDS.updateIndexOptions(existingOptions, true); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals( + "Conversion to incremental refresh index cannot proceed due to missing" + + " attributes: checkpoint_location.", + asyncQueryExecutionResponse.getError()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("true", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryOfIncrementalRefreshWithInsufficientOptionsForMV() { + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," + + " incremental_refresh=true) "); + ImmutableList.of(ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + existingOptions.put("incremental_refresh", "false"); + mockDS.updateIndexOptions(existingOptions, true); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals( + "Conversion to incremental refresh index cannot proceed due to missing" + + " attributes: checkpoint_location, watermark_delay.", + asyncQueryExecutionResponse.getError()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("true", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryOfIncrementalRefreshWithEmptyExistingOptionsForMV() { + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," + + " incremental_refresh=true) "); + ImmutableList.of(ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + existingOptions.put("incremental_refresh", "false"); + existingOptions.put("watermark_delay", ""); + existingOptions.put("checkpoint_location", ""); + mockDS.updateIndexOptions(existingOptions, true); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals( + "Conversion to incremental refresh index cannot proceed due to missing" + + " attributes: checkpoint_location, watermark_delay.", + asyncQueryExecutionResponse.getError()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("true", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryOfIncrementalRefresh() { + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," + + " incremental_refresh=true) "); + ImmutableList.of(ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + existingOptions.put("incremental_refresh", "false"); + existingOptions.put("watermark_delay", "watermark_delay"); + existingOptions.put("checkpoint_location", "s3://checkpoint/location"); + mockDS.updateIndexOptions(existingOptions, true); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); + emrsClient.startJobRunCalled(0); + emrsClient.getJobRunResultCalled(1); + emrsClient.cancelJobRunCalled(1); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + Assertions.assertEquals("true", options.get("incremental_refresh")); + }); + } + + @Test + public void testAlterIndexQueryWithIncrementalRefreshAlreadyExisting() { + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false) "); + ImmutableList.of(ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + existingOptions.put("incremental_refresh", "true"); + existingOptions.put("watermark_delay", "watermark_delay"); + existingOptions.put("checkpoint_location", "s3://checkpoint/location"); + mockDS.updateIndexOptions(existingOptions, true); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); + emrsClient.startJobRunCalled(0); + emrsClient.getJobRunResultCalled(1); + emrsClient.cancelJobRunCalled(1); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + Assertions.assertEquals("true", options.get("incremental_refresh")); + }); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterTest.java new file mode 100644 index 00000000000..f7e14c80bb5 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecAlterTest.java @@ -0,0 +1,517 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import com.amazonaws.services.emrserverless.model.CancelJobRunResult; +import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.amazonaws.services.emrserverless.model.JobRun; +import com.amazonaws.services.emrserverless.model.ValidationException; +import com.google.common.collect.ImmutableList; +import java.util.HashMap; +import java.util.Map; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorServiceSpec; +import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.client.EMRServerlessClientFactory; +import org.opensearch.sql.spark.client.StartJobRequest; +import org.opensearch.sql.spark.flint.FlintIndexState; +import org.opensearch.sql.spark.flint.FlintIndexType; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; +import org.opensearch.sql.spark.rest.model.LangType; + +public class IndexQuerySpecAlterTest extends AsyncQueryExecutorServiceSpec { + @Test + public void testAlterIndexQueryWithRedundantOperation() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false)"); + MockFlintIndex ALTER_COVERING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false)"); + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," + + " incremental_refresh=false) "); + ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public String startJobRun(StartJobRequest startJobRequest) { + return "jobId"; + } + + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + super.cancelJobRun(applicationId, jobId, allowExceptionPropagation); + throw new ValidationException("Job run is not in a cancellable state"); + } + }; + EMRServerlessClientFactory emrServerlessCientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessCientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "false"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(1); + emrsClient.getJobRunResultCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryWithOutAnyAutoRefresh() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (" + + " incremental_refresh=false)"); + MockFlintIndex ALTER_COVERING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (" + + " incremental_refresh=false)"); + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (" + " incremental_refresh=false) "); + ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = new LocalEMRSClient(); + EMRServerlessClientFactory clientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(clientFactory); + + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "false"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + assertEquals( + "RUNNING", + asyncQueryExecutorService + .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) + .getStatus()); + + flintIndexJob.assertState(FlintIndexState.ACTIVE); + emrsClient.startJobRunCalled(1); + emrsClient.getJobRunResultCalled(1); + emrsClient.cancelJobRunCalled(0); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryOfFullRefreshWithInvalidOptions() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false, checkpoint_location=\"s3://ckp/skp\")"); + MockFlintIndex ALTER_COVERING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false, checkpoint_location=\"s3://ckp/skp\")"); + MockFlintIndex ALTER_MV = + new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," + + " incremental_refresh=false, checkpoint_location=\"s3://ckp/skp\") "); + ImmutableList.of(ALTER_SKIPPING, ALTER_COVERING, ALTER_MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals( + "Altering to full refresh only allows: [auto_refresh, incremental_refresh]" + + " options", + asyncQueryExecutionResponse.getError()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("true", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryWithInvalidInitialState() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false)"); + ImmutableList.of(ALTER_SKIPPING) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.updating(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals( + "Transaction failed as flint index is not in a valid state.", + asyncQueryExecutionResponse.getError()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + flintIndexJob.assertState(FlintIndexState.UPDATING); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("true", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryWithValidationExceptionWithSuccess() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false)"); + ImmutableList.of(ALTER_SKIPPING) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + super.cancelJobRun(applicationId, jobId, allowExceptionPropagation); + throw new ValidationException("Job run is not in a cancellable state"); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(1); + emrsClient.getJobRunResultCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryWithResourceNotFoundExceptionWithSuccess() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false)"); + ImmutableList.of(ALTER_SKIPPING) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + super.cancelJobRun(applicationId, jobId, allowExceptionPropagation); + throw new ValidationException("Random validation exception"); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals("Internal Server Error.", asyncQueryExecutionResponse.getError()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(1); + emrsClient.getJobRunResultCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + } + + @Test + public void testAlterIndexQueryWithUnknownError() { + MockFlintIndex ALTER_SKIPPING = + new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=false)"); + ImmutableList.of(ALTER_SKIPPING) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + super.cancelJobRun(applicationId, jobId, allowExceptionPropagation); + throw new IllegalArgumentException("Unknown Error"); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + // Mock flint index + mockDS.createIndex(); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + mockDS.updateIndexOptions(existingOptions, false); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.getLatestId(), MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. alter index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.getQuery(), MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals("Internal Server Error.", asyncQueryExecutionResponse.getError()); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(1); + emrsClient.getJobRunResultCalled(0); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = mockDS.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecBase.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecBase.java new file mode 100644 index 00000000000..07441b21600 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecBase.java @@ -0,0 +1,93 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorServiceSpec; +import org.opensearch.sql.spark.flint.FlintIndexType; + +public class IndexQuerySpecBase extends AsyncQueryExecutorServiceSpec { + private final String specialName = "`test ,:\"+/\\|?#><`"; + private final String encodedName = "test%20%2c%3a%22%2b%2f%5c%7c%3f%23%3e%3c"; + + public final String REFRESH_SI = "REFRESH SKIPPING INDEX on mys3.default.http_logs"; + public final String REFRESH_CI = "REFRESH INDEX covering ON mys3.default.http_logs"; + public final String REFRESH_MV = "REFRESH MATERIALIZED VIEW mys3.default.http_logs_metrics"; + public final String REFRESH_SCI = "REFRESH SKIPPING INDEX on mys3.default." + specialName; + + public final FlintDatasetMock LEGACY_SKIPPING = + new FlintDatasetMock( + "DROP SKIPPING INDEX ON mys3.default.http_logs", + REFRESH_SI, + FlintIndexType.SKIPPING, + "flint_mys3_default_http_logs_skipping_index") + .isLegacy(true); + public final FlintDatasetMock LEGACY_COVERING = + new FlintDatasetMock( + "DROP INDEX covering ON mys3.default.http_logs", + REFRESH_CI, + FlintIndexType.COVERING, + "flint_mys3_default_http_logs_covering_index") + .isLegacy(true); + public final FlintDatasetMock LEGACY_MV = + new FlintDatasetMock( + "DROP MATERIALIZED VIEW mys3.default.http_logs_metrics", + REFRESH_MV, + FlintIndexType.MATERIALIZED_VIEW, + "flint_mys3_default_http_logs_metrics") + .isLegacy(true); + + public final FlintDatasetMock LEGACY_SPECIAL_CHARACTERS = + new FlintDatasetMock( + "DROP SKIPPING INDEX ON mys3.default." + specialName, + REFRESH_SCI, + FlintIndexType.SKIPPING, + "flint_mys3_default_" + encodedName + "_skipping_index") + .isLegacy(true) + .isSpecialCharacter(true); + + public final FlintDatasetMock SKIPPING = + new FlintDatasetMock( + "DROP SKIPPING INDEX ON mys3.default.http_logs", + REFRESH_SI, + FlintIndexType.SKIPPING, + "flint_mys3_default_http_logs_skipping_index") + .latestId("ZmxpbnRfbXlzM19kZWZhdWx0X2h0dHBfbG9nc19za2lwcGluZ19pbmRleA=="); + public final FlintDatasetMock COVERING = + new FlintDatasetMock( + "DROP INDEX covering ON mys3.default.http_logs", + REFRESH_CI, + FlintIndexType.COVERING, + "flint_mys3_default_http_logs_covering_index") + .latestId("ZmxpbnRfbXlzM19kZWZhdWx0X2h0dHBfbG9nc19jb3ZlcmluZ19pbmRleA=="); + public final FlintDatasetMock MV = + new FlintDatasetMock( + "DROP MATERIALIZED VIEW mys3.default.http_logs_metrics", + REFRESH_MV, + FlintIndexType.MATERIALIZED_VIEW, + "flint_mys3_default_http_logs_metrics") + .latestId("ZmxpbnRfbXlzM19kZWZhdWx0X2h0dHBfbG9nc19tZXRyaWNz"); + public final FlintDatasetMock SPECIAL_CHARACTERS = + new FlintDatasetMock( + "DROP SKIPPING INDEX ON mys3.default." + specialName, + REFRESH_SCI, + FlintIndexType.SKIPPING, + "flint_mys3_default_" + encodedName + "_skipping_index") + .isSpecialCharacter(true) + .latestId( + "ZmxpbnRfbXlzM19kZWZhdWx0X3Rlc3QlMjAlMmMlM2ElMjIlMmIlMmYlNWMlN2MlM2YlMjMlM2UlM2Nfc2tpcHBpbmdfaW5kZXg="); + + public final String CREATE_SI_AUTO = + "CREATE SKIPPING INDEX ON mys3.default.http_logs" + + "(l_orderkey VALUE_SET) WITH (auto_refresh = true)"; + + public final String CREATE_CI_AUTO = + "CREATE INDEX covering ON mys3.default.http_logs " + + "(l_orderkey, l_quantity) WITH (auto_refresh = true)"; + + public final String CREATE_MV_AUTO = + "CREATE MATERIALIZED VIEW mys3.default.http_logs_metrics AS select * " + + "from mys3.default.https WITH (auto_refresh = true)"; +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecCancelTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecCancelTest.java new file mode 100644 index 00000000000..571906387d6 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecCancelTest.java @@ -0,0 +1,209 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import com.amazonaws.services.emrserverless.model.CancelJobRunResult; +import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.amazonaws.services.emrserverless.model.JobRun; +import com.google.common.collect.ImmutableList; +import org.junit.Assert; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService; +import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.client.EMRServerlessClientFactory; +import org.opensearch.sql.spark.flint.FlintIndexState; +import org.opensearch.sql.spark.flint.FlintIndexType; +import org.opensearch.sql.spark.leasemanager.ConcurrencyLimitExceededException; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; +import org.opensearch.sql.spark.rest.model.LangType; + +public class IndexQuerySpecCancelTest extends IndexQuerySpecBase { + /** Cancel create flint index statement with auto_refresh=true, should throw exception. */ + @Test + public void cancelAutoRefreshCreateFlintIndexShouldThrowException() { + ImmutableList.of(CREATE_SI_AUTO, CREATE_CI_AUTO, CREATE_MV_AUTO) + .forEach( + query -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + Assert.fail("should not call cancelJobRun"); + return null; + } + + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + Assert.fail("should not call getJobRunResult"); + return null; + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // 1. submit create / refresh index query + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. cancel query + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + asyncQueryExecutorService.cancelQuery( + response.getQueryId(), asyncQueryRequestContext)); + assertEquals( + "can't cancel index DML query, using ALTER auto_refresh=off statement to stop" + + " job, using VACUUM statement to stop job and delete data", + exception.getMessage()); + }); + } + + /** Cancel REFRESH statement should success */ + @Test + public void cancelRefreshStatement() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService( + (accountId) -> + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult( + String applicationId, String jobId) { + return new GetJobRunResult() + .withJobRun(new JobRun().withState("Cancelled")); + } + }); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + + // 1. Submit REFRESH statement + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.refreshQuery, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + // mock index state. + flintIndexJob.refreshing(); + + // 2. Cancel query + String cancelResponse = + asyncQueryExecutorService.cancelQuery( + response.getQueryId(), asyncQueryRequestContext); + + assertNotNull(cancelResponse); + assertTrue(clusterService.state().routingTable().hasIndex(mockDS.indexName)); + + // assert state is active + flintIndexJob.assertState(FlintIndexState.ACTIVE); + }); + } + + /** Cancel REFRESH statement should success */ + @Test + public void cancelRefreshStatementWithActiveState() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService( + (accountId) -> + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult( + String applicationId, String jobId) { + return new GetJobRunResult() + .withJobRun(new JobRun().withState("Cancelled")); + } + }); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + + // 1. Submit REFRESH statement + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.refreshQuery, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + // mock index state. + flintIndexJob.active(); + + // 2. Cancel query + IllegalStateException illegalStateException = + Assertions.assertThrows( + IllegalStateException.class, + () -> + asyncQueryExecutorService.cancelQuery( + response.getQueryId(), asyncQueryRequestContext)); + Assertions.assertEquals( + "Transaction failed as flint index is not in a valid state.", + illegalStateException.getMessage()); + + // assert state is active + flintIndexJob.assertState(FlintIndexState.ACTIVE); + }); + } + + @Test + public void cancelRefreshStatementWithFailureInFetchingIndexMetadata() { + String indexName = "flint_my_glue_mydb_http_logs_covering_corrupted_index"; + MockFlintIndex mockFlintIndex = + new MockFlintIndex(client(), indexName, FlintIndexType.COVERING, null); + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService( + (accountId) -> + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); + } + }); + + mockFlintIndex.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, indexName + "_latest_id", MYS3_DATASOURCE); + + // 1. Submit REFRESH statement + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + "REFRESH INDEX covering_corrupted ON my_glue.mydb.http_logs", + MYS3_DATASOURCE, + LangType.SQL, + null), + asyncQueryRequestContext); + // mock index state. + flintIndexJob.refreshing(); + + // 2. Cancel query + Assertions.assertThrows( + IllegalStateException.class, + () -> + asyncQueryExecutorService.cancelQuery(response.getQueryId(), asyncQueryRequestContext)); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecConcurrentTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecConcurrentTest.java new file mode 100644 index 00000000000..98395907207 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecConcurrentTest.java @@ -0,0 +1,119 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import org.junit.Test; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.client.EMRServerlessClientFactory; +import org.opensearch.sql.spark.leasemanager.ConcurrencyLimitExceededException; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; +import org.opensearch.sql.spark.rest.model.LangType; + +public class IndexQuerySpecConcurrentTest extends IndexQuerySpecBase { + @Test + public void concurrentRefreshJobLimitNotApplied() { + EMRServerlessClientFactory emrServerlessClientFactory = new LocalEMRServerlessClientFactory(); + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + COVERING.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob(flintIndexStateModelService, COVERING.latestId, MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // query with auto refresh + String query = + "CREATE INDEX covering ON mys3.default.http_logs(l_orderkey, " + + "l_quantity) WITH (auto_refresh = true)"; + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + assertNull(response.getSessionId()); + } + + @Test + public void concurrentRefreshJobLimitAppliedToDDLWithAuthRefresh() { + EMRServerlessClientFactory emrServerlessClientFactory = new LocalEMRServerlessClientFactory(); + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + setConcurrentRefreshJob(1); + + // Mock flint index + COVERING.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob(flintIndexStateModelService, COVERING.latestId, MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // query with auto_refresh = true. + String query = + "CREATE INDEX covering ON mys3.default.http_logs(l_orderkey, " + + "l_quantity) WITH (auto_refresh = true)"; + ConcurrencyLimitExceededException exception = + assertThrows( + ConcurrencyLimitExceededException.class, + () -> + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext)); + assertEquals("domain concurrent refresh job can not exceed 1", exception.getMessage()); + } + + @Test + public void concurrentRefreshJobLimitAppliedToRefresh() { + EMRServerlessClientFactory emrServerlessClientFactory = new LocalEMRServerlessClientFactory(); + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + setConcurrentRefreshJob(1); + + // Mock flint index + COVERING.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob(flintIndexStateModelService, COVERING.latestId, MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // query with auto_refresh = true. + String query = "REFRESH INDEX covering ON mys3.default.http_logs"; + ConcurrencyLimitExceededException exception = + assertThrows( + ConcurrencyLimitExceededException.class, + () -> + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext)); + assertEquals("domain concurrent refresh job can not exceed 1", exception.getMessage()); + } + + @Test + public void concurrentRefreshJobLimitNotAppliedToDDL() { + String query = "CREATE INDEX covering ON mys3.default.http_logs(l_orderkey, l_quantity)"; + EMRServerlessClientFactory emrServerlessClientFactory = new LocalEMRServerlessClientFactory(); + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + setConcurrentRefreshJob(1); + + // Mock flint index + COVERING.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob(flintIndexStateModelService, COVERING.latestId, MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + CreateAsyncQueryResponse asyncQueryResponse = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest(query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecDropTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecDropTest.java new file mode 100644 index 00000000000..ace0356e909 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecDropTest.java @@ -0,0 +1,229 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import com.amazonaws.services.emrserverless.model.CancelJobRunResult; +import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.amazonaws.services.emrserverless.model.JobRun; +import com.amazonaws.services.emrserverless.model.ValidationException; +import com.google.common.collect.ImmutableList; +import org.junit.Assert; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService; +import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.client.EMRServerlessClientFactory; +import org.opensearch.sql.spark.flint.FlintIndexState; +import org.opensearch.sql.spark.flint.FlintIndexType; +import org.opensearch.sql.spark.leasemanager.ConcurrencyLimitExceededException; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; +import org.opensearch.sql.spark.rest.model.LangType; + +public class IndexQuerySpecDropTest extends IndexQuerySpecBase { + /** + * Happy case. expectation is + * + *

(1) Drop Index response is SUCCESS (2) change index state to: DELETED + */ + @Test + public void dropAndFetchAndCancel() { + ImmutableList.of(SKIPPING, COVERING, MV, SPECIAL_CHARACTERS) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // 1.drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + assertNotNull(response.getQueryId()); + assertTrue(clusterService.state().routingTable().hasIndex(mockDS.indexName)); + + // assert state is DELETED + flintIndexJob.assertState(FlintIndexState.DELETED); + + // 2.fetch result + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryResults.getStatus()); + assertNull(asyncQueryResults.getError()); + emrsClient.cancelJobRunCalled(1); + + // 3.cancel + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + asyncQueryExecutorService.cancelQuery( + response.getQueryId(), asyncQueryRequestContext)); + assertEquals("can't cancel index DML query", exception.getMessage()); + }); + } + + /** + * Cancel EMR-S job, but not job running. expectation is + * + *

(1) Drop Index response is SUCCESS (2) change index state to: DELETED + */ + @Test + public void dropIndexNoJobRunning() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + // Mock EMR-S job is not running + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + throw new ValidationException("Job run is not in a cancellable state"); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state in refresh state. + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // 1.drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2.fetch result. + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryResults.getStatus()); + assertNull(asyncQueryResults.getError()); + + flintIndexJob.assertState(FlintIndexState.DELETED); + }); + } + + /** + * Cancel EMR-S job call timeout, expectation is + * + *

(1) Drop Index response is failed, (2) change index state to: CANCELLING + */ + @Test + public void dropIndexCancelJobTimeout() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + // Mock EMR-S always return running. + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + return new GetJobRunResult().withJobRun(new JobRun().withState("Running")); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // 1. drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryResults.getStatus()); + assertEquals("Cancel job operation timed out.", asyncQueryResults.getError()); + flintIndexJob.assertState(FlintIndexState.REFRESHING); + }); + } + + /** + * Cancel EMR-S job, but not job running. expectation is + * + *

(1) Drop Index response is SUCCESS (2) change index state to: DELETED + */ + @Test + public void dropIndexSpecialCharacter() { + FlintDatasetMock mockDS = SPECIAL_CHARACTERS; + // Mock EMR-S job is not running + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + throw new IllegalArgumentException("Job run is not in a cancellable state"); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state in refresh state. + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob(flintIndexStateModelService, mockDS.latestId, MYGLUE_DATASOURCE); + flintIndexJob.refreshing(); + + // 1.drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest(mockDS.query, MYGLUE_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2.fetch result. + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryResults.getStatus()); + assertEquals("Internal Server Error.", asyncQueryResults.getError()); + + flintIndexJob.assertState(FlintIndexState.REFRESHING); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecLegacyTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecLegacyTest.java new file mode 100644 index 00000000000..ed9bbc8b1d0 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecLegacyTest.java @@ -0,0 +1,200 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import com.amazonaws.services.emrserverless.model.CancelJobRunResult; +import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.amazonaws.services.emrserverless.model.JobRun; +import com.amazonaws.services.emrserverless.model.ValidationException; +import com.google.common.collect.ImmutableList; +import org.junit.Assert; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService; +import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.client.EMRServerlessClientFactory; +import org.opensearch.sql.spark.flint.FlintIndexState; +import org.opensearch.sql.spark.flint.FlintIndexType; +import org.opensearch.sql.spark.leasemanager.ConcurrencyLimitExceededException; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; +import org.opensearch.sql.spark.rest.model.LangType; + +public class IndexQuerySpecLegacyTest extends IndexQuerySpecBase { + /** + * Happy case. expectation is + * + *

(1) Drop Index response is SUCCESS + */ + @Test + public void legacyBasicDropAndFetchAndCancel() { + ImmutableList.of(LEGACY_SKIPPING, LEGACY_COVERING, LEGACY_SPECIAL_CHARACTERS) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + + // 1.drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + assertNotNull(response.getQueryId()); + assertTrue(clusterService.state().routingTable().hasIndex(mockDS.indexName)); + + // 2.fetch result + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryResults.getStatus()); + assertNull(asyncQueryResults.getError()); + emrsClient.cancelJobRunCalled(1); + + // 3.cancel + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + asyncQueryExecutorService.cancelQuery( + response.getQueryId(), asyncQueryRequestContext)); + assertEquals("can't cancel index DML query", exception.getMessage()); + }); + } + + /** + * Legacy Test, without state index support. Not EMR-S job running. expectation is + * + *

(1) Drop Index response is SUCCESS + */ + @Test + public void legacyDropIndexNoJobRunning() { + ImmutableList.of(LEGACY_SKIPPING, LEGACY_COVERING, LEGACY_MV, LEGACY_SPECIAL_CHARACTERS) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + throw new ValidationException("Job run is not in a cancellable state"); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + + // 1.drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2.fetch result. + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryResults.getStatus()); + assertNull(asyncQueryResults.getError()); + }); + } + + /** + * Legacy Test, without state index support. Cancel EMR-S job call timeout. expectation is + * + *

(1) Drop Index response is FAILED + */ + @Test + public void legacyDropIndexCancelJobTimeout() { + ImmutableList.of(LEGACY_SKIPPING, LEGACY_COVERING, LEGACY_MV, LEGACY_SPECIAL_CHARACTERS) + .forEach( + mockDS -> { + // Mock EMR-S always return running. + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + return new GetJobRunResult().withJobRun(new JobRun().withState("Running")); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + + // 1. drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryResults.getStatus()); + assertEquals("Cancel job operation timed out.", asyncQueryResults.getError()); + }); + } + + /** + * Legacy Test, without state index support. Not EMR-S job running. expectation is + * + *

(1) Drop Index response is SUCCESS + */ + @Test + public void legacyDropIndexSpecialCharacter() { + FlintDatasetMock mockDS = LEGACY_SPECIAL_CHARACTERS; + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + throw new ValidationException("Job run is not in a cancellable state"); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + + // 1.drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest(mockDS.query, MYGLUE_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2.fetch result. + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryResults.getStatus()); + assertNull(asyncQueryResults.getError()); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecStateTest.java b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecStateTest.java new file mode 100644 index 00000000000..14040663d16 --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/asyncquery/indexquery/IndexQuerySpecStateTest.java @@ -0,0 +1,334 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.asyncquery.indexquery; + +import com.amazonaws.services.emrserverless.model.CancelJobRunResult; +import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.amazonaws.services.emrserverless.model.JobRun; +import com.amazonaws.services.emrserverless.model.ValidationException; +import com.google.common.collect.ImmutableList; +import org.jetbrains.annotations.NotNull; +import org.junit.Assert; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService; +import org.opensearch.sql.spark.asyncquery.model.AsyncQueryExecutionResponse; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.client.EMRServerlessClientFactory; +import org.opensearch.sql.spark.flint.FlintIndexState; +import org.opensearch.sql.spark.flint.FlintIndexType; +import org.opensearch.sql.spark.leasemanager.ConcurrencyLimitExceededException; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest; +import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse; +import org.opensearch.sql.spark.rest.model.LangType; + +public class IndexQuerySpecStateTest extends IndexQuerySpecBase { + /** + * No Job running, expectation is + * + *

(1) not call EMR-S (2) change index state to: DELETED + */ + @Test + public void edgeCaseNoIndexStateDoc() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + EMRServerlessClientFactory emrServerlessClientFactory = getEmrServerlessClientFactory(); + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + + // 1. drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryResults = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("FAILED", asyncQueryResults.getStatus()); + assertTrue(asyncQueryResults.getError().contains("no state found")); + }); + } + + /** + * Drop Index operation is retryable, expectation is + * + *

(1) call EMR-S (2) change index state to: DELETED + */ + @Test + public void dropIndexWithIndexInRefreshingState() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + flintIndexJob.refreshing(); + + // 1. drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + assertEquals( + "SUCCESS", + asyncQueryExecutorService + .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) + .getStatus()); + + flintIndexJob.assertState(FlintIndexState.DELETED); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(1); + emrsClient.getJobRunResultCalled(1); + }); + } + + /** + * Index state is stable, Drop Index operation is retryable, expectation is + * + *

(1) call EMR-S (2) change index state to: DELETED + */ + @Test + public void dropIndexWithIndexInActiveState() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + flintIndexJob.active(); + + // 1. drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + assertEquals("SUCCESS", asyncQueryExecutionResponse.getStatus()); + flintIndexJob.assertState(FlintIndexState.DELETED); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(1); + emrsClient.getJobRunResultCalled(1); + }); + } + + /** + * Index state is stable, expectation is + * + *

(1) call EMR-S (2) change index state to: DELETED + */ + @Test + public void dropIndexWithIndexInCreatingState() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + flintIndexJob.creating(); + + // 1. drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + assertEquals( + "SUCCESS", + asyncQueryExecutorService + .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) + .getStatus()); + + flintIndexJob.assertState(FlintIndexState.DELETED); + }); + } + + /** + * Index state is stable, Drop Index operation is retryable, expectation is + * + *

(1) call EMR-S (2) change index state to: DELETED + */ + @Test + public void dropIndexWithIndexInEmptyState() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + return new GetJobRunResult().withJobRun(new JobRun().withState("Cancelled")); + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + + // 1. drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + // 2. fetch result + assertEquals( + "SUCCESS", + asyncQueryExecutorService + .getAsyncQueryResults(response.getQueryId(), asyncQueryRequestContext) + .getStatus()); + + flintIndexJob.assertState(FlintIndexState.DELETED); + }); + } + + /** + * Couldn't acquire lock as the index is in transitioning state. Will result in error. + * + *

(1) not call EMR-S (2) change index state to: DELETED + */ + @Test + public void dropIndexWithIndexInDeletedState() { + ImmutableList.of(SKIPPING, COVERING, MV) + .forEach( + mockDS -> { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + Assert.fail("should not call cancelJobRun"); + return null; + } + + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + Assert.fail("should not call getJobRunResult"); + return null; + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + AsyncQueryExecutorService asyncQueryExecutorService = + createAsyncQueryExecutorService(emrServerlessClientFactory); + + // Mock flint index + mockDS.createIndex(); + // Mock index state + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, mockDS.latestId, MYS3_DATASOURCE); + flintIndexJob.deleting(); + + // 1. drop index + CreateAsyncQueryResponse response = + asyncQueryExecutorService.createAsyncQuery( + new CreateAsyncQueryRequest( + mockDS.query, MYS3_DATASOURCE, LangType.SQL, null), + asyncQueryRequestContext); + + AsyncQueryExecutionResponse asyncQueryExecutionResponse = + asyncQueryExecutorService.getAsyncQueryResults( + response.getQueryId(), asyncQueryRequestContext); + // 2. fetch result + assertEquals("FAILED", asyncQueryExecutionResponse.getStatus()); + assertEquals( + "Transaction failed as flint index is not in a valid state.", + asyncQueryExecutionResponse.getError()); + flintIndexJob.assertState(FlintIndexState.DELETING); + }); + } + + @NotNull + private static EMRServerlessClientFactory getEmrServerlessClientFactory() { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public CancelJobRunResult cancelJobRun( + String applicationId, String jobId, boolean allowExceptionPropagation) { + Assert.fail("should not call cancelJobRun"); + return null; + } + + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + Assert.fail("should not call getJobRunResult"); + return null; + } + }; + EMRServerlessClientFactory emrServerlessClientFactory = (accountId) -> emrsClient; + return emrServerlessClientFactory; + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperBase.java b/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperBase.java new file mode 100644 index 00000000000..66932ce075b --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperBase.java @@ -0,0 +1,55 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.cluster; + +import com.google.common.collect.ImmutableList; +import org.opensearch.sql.datasource.model.DataSourceStatus; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorServiceSpec; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.flint.FlintIndexType; + +import java.util.HashMap; + +public class FlintStreamingJobHouseKeeperBase extends AsyncQueryExecutorServiceSpec { + protected void changeDataSourceStatus(String dataSourceName, DataSourceStatus dataSourceStatus) { + HashMap datasourceMap = new HashMap<>(); + datasourceMap.put("name", dataSourceName); + datasourceMap.put("status", dataSourceStatus); + this.dataSourceService.patchDataSource(datasourceMap); + } + + + protected ImmutableList getMockFlintIndices() { + return ImmutableList.of(getSkipping(), getCovering(), getMv()); + } + + protected MockFlintIndex getMv() { + return new MockFlintIndex( + client, + "flint_my_glue_mydb_mv", + FlintIndexType.MATERIALIZED_VIEW, + "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," + + " incremental_refresh=true, output_mode=\"complete\") "); + } + + protected MockFlintIndex getCovering() { + return new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_covering_index", + FlintIndexType.COVERING, + "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=true, output_mode=\"complete\")"); + } + + protected MockFlintIndex getSkipping() { + return new MockFlintIndex( + client, + "flint_my_glue_mydb_http_logs_skipping_index", + FlintIndexType.SKIPPING, + "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," + + " incremental_refresh=true, output_mode=\"complete\")"); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskBatch1Test.java b/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskBatch1Test.java new file mode 100644 index 00000000000..34b77bf742d --- /dev/null +++ b/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskBatch1Test.java @@ -0,0 +1,293 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.spark.cluster; + +import static org.opensearch.sql.datasource.model.DataSourceStatus.DISABLED; + +import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.amazonaws.services.emrserverless.model.JobRun; +import com.google.common.collect.ImmutableList; +import java.util.HashMap; +import java.util.Map; +import lombok.SneakyThrows; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.ExecutionMode; +import org.opensearch.sql.datasource.model.DataSourceStatus; +import org.opensearch.sql.legacy.metrics.MetricName; +import org.opensearch.sql.legacy.metrics.Metrics; +import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorServiceSpec; +import org.opensearch.sql.spark.asyncquery.model.AsyncQueryRequestContext; +import org.opensearch.sql.spark.asyncquery.model.MockFlintIndex; +import org.opensearch.sql.spark.asyncquery.model.MockFlintSparkJob; +import org.opensearch.sql.spark.dispatcher.model.FlintIndexOptions; +import org.opensearch.sql.spark.flint.FlintIndexMetadata; +import org.opensearch.sql.spark.flint.FlintIndexMetadataService; +import org.opensearch.sql.spark.flint.FlintIndexMetadataServiceImpl; +import org.opensearch.sql.spark.flint.FlintIndexState; +import org.opensearch.sql.spark.flint.FlintIndexType; + +public class FlintStreamingJobHouseKeeperTaskBatch1Test extends FlintStreamingJobHouseKeeperBase { + @Test + public void testStreamingJobHouseKeeperWhenS3GlueIsDisabledButNotStreamingJobQueries() + throws InterruptedException { + changeDataSourceStatus(MYGLUE_DATASOURCE, DISABLED); + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + FlintIndexMetadataService flintIndexMetadataService = new FlintIndexMetadataServiceImpl(client); + FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = + new FlintStreamingJobHouseKeeperTask( + dataSourceService, + flintIndexMetadataService, + getFlintIndexOpFactory((accountId) -> emrsClient)); + + Thread thread = new Thread(flintStreamingJobHouseKeeperTask); + thread.start(); + thread.join(); + + emrsClient.getJobRunResultCalled(0); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + Assertions.assertEquals( + 0L, + Metrics.getInstance() + .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) + .getValue()); + } + + @Test + public void testStreamingJobHouseKeeperWhenFlintIndexIsCorrupted() throws InterruptedException { + String indexName = "flint_my_glue_mydb_http_logs_covering_error_index"; + MockFlintIndex mockFlintIndex = + new MockFlintIndex(client(), indexName, FlintIndexType.COVERING, null); + mockFlintIndex.createIndex(); + changeDataSourceStatus(MYGLUE_DATASOURCE, DISABLED); + LocalEMRSClient emrsClient = getCancelledLocalEmrsClient(); + FlintIndexMetadataService flintIndexMetadataService = new FlintIndexMetadataServiceImpl(client); + FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = + new FlintStreamingJobHouseKeeperTask( + dataSourceService, + flintIndexMetadataService, + getFlintIndexOpFactory((accountId) -> emrsClient)); + + Thread thread = new Thread(flintStreamingJobHouseKeeperTask); + thread.start(); + thread.join(); + + emrsClient.getJobRunResultCalled(0); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + Assertions.assertEquals( + 1L, + Metrics.getInstance() + .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) + .getValue()); + } + + @SneakyThrows + @Test + public void testErrorScenario() { + LocalEMRSClient emrsClient = + new LocalEMRSClient() { + @Override + public GetJobRunResult getJobRunResult(String applicationId, String jobId) { + super.getJobRunResult(applicationId, jobId); + JobRun jobRun = new JobRun(); + jobRun.setState("cancelled"); + return new GetJobRunResult().withJobRun(jobRun); + } + }; + FlintIndexMetadataService flintIndexMetadataService = + new FlintIndexMetadataService() { + @Override + public Map getFlintIndexMetadata( + String indexPattern, AsyncQueryRequestContext asyncQueryRequestContext) { + throw new RuntimeException("Couldn't fetch details from ElasticSearch"); + } + + @Override + public void updateIndexToManualRefresh( + String indexName, + FlintIndexOptions flintIndexOptions, + AsyncQueryRequestContext asyncQueryRequestContext) {} + }; + FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = + new FlintStreamingJobHouseKeeperTask( + dataSourceService, + flintIndexMetadataService, + getFlintIndexOpFactory((accountId) -> emrsClient)); + + Thread thread = new Thread(flintStreamingJobHouseKeeperTask); + thread.start(); + thread.join(); + + Assertions.assertFalse(FlintStreamingJobHouseKeeperTask.isRunning.get()); + emrsClient.getJobRunResultCalled(0); + emrsClient.startJobRunCalled(0); + emrsClient.cancelJobRunCalled(0); + Assertions.assertEquals( + 1L, + Metrics.getInstance() + .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) + .getValue()); + } + + @Test + @SneakyThrows + public void testStreamingJobHouseKeeperMultipleTimesWhenDataSourceDisabled() { + ImmutableList mockFlintIndices = getMockFlintIndices(); + Map indexJobMapping = new HashMap<>(); + mockFlintIndices.forEach( + INDEX -> { + INDEX.createIndex(); + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, INDEX.getLatestId(), MYGLUE_DATASOURCE); + indexJobMapping.put(INDEX, flintIndexJob); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + // Making Index Auto Refresh + INDEX.updateIndexOptions(existingOptions, false); + flintIndexJob.refreshing(); + }); + changeDataSourceStatus(MYGLUE_DATASOURCE, DISABLED); + LocalEMRSClient emrsClient = getCancelledLocalEmrsClient(); + FlintIndexMetadataService flintIndexMetadataService = new FlintIndexMetadataServiceImpl(client); + FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = + new FlintStreamingJobHouseKeeperTask( + dataSourceService, + flintIndexMetadataService, + getFlintIndexOpFactory((accountId) -> emrsClient)); + + Thread thread = new Thread(flintStreamingJobHouseKeeperTask); + thread.start(); + thread.join(); + + mockFlintIndices.forEach( + INDEX -> { + MockFlintSparkJob flintIndexJob = indexJobMapping.get(INDEX); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = INDEX.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + emrsClient.cancelJobRunCalled(3); + emrsClient.getJobRunResultCalled(3); + emrsClient.startJobRunCalled(0); + Assertions.assertEquals( + 0L, + Metrics.getInstance() + .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) + .getValue()); + + // Second Run + Thread thread2 = new Thread(flintStreamingJobHouseKeeperTask); + thread2.start(); + thread2.join(); + mockFlintIndices.forEach( + INDEX -> { + MockFlintSparkJob flintIndexJob = indexJobMapping.get(INDEX); + flintIndexJob.assertState(FlintIndexState.ACTIVE); + Map mappings = INDEX.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("false", options.get("auto_refresh")); + }); + + // No New Calls and Errors + emrsClient.cancelJobRunCalled(3); + emrsClient.getJobRunResultCalled(3); + emrsClient.startJobRunCalled(0); + Assertions.assertEquals( + 0L, + Metrics.getInstance() + .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) + .getValue()); + } + + @SneakyThrows + @Test + public void testRunStreamingJobHouseKeeperWhenDataSourceIsDeleted() { + ImmutableList mockFlintIndices = getMockFlintIndices(); + Map indexJobMapping = new HashMap<>(); + mockFlintIndices.forEach( + INDEX -> { + INDEX.createIndex(); + MockFlintSparkJob flintIndexJob = + new MockFlintSparkJob( + flintIndexStateModelService, INDEX.getLatestId(), MYGLUE_DATASOURCE); + indexJobMapping.put(INDEX, flintIndexJob); + HashMap existingOptions = new HashMap<>(); + existingOptions.put("auto_refresh", "true"); + // Making Index Auto Refresh + INDEX.updateIndexOptions(existingOptions, false); + flintIndexJob.refreshing(); + }); + this.dataSourceService.deleteDataSource(MYGLUE_DATASOURCE); + LocalEMRSClient emrsClient = getCancelledLocalEmrsClient(); + FlintIndexMetadataService flintIndexMetadataService = new FlintIndexMetadataServiceImpl(client); + FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = + new FlintStreamingJobHouseKeeperTask( + dataSourceService, + flintIndexMetadataService, + getFlintIndexOpFactory((accountId) -> emrsClient)); + + Thread thread = new Thread(flintStreamingJobHouseKeeperTask); + thread.start(); + thread.join(); + + mockFlintIndices.forEach( + INDEX -> { + MockFlintSparkJob flintIndexJob = indexJobMapping.get(INDEX); + flintIndexJob.assertState(FlintIndexState.DELETED); + Map mappings = INDEX.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("true", options.get("auto_refresh")); + }); + emrsClient.cancelJobRunCalled(3); + emrsClient.getJobRunResultCalled(3); + emrsClient.startJobRunCalled(0); + Assertions.assertEquals( + 0L, + Metrics.getInstance() + .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) + .getValue()); + + // Second Run + Thread thread2 = new Thread(flintStreamingJobHouseKeeperTask); + thread2.start(); + thread2.join(); + mockFlintIndices.forEach( + INDEX -> { + MockFlintSparkJob flintIndexJob = indexJobMapping.get(INDEX); + flintIndexJob.assertState(FlintIndexState.DELETED); + Map mappings = INDEX.getIndexMappings(); + Map meta = (HashMap) mappings.get("_meta"); + Map options = (Map) meta.get("options"); + Assertions.assertEquals("true", options.get("auto_refresh")); + }); + // No New Calls and Errors + emrsClient.cancelJobRunCalled(3); + emrsClient.getJobRunResultCalled(3); + emrsClient.startJobRunCalled(0); + Assertions.assertEquals( + 0L, + Metrics.getInstance() + .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) + .getValue()); + } +} diff --git a/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskTest.java b/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskBatch2Test.java similarity index 50% rename from async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskTest.java rename to async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskBatch2Test.java index 0a3a1809321..22e5c0be9cd 100644 --- a/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskTest.java +++ b/async-query/src/test/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTaskBatch2Test.java @@ -5,14 +5,9 @@ package org.opensearch.sql.spark.cluster; -import static org.opensearch.sql.datasource.model.DataSourceStatus.DISABLED; -import static org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorServiceSpec.MYGLUE_DATASOURCE; - import com.amazonaws.services.emrserverless.model.GetJobRunResult; import com.amazonaws.services.emrserverless.model.JobRun; import com.google.common.collect.ImmutableList; -import java.util.HashMap; -import java.util.Map; import lombok.SneakyThrows; import org.junit.Test; import org.junit.jupiter.api.Assertions; @@ -30,8 +25,12 @@ import org.opensearch.sql.spark.flint.FlintIndexState; import org.opensearch.sql.spark.flint.FlintIndexType; -public class FlintStreamingJobHouseKeeperTaskTest extends AsyncQueryExecutorServiceSpec { +import java.util.HashMap; +import java.util.Map; + +import static org.opensearch.sql.datasource.model.DataSourceStatus.DISABLED; +public class FlintStreamingJobHouseKeeperTaskBatch2Test extends FlintStreamingJobHouseKeeperBase { @Test @SneakyThrows public void testStreamingJobHouseKeeperWhenDataSourceDisabled() { @@ -82,37 +81,6 @@ public void testStreamingJobHouseKeeperWhenDataSourceDisabled() { .getValue()); } - private ImmutableList getMockFlintIndices() { - return ImmutableList.of(getSkipping(), getCovering(), getMv()); - } - - private MockFlintIndex getMv() { - return new MockFlintIndex( - client, - "flint_my_glue_mydb_mv", - FlintIndexType.MATERIALIZED_VIEW, - "ALTER MATERIALIZED VIEW my_glue.mydb.mv WITH (auto_refresh=false," - + " incremental_refresh=true, output_mode=\"complete\") "); - } - - private MockFlintIndex getCovering() { - return new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_covering_index", - FlintIndexType.COVERING, - "ALTER INDEX covering ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=true, output_mode=\"complete\")"); - } - - private MockFlintIndex getSkipping() { - return new MockFlintIndex( - client, - "flint_my_glue_mydb_http_logs_skipping_index", - FlintIndexType.SKIPPING, - "ALTER SKIPPING INDEX ON my_glue.mydb.http_logs WITH (auto_refresh=false," - + " incremental_refresh=true, output_mode=\"complete\")"); - } - @Test @SneakyThrows public void testStreamingJobHouseKeeperWhenCancelJobGivesTimeout() { @@ -313,270 +281,4 @@ public void testStreamingJobHouseKeeperWhenDataSourceIsNeitherDisabledNorDeleted .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) .getValue()); } - - @Test - public void testStreamingJobHouseKeeperWhenS3GlueIsDisabledButNotStreamingJobQueries() - throws InterruptedException { - changeDataSourceStatus(MYGLUE_DATASOURCE, DISABLED); - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - FlintIndexMetadataService flintIndexMetadataService = new FlintIndexMetadataServiceImpl(client); - FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = - new FlintStreamingJobHouseKeeperTask( - dataSourceService, - flintIndexMetadataService, - getFlintIndexOpFactory((accountId) -> emrsClient)); - - Thread thread = new Thread(flintStreamingJobHouseKeeperTask); - thread.start(); - thread.join(); - - emrsClient.getJobRunResultCalled(0); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - Assertions.assertEquals( - 0L, - Metrics.getInstance() - .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) - .getValue()); - } - - @Test - public void testStreamingJobHouseKeeperWhenFlintIndexIsCorrupted() throws InterruptedException { - String indexName = "flint_my_glue_mydb_http_logs_covering_error_index"; - MockFlintIndex mockFlintIndex = - new MockFlintIndex(client(), indexName, FlintIndexType.COVERING, null); - mockFlintIndex.createIndex(); - changeDataSourceStatus(MYGLUE_DATASOURCE, DISABLED); - LocalEMRSClient emrsClient = getCancelledLocalEmrsClient(); - FlintIndexMetadataService flintIndexMetadataService = new FlintIndexMetadataServiceImpl(client); - FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = - new FlintStreamingJobHouseKeeperTask( - dataSourceService, - flintIndexMetadataService, - getFlintIndexOpFactory((accountId) -> emrsClient)); - - Thread thread = new Thread(flintStreamingJobHouseKeeperTask); - thread.start(); - thread.join(); - - emrsClient.getJobRunResultCalled(0); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - Assertions.assertEquals( - 1L, - Metrics.getInstance() - .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) - .getValue()); - } - - @SneakyThrows - @Test - public void testErrorScenario() { - LocalEMRSClient emrsClient = - new LocalEMRSClient() { - @Override - public GetJobRunResult getJobRunResult(String applicationId, String jobId) { - super.getJobRunResult(applicationId, jobId); - JobRun jobRun = new JobRun(); - jobRun.setState("cancelled"); - return new GetJobRunResult().withJobRun(jobRun); - } - }; - FlintIndexMetadataService flintIndexMetadataService = - new FlintIndexMetadataService() { - @Override - public Map getFlintIndexMetadata( - String indexPattern, AsyncQueryRequestContext asyncQueryRequestContext) { - throw new RuntimeException("Couldn't fetch details from ElasticSearch"); - } - - @Override - public void updateIndexToManualRefresh( - String indexName, - FlintIndexOptions flintIndexOptions, - AsyncQueryRequestContext asyncQueryRequestContext) {} - }; - FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = - new FlintStreamingJobHouseKeeperTask( - dataSourceService, - flintIndexMetadataService, - getFlintIndexOpFactory((accountId) -> emrsClient)); - - Thread thread = new Thread(flintStreamingJobHouseKeeperTask); - thread.start(); - thread.join(); - - Assertions.assertFalse(FlintStreamingJobHouseKeeperTask.isRunning.get()); - emrsClient.getJobRunResultCalled(0); - emrsClient.startJobRunCalled(0); - emrsClient.cancelJobRunCalled(0); - Assertions.assertEquals( - 1L, - Metrics.getInstance() - .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) - .getValue()); - } - - @Test - @SneakyThrows - public void testStreamingJobHouseKeeperMultipleTimesWhenDataSourceDisabled() { - ImmutableList mockFlintIndices = getMockFlintIndices(); - Map indexJobMapping = new HashMap<>(); - mockFlintIndices.forEach( - INDEX -> { - INDEX.createIndex(); - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, INDEX.getLatestId(), MYGLUE_DATASOURCE); - indexJobMapping.put(INDEX, flintIndexJob); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - // Making Index Auto Refresh - INDEX.updateIndexOptions(existingOptions, false); - flintIndexJob.refreshing(); - }); - changeDataSourceStatus(MYGLUE_DATASOURCE, DISABLED); - LocalEMRSClient emrsClient = getCancelledLocalEmrsClient(); - FlintIndexMetadataService flintIndexMetadataService = new FlintIndexMetadataServiceImpl(client); - FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = - new FlintStreamingJobHouseKeeperTask( - dataSourceService, - flintIndexMetadataService, - getFlintIndexOpFactory((accountId) -> emrsClient)); - - Thread thread = new Thread(flintStreamingJobHouseKeeperTask); - thread.start(); - thread.join(); - - mockFlintIndices.forEach( - INDEX -> { - MockFlintSparkJob flintIndexJob = indexJobMapping.get(INDEX); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = INDEX.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - emrsClient.cancelJobRunCalled(3); - emrsClient.getJobRunResultCalled(3); - emrsClient.startJobRunCalled(0); - Assertions.assertEquals( - 0L, - Metrics.getInstance() - .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) - .getValue()); - - // Second Run - Thread thread2 = new Thread(flintStreamingJobHouseKeeperTask); - thread2.start(); - thread2.join(); - mockFlintIndices.forEach( - INDEX -> { - MockFlintSparkJob flintIndexJob = indexJobMapping.get(INDEX); - flintIndexJob.assertState(FlintIndexState.ACTIVE); - Map mappings = INDEX.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("false", options.get("auto_refresh")); - }); - - // No New Calls and Errors - emrsClient.cancelJobRunCalled(3); - emrsClient.getJobRunResultCalled(3); - emrsClient.startJobRunCalled(0); - Assertions.assertEquals( - 0L, - Metrics.getInstance() - .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) - .getValue()); - } - - @SneakyThrows - @Test - public void testRunStreamingJobHouseKeeperWhenDataSourceIsDeleted() { - ImmutableList mockFlintIndices = getMockFlintIndices(); - Map indexJobMapping = new HashMap<>(); - mockFlintIndices.forEach( - INDEX -> { - INDEX.createIndex(); - MockFlintSparkJob flintIndexJob = - new MockFlintSparkJob( - flintIndexStateModelService, INDEX.getLatestId(), MYGLUE_DATASOURCE); - indexJobMapping.put(INDEX, flintIndexJob); - HashMap existingOptions = new HashMap<>(); - existingOptions.put("auto_refresh", "true"); - // Making Index Auto Refresh - INDEX.updateIndexOptions(existingOptions, false); - flintIndexJob.refreshing(); - }); - this.dataSourceService.deleteDataSource(MYGLUE_DATASOURCE); - LocalEMRSClient emrsClient = getCancelledLocalEmrsClient(); - FlintIndexMetadataService flintIndexMetadataService = new FlintIndexMetadataServiceImpl(client); - FlintStreamingJobHouseKeeperTask flintStreamingJobHouseKeeperTask = - new FlintStreamingJobHouseKeeperTask( - dataSourceService, - flintIndexMetadataService, - getFlintIndexOpFactory((accountId) -> emrsClient)); - - Thread thread = new Thread(flintStreamingJobHouseKeeperTask); - thread.start(); - thread.join(); - - mockFlintIndices.forEach( - INDEX -> { - MockFlintSparkJob flintIndexJob = indexJobMapping.get(INDEX); - flintIndexJob.assertState(FlintIndexState.DELETED); - Map mappings = INDEX.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("true", options.get("auto_refresh")); - }); - emrsClient.cancelJobRunCalled(3); - emrsClient.getJobRunResultCalled(3); - emrsClient.startJobRunCalled(0); - Assertions.assertEquals( - 0L, - Metrics.getInstance() - .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) - .getValue()); - - // Second Run - Thread thread2 = new Thread(flintStreamingJobHouseKeeperTask); - thread2.start(); - thread2.join(); - mockFlintIndices.forEach( - INDEX -> { - MockFlintSparkJob flintIndexJob = indexJobMapping.get(INDEX); - flintIndexJob.assertState(FlintIndexState.DELETED); - Map mappings = INDEX.getIndexMappings(); - Map meta = (HashMap) mappings.get("_meta"); - Map options = (Map) meta.get("options"); - Assertions.assertEquals("true", options.get("auto_refresh")); - }); - // No New Calls and Errors - emrsClient.cancelJobRunCalled(3); - emrsClient.getJobRunResultCalled(3); - emrsClient.startJobRunCalled(0); - Assertions.assertEquals( - 0L, - Metrics.getInstance() - .getNumericalMetric(MetricName.STREAMING_JOB_HOUSEKEEPER_TASK_FAILURE_COUNT) - .getValue()); - } - - private void changeDataSourceStatus(String dataSourceName, DataSourceStatus dataSourceStatus) { - HashMap datasourceMap = new HashMap<>(); - datasourceMap.put("name", dataSourceName); - datasourceMap.put("status", dataSourceStatus); - this.dataSourceService.patchDataSource(datasourceMap); - } } diff --git a/core/build.gradle b/core/build.gradle index ecd91040995..263a8e73ddb 100644 --- a/core/build.gradle +++ b/core/build.gradle @@ -66,6 +66,7 @@ dependencies { } test { + maxParallelForks = Runtime.runtime.availableProcessors() useJUnitPlatform() testLogging { events "skipped", "failed" diff --git a/datasources/build.gradle b/datasources/build.gradle index c6915ee4f33..575db4d499f 100644 --- a/datasources/build.gradle +++ b/datasources/build.gradle @@ -39,6 +39,7 @@ dependencies { } test { + maxParallelForks = Runtime.runtime.availableProcessors() useJUnitPlatform() testLogging { events "skipped", "failed" diff --git a/doctest/build.gradle b/doctest/build.gradle index 09c9421409c..6680411c752 100644 --- a/doctest/build.gradle +++ b/doctest/build.gradle @@ -65,6 +65,8 @@ task startPrometheus(type: SpawnProcessTask) { //evaluationDependsOn(':') task startOpenSearch(type: SpawnProcessTask) { + dependsOn ':opensearch-sql-plugin:bundlePlugin' + if( getOSFamilyType() == "windows") { command "${path}\\gradlew.bat -p ${plugin_path} runRestTestCluster" } diff --git a/doctest/test_docs.py b/doctest/test_docs.py index d79e5b1cf55..636902f6390 100644 --- a/doctest/test_docs.py +++ b/doctest/test_docs.py @@ -11,6 +11,7 @@ import subprocess import sys import unittest +from concurrent.futures import ThreadPoolExecutor, as_completed from functools import partial import click @@ -205,6 +206,7 @@ class TestDataManager: def __init__(self): self.client = OpenSearch([ENDPOINT], verify_certs=True) + self.is_loaded = False def load_file(self, filename, index_name): mapping_file_path = './test_mapping/' + filename @@ -218,18 +220,33 @@ def load_json(): for line in f: yield json.loads(line) - helpers.bulk(self.client, load_json(), stats_only=True, index=index_name, refresh='wait_for') + helpers.bulk(self.client, load_json(), stats_only=True, index=index_name, refresh="wait_for") def load_all_test_data(self): - for index_name, filename in TEST_DATA.items(): + if self.is_loaded: + return + + def load_index(index_name, filename): if filename is not None: self.load_file(filename, index_name) else: debug(f"Skipping index '{index_name}' - filename is None") - def cleanup_indices(self): - indices_to_delete = list(TEST_DATA.keys()) - self.client.indices.delete(index=indices_to_delete, ignore_unavailable=True) + with ThreadPoolExecutor() as executor: + futures = { + executor.submit(load_index, index_name, filename): index_name + for index_name, filename in TEST_DATA.items() + } + + for future in as_completed(futures): + index_name = futures[future] + try: + future.result() + except Exception as e: + debug(f"Error loading index '{index_name}': {str(e)}") + raise + + self.is_loaded = True def sql_cli_transform(s): @@ -282,7 +299,7 @@ def set_up_test_indices_without_calcite(test): def tear_down(test): - get_test_data_manager().cleanup_indices() + pass docsuite = partial(doctest.DocFileSuite, diff --git a/gradle.properties b/gradle.properties index 5cf428415a8..cca553e80bf 100644 --- a/gradle.properties +++ b/gradle.properties @@ -4,3 +4,4 @@ version=1.13.0 org.gradle.jvmargs=-Duser.language=en -Duser.country=US +org.gradle.parallel=true diff --git a/legacy/build.gradle b/legacy/build.gradle index 72829e1be79..0359b5c9dd9 100644 --- a/legacy/build.gradle +++ b/legacy/build.gradle @@ -66,6 +66,7 @@ compileTestJava { // TODO: Need to update integration test to use OpenSearch test framework test { + maxParallelForks = Runtime.runtime.availableProcessors() include '**/*Test.class' exclude 'org/opensearch/sql/intgtest/**' // Gradle runs unit test using a working directory other and project root diff --git a/opensearch/build.gradle b/opensearch/build.gradle index 918c8b9853d..f86cd706f5b 100644 --- a/opensearch/build.gradle +++ b/opensearch/build.gradle @@ -57,6 +57,7 @@ dependencies { } test { + maxParallelForks = Runtime.runtime.availableProcessors() useJUnitPlatform() testLogging { events "passed", "skipped", "failed" diff --git a/plugin/build.gradle b/plugin/build.gradle index 54573dcc121..86f2d2e8b80 100644 --- a/plugin/build.gradle +++ b/plugin/build.gradle @@ -155,6 +155,7 @@ dependencies { } test { + maxParallelForks = Runtime.runtime.availableProcessors() include '**/*Test.class' testLogging { events "passed", "skipped", "failed" diff --git a/ppl/build.gradle b/ppl/build.gradle index 3237c394470..6e0a77ef811 100644 --- a/ppl/build.gradle +++ b/ppl/build.gradle @@ -86,6 +86,7 @@ spotless { } test { + maxParallelForks = Runtime.runtime.availableProcessors() testLogging { events "passed", "skipped", "failed" exceptionFormat "full" diff --git a/prometheus/build.gradle b/prometheus/build.gradle index 7a3b3f7af6e..4d39c8076f3 100644 --- a/prometheus/build.gradle +++ b/prometheus/build.gradle @@ -32,6 +32,7 @@ dependencies { } test { + maxParallelForks = Runtime.runtime.availableProcessors() useJUnitPlatform() testLogging { events "passed", "skipped", "failed" diff --git a/protocol/build.gradle b/protocol/build.gradle index 3da6d56770c..c9de41284bd 100644 --- a/protocol/build.gradle +++ b/protocol/build.gradle @@ -49,6 +49,7 @@ configurations.all { } test { + maxParallelForks = Runtime.runtime.availableProcessors() useJUnitPlatform() testLogging { events "passed", "skipped", "failed" diff --git a/sql/build.gradle b/sql/build.gradle index d7509db0994..fea44918977 100644 --- a/sql/build.gradle +++ b/sql/build.gradle @@ -60,6 +60,7 @@ dependencies { } test { + maxParallelForks = Runtime.runtime.availableProcessors() useJUnitPlatform() testLogging { events "passed", "skipped", "failed"