diff --git a/.github/workflows/pinot_multi_stage_query_engine_compatibility_tests.yml b/.github/workflows/pinot_multi_stage_query_engine_compatibility_tests.yml new file mode 100644 index 000000000000..52b7773033de --- /dev/null +++ b/.github/workflows/pinot_multi_stage_query_engine_compatibility_tests.yml @@ -0,0 +1,84 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +name: Pinot Multi-Stage Query Engine Compatibility Regression Test + +on: + workflow_dispatch: + inputs: + oldCommit: + description: "Git hash (or tag) for old commit. (required)" + required: true + newCommit: + description: "Git hash (or tag) for new commit. (required)" + required: true + +jobs: + compatibility-verifier: + runs-on: ubuntu-latest + strategy: + matrix: + test_suite: [ "compatibility-verifier/multi-stage-query-engine-test-suite" ] + name: Pinot Multi-Stage Query Engine Compatibility Regression Testing against ${{ github.event.inputs.oldCommit }} and ${{ github.event.inputs.newCommit }} on ${{ matrix.test_suite }} + steps: + - uses: actions/checkout@v4 + - name: Set up JDK 11 + uses: actions/setup-java@v4 + with: + java-version: 11 + distribution: 'temurin' + cache: 'maven' + - name: Setup node + uses: actions/setup-node@v4 + with: + node-version: v16.15.0 + cache: 'npm' + cache-dependency-path: pinot-controller/src/main/resources/package-lock.json + - name: Install npm + run: | + npm install -g npm@8.5.5 + npm --version + - name: Pinot Multi-Stage Query Engine Compatibility Regression Testing + if : ${{github.event_name == 'workflow_dispatch'}} + env: + OLD_COMMIT: ${{ github.event.inputs.oldCommit }} + NEW_COMMIT: ${{ github.event.inputs.newCommit }} + WORKING_DIR: /tmp/multi-stage-compatibility-verifier + TEST_SUITE: ${{ matrix.test_suite }} + MAVEN_OPTS: > + -Xmx2G -DskipShade -DfailIfNoTests=false -Dmaven.wagon.httpconnectionManager.ttlSeconds=25 + -Dmaven.wagon.http.retryHandler.count=30 -Dhttp.keepAlive=false -Dmaven.wagon.http.pool=false + -XX:+IgnoreUnrecognizedVMOptions + --add-exports=jdk.compiler/com.sun.tools.javac.api=ALL-UNNAMED + --add-exports=jdk.compiler/com.sun.tools.javac.file=ALL-UNNAMED + --add-exports=jdk.compiler/com.sun.tools.javac.parser=ALL-UNNAMED + --add-exports=jdk.compiler/com.sun.tools.javac.tree=ALL-UNNAMED + --add-exports=jdk.compiler/com.sun.tools.javac.util=ALL-UNNAMED + run: .github/workflows/scripts/.pinot_compatibility_verifier.sh + - name: Archive artifacts into zip + if: always() + run: | + zip -1 -r artifacts.zip /tmp/multi-stage-compatibility-verifier/* + - uses: actions/upload-artifact@v4 + name: Store multi-stage compatibility verifier work directory + if: always() + with: + ## TODO: currently matrix.test_suite cannot be used as part of name due to invalid path character. + name: multi_stage_compatibility_verifier_work_dir + retention-days: 3 + path: artifacts.zip diff --git a/.github/workflows/pinot_tests.yml b/.github/workflows/pinot_tests.yml index ce170f31900a..ef2828d549b9 100644 --- a/.github/workflows/pinot_tests.yml +++ b/.github/workflows/pinot_tests.yml @@ -267,7 +267,7 @@ jobs: matrix: test_suite: [ "compatibility-verifier/sample-test-suite" ] old_commit: [ - "release-0.12.1", "release-1.0.0", "release-1.1.0", "master" + "release-1.0.0", "release-1.1.0", "master" ] name: Pinot Compatibility Regression Testing against ${{ matrix.old_commit }} on ${{ matrix.test_suite }} steps: @@ -314,6 +314,62 @@ jobs: --add-exports=jdk.compiler/com.sun.tools.javac.util=ALL-UNNAMED run: .github/workflows/scripts/.pinot_compatibility_verifier.sh + multi-stage-compatibility-verifier: + if: github.repository == 'apache/pinot' + runs-on: ubuntu-latest + strategy: + # Changed to false in order to improve coverage using unsafe buffers + fail-fast: false + matrix: + test_suite: [ "compatibility-verifier/multi-stage-query-engine-test-suite" ] + old_commit: [ + "master" + ] + name: Pinot Multi-Stage Query Engine Compatibility Regression Testing against ${{ matrix.old_commit }} on ${{ matrix.test_suite }} + steps: + - uses: actions/checkout@v4 + - name: Set up JDK 11 + uses: actions/setup-java@v4 + with: + java-version: 11 + distribution: 'temurin' + cache: 'maven' + - name: Setup node + uses: actions/setup-node@v4 + with: + node-version: v16.15.0 + cache: 'npm' + cache-dependency-path: pinot-controller/src/main/resources/package-lock.json + - name: Install npm + run: | + npm install -g npm@8.5.5 + npm --version + # Step that does that actual cache save and restore + - uses: actions/cache@v4 + env: + SEGMENT_DOWNLOAD_TIMEOUT_MINS: 10 + with: + path: ~/.m2/repository + key: ${{ runner.os }}-maven-${{ hashFiles('**/pom.xml') }} + restore-keys: | + ${{ runner.os }}-maven- + - name: Pinot Multi-Stage Query Engine Compatibility Regression Testing + env: + OLD_COMMIT: ${{ matrix.old_commit }} + WORKING_DIR: /tmp/multi-stage-compatibility-verifier + TEST_SUITE: ${{ matrix.test_suite }} + MAVEN_OPTS: > + -Xmx2G -DskipShade -DfailIfNoTests=false -Dmaven.wagon.httpconnectionManager.ttlSeconds=25 + -Dmaven.wagon.http.retryHandler.count=30 -Dhttp.keepAlive=false -Dmaven.wagon.http.pool=false + -B -ntp + -XX:+IgnoreUnrecognizedVMOptions + --add-exports=jdk.compiler/com.sun.tools.javac.api=ALL-UNNAMED + --add-exports=jdk.compiler/com.sun.tools.javac.file=ALL-UNNAMED + --add-exports=jdk.compiler/com.sun.tools.javac.parser=ALL-UNNAMED + --add-exports=jdk.compiler/com.sun.tools.javac.tree=ALL-UNNAMED + --add-exports=jdk.compiler/com.sun.tools.javac.util=ALL-UNNAMED + run: .github/workflows/scripts/.pinot_compatibility_verifier.sh + quickstarts: if: github.repository == 'apache/pinot' runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index 25c9ab85411d..c3b52d6e047a 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ cscope.* .classpath .project .svn +.java-version .externalToolBuilders/ maven-eclipse.xml target/ diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/BrokerConfig.properties b/compatibility-verifier/multi-stage-query-engine-test-suite/config/BrokerConfig.properties new file mode 100644 index 000000000000..e56f23c287b6 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/BrokerConfig.properties @@ -0,0 +1,23 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +pinot.broker.client.queryPort = 8099 +pinot.zk.server = localhost:2181 +pinot.cluster.name = PinotCluster +pinot.broker.disable.query.groovy=false diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/ControllerConfig.properties b/compatibility-verifier/multi-stage-query-engine-test-suite/config/ControllerConfig.properties new file mode 100644 index 000000000000..9949b2519ea6 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/ControllerConfig.properties @@ -0,0 +1,25 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +controller.host = localhost +controller.port = 9000 +controller.zk.str = localhost:2181 +controller.data.dir = /tmp/PinotController +controller.helix.cluster.name = PinotCluster +controller.disable.ingestion.groovy = false diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/FeatureTest1-schema.json b/compatibility-verifier/multi-stage-query-engine-test-suite/config/FeatureTest1-schema.json new file mode 100644 index 000000000000..85378dd3a04d --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/FeatureTest1-schema.json @@ -0,0 +1,94 @@ +{ + "dimensionFieldSpecs": [ + { + "dataType": "INT", + "name": "generationNumber" + }, + { + "dataType": "STRING", + "name": "stringDimSV1" + }, + { + "dataType": "STRING", + "name": "stringDimSV2" + }, + { + "dataType": "LONG", + "name": "longDimSV1" + }, + { + "dataType": "LONG", + "name": "longDimSV2" + }, + { + "dataType": "STRING", + "name": "stringDimMV1", + "singleValueField": false + }, + { + "dataType": "STRING", + "name": "stringDimMV2", + "singleValueField": false + }, + { + "dataType": "INT", + "name": "intDimMV1", + "singleValueField": false + }, + { + "dataType": "INT", + "name": "intDimMV2", + "singleValueField": false + }, + { + "dataType": "STRING", + "maxLength": 1000, + "name": "textDim1" + }, + { + "dataType": "BYTES", + "name": "bytesDimSV1" + }, + { + "dataType": "STRING", + "name": "mapDim1__KEYS", + "singleValueField": false + }, + { + "dataType": "INT", + "name": "mapDim1__VALUES", + "singleValueField": false + }, + { + "dataType": "STRING", + "name": "mapDim2json" + } + ], + "metricFieldSpecs": [ + { + "dataType": "INT", + "name": "intMetric1" + }, + { + "dataType": "LONG", + "name": "longMetric1" + }, + { + "dataType": "FLOAT", + "name": "floatMetric1" + }, + { + "dataType": "DOUBLE", + "name": "doubleMetric1" + } + ], + "dateTimeFieldSpecs" : [ + { + "name" : "HoursSinceEpoch", + "dataType" : "INT", + "format" : "1:HOURS:EPOCH", + "granularity": "1:HOURS" + } + ], + "schemaName": "FeatureTest1" +} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/FeatureTest2-schema.json b/compatibility-verifier/multi-stage-query-engine-test-suite/config/FeatureTest2-schema.json new file mode 100644 index 000000000000..f53c5a2c863f --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/FeatureTest2-schema.json @@ -0,0 +1,94 @@ +{ + "dimensionFieldSpecs": [ + { + "dataType": "INT", + "name": "generationNumber" + }, + { + "dataType": "STRING", + "name": "stringDimSV1" + }, + { + "dataType": "STRING", + "name": "stringDimSV2" + }, + { + "dataType": "LONG", + "name": "longDimSV1" + }, + { + "dataType": "LONG", + "name": "longDimSV2" + }, + { + "dataType": "STRING", + "name": "stringDimMV1", + "singleValueField": false + }, + { + "dataType": "STRING", + "name": "stringDimMV2", + "singleValueField": false + }, + { + "dataType": "INT", + "name": "intDimMV1", + "singleValueField": false + }, + { + "dataType": "INT", + "name": "intDimMV2", + "singleValueField": false + }, + { + "dataType": "STRING", + "maxLength": 1000, + "name": "textDim1" + }, + { + "dataType": "BYTES", + "name": "bytesDimSV1" + }, + { + "dataType": "STRING", + "name": "mapDim1__KEYS", + "singleValueField": false + }, + { + "dataType": "INT", + "name": "mapDim1__VALUES", + "singleValueField": false + }, + { + "dataType": "STRING", + "name": "mapDim2json" + } + ], + "metricFieldSpecs": [ + { + "dataType": "INT", + "name": "intMetric1" + }, + { + "dataType": "LONG", + "name": "longMetric1" + }, + { + "dataType": "FLOAT", + "name": "floatMetric1" + }, + { + "dataType": "DOUBLE", + "name": "doubleMetric1" + } + ], + "dateTimeFieldSpecs" : [ + { + "name" : "HoursSinceEpoch", + "dataType" : "INT", + "format" : "1:HOURS:EPOCH", + "granularity": "1:HOURS" + } + ], + "schemaName": "FeatureTest2" +} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/ServerConfig.properties b/compatibility-verifier/multi-stage-query-engine-test-suite/config/ServerConfig.properties new file mode 100644 index 000000000000..0bf3a9b47bb8 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/ServerConfig.properties @@ -0,0 +1,25 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +pinot.server.adminapi.port = 8097 +pinot.server.netty.port = 8098 +pinot.zk.server = localhost:2181 +pinot.cluster.name = PinotCluster +pinot.server.instance.dataDir = /tmp/PinotServer/data +pinot.server.instance.segmentTarDir = /tmp/PinotServer/segments diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/FeatureTest1-data-00.csv b/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/FeatureTest1-data-00.csv new file mode 100644 index 000000000000..bb537d9ba333 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/FeatureTest1-data-00.csv @@ -0,0 +1,12 @@ +# HoursSinceEpoch generationNumber stringDimSV1 stringDimSV2 longDimSV1 longDimSV2 stringDimMV1 stringDimMV2 intDimMV1 intDimMV2 textDim1 mapDim1__KEYS mapDim1__VALUES mapDim2json intMetric1 longMetric1 floatMetric1 doubleMetric1 +# Add some common rows from first segment, and some new rows as well +123456,__GENERATION_NUMBER__,"s1-0",s2-0,1,2,m1-0-0;m1-0-1,m2-0-0;m2-0-1,3;4,6;7,Java C++ Python,01a0bc,k1;k2;k3;k4;k5,1;1;2;2;2,"{""k1"":1,""k2"":1,""k3"":2,""k4"":2,""k5"":2}",10,11,12.1,13.1 +123456,__GENERATION_NUMBER__,"s1-0",s2-0,1,2,m1-0-0;m1-0-1,m2-0-0;m2-0-1,3;4,6;7,Java C++ Python,4877625602,k1;k2;k3;k4;k5,3;3;3;3;3,"{""k1"":3,""k2"":3,""k3"":3,""k4"":3,""k5"":3}",10,11,12.1,13.1, # Dupliate of row 0 1 +123456,__GENERATION_NUMBER__,s1-2,s2-2,11,21,m1-2-0;m1-2-1,m2-2-0;m2-2-1,32;42,62;72,Java C++ golang,13225573e3f5,k1;k2;k3;k4;k5,4;5;6;7;7,"{""k1"":4,""k2"":5,""k3"":6,""k4"":7,""k5"":7}",10,21,22.1,23.10 +123456,__GENERATION_NUMBER__,s1-2,s2-2,11,21,m1-3-0;m1-3-1,m2-3-0;m2-3-1,32;42,62;72,Java C++ golang,deadbeef,k1;k2;k3;k4;k5,7;7;7;7;7,"{""k1"":7,""k2"":7,""k3"":7,""k4"":7,""k5"":7}",10,21,22.1,23.10, # All sv cols same as prev +123456,__GENERATION_NUMBER__,s1-4,s2-4,41,22,m1-2-0;m1-2-1,m2-2-0;m2-2-1,42;52,72;82,Java C++ golang,deed0507,k1;k2;k3;k4;k5,7;7;8;8;8,"{""k1"":7,""k2"":7,""k3"":8,""k4"":8,""k5"":8}",14,24,24.1,24.10, # All mv cols same as row 2 +123456,__GENERATION_NUMBER__,s1-5,,,32,m1-5-0,m2-2-0,,92;22,golang shell bash,,k1;k2;k3;k4;k5,7;7;7;7;7,"{""k1"":7,""k2"":7,""k3"":7,""k4"":7,""k5"":7}",,24,,24.10, # Default values for some columns +123456,__GENERATION_NUMBER__,s1-6,s2-6,7611,7621,m1-5-0;m1-5-1;m1-5-2,m2-2-0,392;462,6662;782,C++ golang python,deed0507,k1;k2;k3;k4;k5,7;8;9;10;20,"{""k1"":7,""k2"":8,""k3"":9,""k4"":10,""k5"":20}",101,251,262.1,263.10, # 3 values in MV +123456,__GENERATION_NUMBER__,s1-6,s2-6,7611,7621,m1-5-0;m1-5-1;m1-5-2,m2-2-0,392;462,6662;782,C++ golang python,deed0507,k1;k2;k3;k4;k5,7;8;9;10;20,"{""k1"":7,""k2"":8,""k3"":9,""k4"":10,""k5"":20}",2147483647,251,262.1,263.10, # MAX_INT in int metric +123456,__GENERATION_NUMBER__,s1-6,s2-6,7611,7621,m1-5-0;m1-5-1;m1-5-2,m2-2-0,392;462,6662;782,C++ golang python,deed0507,k1;k2;k3;k4;k5,7;8;9;10;20,"{""k1"":7,""k2"":8,""k3"":9,""k4"":10,""k5"":20}",2147483647,251,262.1,263.10, # MAX_INT in int metric +123456,__GENERATION_NUMBER__,s1-7,s2-7,6766,6777,m1-6-0;m1-6-1;m1-6-2;m1-6-3,m2-6-0;m2-6-1,392;462,6662;782,golang Java,d54d0507,k1;k2;k3;k4;k5,31;31;32;32;32,"{""k1"":31,""k2"":31,""k3"":32,""k4"":32,""k5"":32}",87,251,262.10,263.10 diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/FeatureTest2-data-realtime-00.csv b/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/FeatureTest2-data-realtime-00.csv new file mode 100644 index 000000000000..e4e36e47051f --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/FeatureTest2-data-realtime-00.csv @@ -0,0 +1,11 @@ +# HoursSinceEpoch generationNumber stringDimSV1 stringDimSV2 longDimSV1 longDimSV2 stringDimMV1 stringDimMV2 intDimMV1 intDimMV2 textDim1 mapDim1__KEYS mapDim1__VALUES mapDim2json intMetric1 longMetric1 floatMetric1 doubleMetric1 +123456,__GENERATION_NUMBER__,"s1-0",s2-0,1,2,m1-0-0;m1-0-1,m2-0-0;m2-0-1,3;4,6;7,Java C++ Python,01a0bc,k1;k2;k3;k4;k5,1;1;2;2;2,"{""k1"":1,""k2"":1,""k3"":2,""k4"":2,""k5"":2}",10,11,12.1,13.1 +123456,__GENERATION_NUMBER__,"s1-0",s2-0,1,2,m1-0-0;m1-0-1,m2-0-0;m2-0-1,3;4,6;7,Java C++ Python,4877625602,k1;k2;k3;k4;k5,3;3;3;3;3,"{""k1"":3,""k2"":3,""k3"":3,""k4"":3,""k5"":3}",10,11,12.1,13.1, # Dupliate of row 0 1 +123456,__GENERATION_NUMBER__,s1-2,s2-2,11,21,m1-2-0;m1-2-1,m2-2-0;m2-2-1,32;42,62;72,Java C++ golang,13225573e3f5,k1;k2;k3;k4;k5,4;5;6;7;7,"{""k1"":4,""k2"":5,""k3"":6,""k4"":7,""k5"":7}",10,21,22.1,23.10 +123456,__GENERATION_NUMBER__,s1-2,s2-2,11,21,m1-3-0;m1-3-1,m2-3-0;m2-3-1,32;42,62;72,Java C++ golang,deadbeef,k1;k2;k3;k4;k5,7;7;7;7;7,"{""k1"":7,""k2"":7,""k3"":7,""k4"":7,""k5"":7}",10,21,22.1,23.10, # All sv cols same as prev +123456,__GENERATION_NUMBER__,s1-4,s2-4,41,22,m1-2-0;m1-2-1,m2-2-0;m2-2-1,42;52,72;82,Java C++ golang,deed0507,k1;k2;k3;k4;k5,7;7;8;8;8,"{""k1"":7,""k2"":7,""k3"":8,""k4"":8,""k5"":8}",14,24,24.1,24.10, # All mv cols same as row 2 +123456,__GENERATION_NUMBER__,s1-5,,,32,m1-5-0,m2-2-0,,92;22,golang shell bash,,k1;k2;k3;k4;k5,7;7;7;7;7,"{""k1"":7,""k2"":7,""k3"":7,""k4"":7,""k5"":7}",,24,,24.10, # Default values for some columns +123456,__GENERATION_NUMBER__,s1-6,s2-6,7611,7621,m1-5-0;m1-5-1;m1-5-2,m2-2-0,392;462,6662;782,C++ golang python,deed0507,k1;k2;k3;k4;k5,7;8;9;10;20,"{""k1"":7,""k2"":8,""k3"":9,""k4"":10,""k5"":20}",101,251,262.1,263.10, # 3 values in MV +123456,__GENERATION_NUMBER__,s1-6,s2-6,7611,7621,m1-5-0;m1-5-1;m1-5-2,m2-2-0,392;462,6662;782,C++ golang python,deed0507,k1;k2;k3;k4;k5,7;8;9;10;20,"{""k1"":7,""k2"":8,""k3"":9,""k4"":10,""k5"":20}",2147483647,251,262.1,263.10, # MAX_INT in int metric +123456,__GENERATION_NUMBER__,s1-6,s2-6,7611,7621,m1-5-0;m1-5-1;m1-5-2,m2-2-0,392;462,6662;782,C++ golang python,deed0507,k1;k2;k3;k4;k5,7;8;9;10;20,"{""k1"":7,""k2"":8,""k3"":9,""k4"":10,""k5"":20}",2147483647,251,262.1,263.10, # MAX_INT in int metric +123456,__GENERATION_NUMBER__,s1-7,s2-7,6766,6777,m1-6-0;m1-6-1;m1-6-2;m1-6-3,m2-6-0;m2-6-1,392;462,6662;782,golang Java,d54d0507,k1;k2;k3;k4;k5,31;31;32;32;32,"{""k1"":31,""k2"":31,""k3"":32,""k4"":32,""k5"":32}",87,251,262.10,263.10 diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/recordReaderConfig.json b/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/recordReaderConfig.json new file mode 100644 index 000000000000..ebaa04e10464 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/data/recordReaderConfig.json @@ -0,0 +1,5 @@ +{ + "commentMarker" : "#", + "header" : + "HoursSinceEpoch,generationNumber,stringDimSV1,stringDimSV2,longDimSV1,longDimSV2,stringDimMV1,stringDimMV2,intDimMV1,intDimMV2,textDim1,bytesDimSV1,mapDim1__KEYS,mapDim1__VALUES,mapDim2json,intMetric1,longMetric1,floatMetric1,doubleMetric1" +} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-1.json b/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-1.json new file mode 100644 index 000000000000..381e6885d19a --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-1.json @@ -0,0 +1,46 @@ +{ + "fieldConfigList": [ + { + "encodingType": "RAW", + "indexType": "TEXT", + "name": "textDim1", + "properties": { + "deriveNumDocsPerChunkForRawIndex": "true", + "rawIndexWriterVersion": "3" + } + } + ], + "metadata": { + "customConfigs": { + "d2Name": "" + } + }, + "segmentsConfig": { + "replication": "1", + "retentionTimeUnit": "", + "retentionTimeValue": "", + "segmentAssignmentStrategy": "BalanceNumSegmentAssignmentStrategy", + "segmentPushFrequency": "daily", + "segmentPushType": "REFRESH", + "timeColumnName": "HoursSinceEpoch", + "timeType": "HOURS" + }, + "tableIndexConfig": { + "aggregateMetrics": false, + "autoGeneratedInvertedIndex": false, + "createInvertedIndexDuringSegmentGeneration": false, + "enableDefaultStarTree": false, + "enableDynamicStarTreeCreation": false, + "loadMode": "MMAP", + "noDictionaryColumns": ["textDim1"], + "nullHandlingEnabled": false, + "sortedColumn": [], + "streamConfigs": {} + }, + "tableName": "FeatureTest1_OFFLINE", + "tableType": "OFFLINE", + "tenants": { + "broker": "DefaultTenant", + "server": "DefaultTenant" + } +} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-2-realtime-stream-config.json b/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-2-realtime-stream-config.json new file mode 100644 index 000000000000..441ce201a903 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-2-realtime-stream-config.json @@ -0,0 +1,8 @@ +{ + "streamType": "kafka", + "stream.kafka.consumer.type": "simple", + "topicName": "PinotRealtimeFeatureTest2Event", + "partitionColumn": "longDimSV1", + "numPartitions": "1", + "stream.kafka.consumer.prop.auto.offset.reset": "smallest" +} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-2-realtime.json b/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-2-realtime.json new file mode 100644 index 000000000000..6df335e48841 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/feature-test-2-realtime.json @@ -0,0 +1,61 @@ +{ + "fieldConfigList": [ + { + "encodingType": "RAW", + "indexType": "TEXT", + "name": "textDim1", + "properties": { + "deriveNumDocsPerChunkForRawIndex": "true", + "rawIndexWriterVersion": "3" + } + } + ], + "metadata": { + "customConfigs": { + "d2Name": "" + } + }, + "segmentsConfig": { + "replicasPerPartition": "1", + "replication": "1", + "retentionTimeUnit": "", + "retentionTimeValue": "", + "schemaName": "FeatureTest2", + "segmentAssignmentStrategy": "BalanceNumSegmentAssignmentStrategy", + "segmentPushFrequency": "daily", + "segmentPushType": "APPEND", + "timeColumnName": "HoursSinceEpoch", + "timeType": "HOURS" + }, + "tableIndexConfig": { + "aggregateMetrics": false, + "autoGeneratedInvertedIndex": false, + "bloomFilterColumns": [], + "createInvertedIndexDuringSegmentGeneration": false, + "enableDefaultStarTree": false, + "enableDynamicStarTreeCreation": false, + "loadMode": "MMAP", + "noDictionaryColumns": [], + "nullHandlingEnabled": false, + "segmentFormatVersion": "v3", + "sortedColumn": [], + "streamConfigs": { + "realtime.segment.flush.threshold.size": "63", + "realtime.segment.flush.threshold.time": "1h", + "streamType": "kafka", + "stream.kafka.topic.name": "PinotRealtimeFeatureTest2Event", + "stream.kafka.consumer.type": "simple", + "stream.kafka.decoder.class.name": "org.apache.pinot.plugin.stream.kafka.KafkaJSONMessageDecoder", + "stream.kafka.consumer.factory.class.name": "org.apache.pinot.plugin.stream.kafka20.KafkaConsumerFactory", + "stream.kafka.broker.list": "localhost:19092", + "stream.kafka.zk.broker.url": "localhost:2181/kafka", + "stream.kafka.consumer.prop.auto.offset.reset": "largest" + } + }, + "tableName": "FeatureTest2", + "tableType": "REALTIME", + "tenants": { + "broker": "DefaultTenant", + "server": "DefaultTenant" + } +} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/feature-test-multi-stage.queries b/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/feature-test-multi-stage.queries new file mode 100644 index 000000000000..00d0c6cb8a74 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/feature-test-multi-stage.queries @@ -0,0 +1,31 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Joins +SELECT COUNT(*) FROM FeatureTest1 ft1 INNER JOIN FeatureTest2 ft2 ON ft1.stringDimSV1 = ft2.stringDimSV1 WHERE ft1.generationNumber = __GENERATION_NUMBER__ AND ft2.generationNumber = __GENERATION_NUMBER__ +SELECT ft1.stringDimSV1, COUNT(ft1.stringDimSV1) FROM FeatureTest1 ft1 INNER JOIN FeatureTest2 ft2 ON ft1.stringDimSV1 = ft2.stringDimSV1 WHERE ft1.generationNumber = __GENERATION_NUMBER__ AND ft2.generationNumber = __GENERATION_NUMBER__ GROUP BY ft1.stringDimSV1 +SELECT ft1.stringDimSV2, SUM(ft1.floatMetric1) FROM FeatureTest1 ft1 INNER JOIN FeatureTest2 ft2 ON ft1.stringDimSV1 = ft2.stringDimSV1 WHERE ft1.generationNumber = __GENERATION_NUMBER__ AND ft2.generationNumber = __GENERATION_NUMBER__ GROUP BY ft1.stringDimSV2 +SELECT ft1.stringDimSV1 FROM FeatureTest1 ft1 WHERE ft1.generationNumber = __GENERATION_NUMBER__ AND EXISTS (SELECT 1 FROM FeatureTest2 ft2 WHERE ft2.generationNumber = __GENERATION_NUMBER__ AND ft2.stringDimSV2 = ft1.stringDimSV1) + +# Set operations +SELECT * FROM (SELECT stringDimSV1 FROM FeatureTest1 WHERE generationNumber = __GENERATION_NUMBER__) INTERSECT (SELECT stringDimSV1 FROM FeatureTest2 WHERE generationNumber = __GENERATION_NUMBER__) +SELECT * FROM (SELECT stringDimSV1 FROM FeatureTest1 WHERE generationNumber = __GENERATION_NUMBER__) UNION (SELECT stringDimSV1 FROM FeatureTest2 WHERE generationNumber = __GENERATION_NUMBER__) + +# Windows +SELECT stringDimSV1, longMetric1, SUM(longMetric1) OVER (PARTITION BY stringDimSV1) FROM FeatureTest1 WHERE generationNumber = __GENERATION_NUMBER__ diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/feature-test-multi-stage.results b/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/feature-test-multi-stage.results new file mode 100644 index 000000000000..b073dd9e1cbb --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/feature-test-multi-stage.results @@ -0,0 +1,31 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Joins +{"resultTable":{"dataSchema":{"columnNames":["EXPR$0"],"columnDataTypes":["LONG"]},"rows":[[130]]},"requestId":"11345778000000000","stageStats":{"type":"MAILBOX_RECEIVE","executionTimeMs":36,"emittedRows":1,"fanIn":1,"rawMessages":2,"deserializedBytes":714,"upstreamWaitMs":40,"children":[{"type":"MAILBOX_SEND","executionTimeMs":36,"emittedRows":1,"stage":1,"parallelism":1,"fanOut":1,"rawMessages":2,"serializedBytes":99,"serializationTimeMs":2,"children":[{"type":"AGGREGATE","executionTimeMs":1,"emittedRows":1,"children":[{"type":"MAILBOX_RECEIVE","executionTimeMs":1,"emittedRows":1,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":7,"children":[{"type":"MAILBOX_SEND","executionTimeMs":1,"emittedRows":1,"stage":2,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"AGGREGATE","executionTimeMs":1,"emittedRows":1,"children":[{"type":"HASH_JOIN","executionTimeMs":1,"emittedRows":130,"timeBuildingHashTableMs":1,"children":[{"type":"MAILBOX_RECEIVE","emittedRows":10,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":4,"children":[{"type":"MAILBOX_SEND","executionTimeMs":4,"emittedRows":10,"stage":3,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"LEAF","table":"FeatureTest1","executionTimeMs":5,"emittedRows":10,"numDocsScanned":10,"numEntriesScannedPostFilter":10,"numSegmentsQueried":1,"numSegmentsProcessed":1,"numSegmentsMatched":1,"totalDocs":10}]}]},{"type":"MAILBOX_RECEIVE","emittedRows":66,"fanIn":1,"inMemoryMessages":3,"upstreamWaitMs":3,"children":[{"type":"MAILBOX_SEND","executionTimeMs":4,"emittedRows":66,"stage":4,"parallelism":1,"fanOut":1,"inMemoryMessages":3,"children":[{"type":"LEAF","table":"FeatureTest2","executionTimeMs":4,"emittedRows":66,"numDocsScanned":66,"numEntriesScannedPostFilter":66,"numSegmentsQueried":2,"numSegmentsProcessed":2,"numSegmentsMatched":2,"numConsumingSegmentsQueried":1,"minConsumingFreshnessTimeMs":1716274896691,"totalDocs":66}]}]}]}]}]}]}]}]}]},"brokerId":"Broker_192.168.29.25_8099","exceptions":[],"numServersQueried":0,"numServersResponded":0,"numSegmentsQueried":3,"numSegmentsProcessed":3,"numSegmentsMatched":3,"numConsumingSegmentsQueried":1,"numConsumingSegmentsProcessed":0,"numConsumingSegmentsMatched":0,"numDocsScanned":76,"numEntriesScannedInFilter":0,"numEntriesScannedPostFilter":76,"numGroupsLimitReached":false,"maxRowsInJoinReached":false,"totalDocs":76,"timeUsedMs":975,"offlineThreadCpuTimeNs":0,"realtimeThreadCpuTimeNs":0,"offlineSystemActivitiesCpuTimeNs":0,"realtimeSystemActivitiesCpuTimeNs":0,"offlineResponseSerializationCpuTimeNs":0,"realtimeResponseSerializationCpuTimeNs":0,"offlineTotalCpuTimeNs":0,"realtimeTotalCpuTimeNs":0,"brokerReduceTimeMs":56,"partialResult":false,"exceptionsSize":0,"numRowsResultSet":1,"maxRowsInOperator":130,"numSegmentsPrunedByBroker":0,"minConsumingFreshnessTimeMs":1716274896691,"numSegmentsPrunedByServer":0,"numSegmentsPrunedInvalid":0,"numSegmentsPrunedByLimit":0,"numSegmentsPrunedByValue":0,"explainPlanNumEmptyFilterSegments":0,"explainPlanNumMatchAllFilterSegments":0} +{"resultTable":{"dataSchema":{"columnNames":["stringDimSV1","EXPR$1"],"columnDataTypes":["STRING","LONG"]},"rows":[["s1-4",7],["s1-6",54],["s1-2",28],["s1-0",28],["s1-5",7],["s1-7",6]]},"requestId":"11345778000000001","stageStats":{"type":"MAILBOX_RECEIVE","executionTimeMs":3,"emittedRows":6,"fanIn":1,"rawMessages":2,"deserializedBytes":824,"upstreamWaitMs":3,"children":[{"type":"MAILBOX_SEND","executionTimeMs":2,"emittedRows":6,"stage":1,"parallelism":1,"fanOut":1,"rawMessages":2,"serializedBytes":245,"children":[{"type":"AGGREGATE","executionTimeMs":2,"emittedRows":6,"children":[{"type":"MAILBOX_RECEIVE","executionTimeMs":1,"emittedRows":6,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":4,"children":[{"type":"MAILBOX_SEND","executionTimeMs":1,"emittedRows":6,"stage":2,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"AGGREGATE","executionTimeMs":1,"emittedRows":6,"children":[{"type":"HASH_JOIN","emittedRows":130,"timeBuildingHashTableMs":1,"children":[{"type":"MAILBOX_RECEIVE","emittedRows":10,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":3,"children":[{"type":"MAILBOX_SEND","emittedRows":10,"stage":3,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"LEAF","table":"FeatureTest1","executionTimeMs":1,"emittedRows":10,"numDocsScanned":10,"numEntriesScannedPostFilter":10,"numSegmentsQueried":1,"numSegmentsProcessed":1,"numSegmentsMatched":1,"totalDocs":10}]}]},{"type":"MAILBOX_RECEIVE","emittedRows":66,"fanIn":1,"inMemoryMessages":3,"upstreamWaitMs":3,"children":[{"type":"MAILBOX_SEND","emittedRows":66,"stage":4,"parallelism":1,"fanOut":1,"inMemoryMessages":3,"children":[{"type":"LEAF","table":"FeatureTest2","executionTimeMs":1,"emittedRows":66,"numDocsScanned":66,"numEntriesScannedPostFilter":66,"numSegmentsQueried":2,"numSegmentsProcessed":2,"numSegmentsMatched":2,"numConsumingSegmentsQueried":1,"minConsumingFreshnessTimeMs":1716276514329,"totalDocs":66}]}]}]}]}]}]}]}]}]},"brokerId":"Broker_192.168.29.25_8099","exceptions":[],"numServersQueried":0,"numServersResponded":0,"numSegmentsQueried":3,"numSegmentsProcessed":3,"numSegmentsMatched":3,"numConsumingSegmentsQueried":1,"numConsumingSegmentsProcessed":0,"numConsumingSegmentsMatched":0,"numDocsScanned":76,"numEntriesScannedInFilter":0,"numEntriesScannedPostFilter":76,"numGroupsLimitReached":false,"maxRowsInJoinReached":false,"totalDocs":76,"timeUsedMs":49,"offlineThreadCpuTimeNs":0,"realtimeThreadCpuTimeNs":0,"offlineSystemActivitiesCpuTimeNs":0,"realtimeSystemActivitiesCpuTimeNs":0,"offlineResponseSerializationCpuTimeNs":0,"realtimeResponseSerializationCpuTimeNs":0,"offlineTotalCpuTimeNs":0,"realtimeTotalCpuTimeNs":0,"brokerReduceTimeMs":3,"partialResult":false,"maxRowsInOperator":130,"numRowsResultSet":6,"exceptionsSize":0,"numSegmentsPrunedByBroker":0,"minConsumingFreshnessTimeMs":1716276514329,"numSegmentsPrunedByServer":0,"numSegmentsPrunedInvalid":0,"numSegmentsPrunedByLimit":0,"numSegmentsPrunedByValue":0,"explainPlanNumEmptyFilterSegments":0,"explainPlanNumMatchAllFilterSegments":0} +{"resultTable":{"dataSchema":{"columnNames":["stringDimSV2","EXPR$1"],"columnDataTypes":["STRING","DOUBLE"]},"rows":[["s2-6",14153.4],["s2-2",618.8],["s2-0",338.8],["s2-4",168.7],["s2-7",1572.6],["null",0.0]]},"requestId":"11345778000000002","stageStats":{"type":"MAILBOX_RECEIVE","executionTimeMs":2,"emittedRows":6,"fanIn":1,"rawMessages":2,"deserializedBytes":826,"upstreamWaitMs":3,"children":[{"type":"MAILBOX_SEND","executionTimeMs":4,"emittedRows":6,"stage":1,"parallelism":1,"fanOut":1,"rawMessages":2,"serializedBytes":247,"children":[{"type":"AGGREGATE","executionTimeMs":3,"emittedRows":6,"children":[{"type":"MAILBOX_RECEIVE","executionTimeMs":3,"emittedRows":6,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":4,"children":[{"type":"MAILBOX_SEND","executionTimeMs":2,"emittedRows":6,"stage":2,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"AGGREGATE","executionTimeMs":2,"emittedRows":6,"children":[{"type":"HASH_JOIN","executionTimeMs":1,"emittedRows":130,"timeBuildingHashTableMs":1,"children":[{"type":"MAILBOX_RECEIVE","emittedRows":10,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":2,"children":[{"type":"MAILBOX_SEND","emittedRows":10,"stage":3,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"LEAF","table":"FeatureTest1","emittedRows":10,"numDocsScanned":10,"numEntriesScannedPostFilter":30,"numSegmentsQueried":1,"numSegmentsProcessed":1,"numSegmentsMatched":1,"totalDocs":10}]}]},{"type":"MAILBOX_RECEIVE","executionTimeMs":1,"emittedRows":66,"fanIn":1,"inMemoryMessages":3,"upstreamWaitMs":2,"children":[{"type":"MAILBOX_SEND","emittedRows":66,"stage":4,"parallelism":1,"fanOut":1,"inMemoryMessages":3,"children":[{"type":"LEAF","table":"FeatureTest2","emittedRows":66,"numDocsScanned":66,"numEntriesScannedPostFilter":66,"numSegmentsQueried":2,"numSegmentsProcessed":2,"numSegmentsMatched":2,"numConsumingSegmentsQueried":1,"minConsumingFreshnessTimeMs":1716276645639,"totalDocs":66}]}]}]}]}]}]}]}]}]},"brokerId":"Broker_192.168.29.25_8099","exceptions":[],"numServersQueried":0,"numServersResponded":0,"numSegmentsQueried":3,"numSegmentsProcessed":3,"numSegmentsMatched":3,"numConsumingSegmentsQueried":1,"numConsumingSegmentsProcessed":0,"numConsumingSegmentsMatched":0,"numDocsScanned":76,"numEntriesScannedInFilter":0,"numEntriesScannedPostFilter":96,"numGroupsLimitReached":false,"maxRowsInJoinReached":false,"totalDocs":76,"timeUsedMs":21,"offlineThreadCpuTimeNs":0,"realtimeThreadCpuTimeNs":0,"offlineSystemActivitiesCpuTimeNs":0,"realtimeSystemActivitiesCpuTimeNs":0,"offlineResponseSerializationCpuTimeNs":0,"realtimeResponseSerializationCpuTimeNs":0,"offlineTotalCpuTimeNs":0,"realtimeTotalCpuTimeNs":0,"brokerReduceTimeMs":3,"partialResult":false,"exceptionsSize":0,"numRowsResultSet":6,"maxRowsInOperator":130,"numSegmentsPrunedByBroker":0,"minConsumingFreshnessTimeMs":1716276645639,"numSegmentsPrunedByServer":0,"numSegmentsPrunedInvalid":0,"numSegmentsPrunedByLimit":0,"numSegmentsPrunedByValue":0,"explainPlanNumEmptyFilterSegments":0,"explainPlanNumMatchAllFilterSegments":0} +{"resultTable":{"dataSchema":{"columnNames":["stringDimSV1"],"columnDataTypes":["STRING"]},"rows":[]},"requestId":"11345778000000003","stageStats":{"type":"MAILBOX_RECEIVE","fanIn":1,"rawMessages":1,"deserializedBytes":132,"children":[{"type":"MAILBOX_SEND","stage":1,"parallelism":1,"fanOut":1,"rawMessages":1,"children":[{"type":"LEAF","table":"FeatureTest1","numSegmentsQueried":1,"totalDocs":10,"numSegmentsPrunedByServer":1}]}]},"brokerId":"Broker_192.168.29.25_8099","exceptions":[],"numServersQueried":0,"numServersResponded":0,"numSegmentsQueried":1,"numSegmentsProcessed":0,"numSegmentsMatched":0,"numConsumingSegmentsQueried":0,"numConsumingSegmentsProcessed":0,"numConsumingSegmentsMatched":0,"numDocsScanned":0,"numEntriesScannedInFilter":0,"numEntriesScannedPostFilter":0,"numGroupsLimitReached":false,"maxRowsInJoinReached":false,"totalDocs":10,"timeUsedMs":43,"offlineThreadCpuTimeNs":0,"realtimeThreadCpuTimeNs":0,"offlineSystemActivitiesCpuTimeNs":0,"realtimeSystemActivitiesCpuTimeNs":0,"offlineResponseSerializationCpuTimeNs":0,"realtimeResponseSerializationCpuTimeNs":0,"offlineTotalCpuTimeNs":0,"realtimeTotalCpuTimeNs":0,"brokerReduceTimeMs":0,"partialResult":false,"maxRowsInOperator":0,"exceptionsSize":0,"numRowsResultSet":0,"numSegmentsPrunedByBroker":0,"minConsumingFreshnessTimeMs":0,"numSegmentsPrunedByServer":1,"numSegmentsPrunedInvalid":0,"numSegmentsPrunedByLimit":0,"numSegmentsPrunedByValue":0,"explainPlanNumEmptyFilterSegments":0,"explainPlanNumMatchAllFilterSegments":0} + +# Set operations +{"resultTable":{"dataSchema":{"columnNames":["stringDimSV1"],"columnDataTypes":["STRING"]},"rows":[["s1-0"],["s1-2"],["s1-4"],["s1-5"],["s1-6"],["s1-7"]]},"requestId":"11345778000000004","stageStats":{"type":"MAILBOX_RECEIVE","executionTimeMs":1,"emittedRows":6,"fanIn":1,"rawMessages":2,"deserializedBytes":628,"upstreamWaitMs":2,"children":[{"type":"MAILBOX_SEND","executionTimeMs":1,"emittedRows":6,"stage":1,"parallelism":1,"fanOut":1,"rawMessages":2,"serializedBytes":174,"children":[{"children":[{"type":"MAILBOX_RECEIVE","emittedRows":6,"children":[{"type":"MAILBOX_SEND","emittedRows":10,"stage":2,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"LEAF","table":"FeatureTest1","emittedRows":10,"numDocsScanned":10,"numEntriesScannedPostFilter":10,"numSegmentsQueried":1,"numSegmentsProcessed":1,"numSegmentsMatched":1,"totalDocs":10}]}]},{"type":"MAILBOX_RECEIVE","emittedRows":10,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":1,"children":[{"type":"MAILBOX_SEND","emittedRows":66,"stage":3,"parallelism":1,"fanOut":1,"inMemoryMessages":3,"children":[{"type":"LEAF","table":"FeatureTest2","emittedRows":66,"numDocsScanned":66,"numEntriesScannedPostFilter":66,"numSegmentsQueried":2,"numSegmentsProcessed":2,"numSegmentsMatched":2,"numConsumingSegmentsQueried":1,"minConsumingFreshnessTimeMs":1716276920346,"totalDocs":66}]}]}]}]}]},"brokerId":"Broker_192.168.29.25_8099","exceptions":[],"numServersQueried":0,"numServersResponded":0,"numSegmentsQueried":3,"numSegmentsProcessed":3,"numSegmentsMatched":3,"numConsumingSegmentsQueried":1,"numConsumingSegmentsProcessed":0,"numConsumingSegmentsMatched":0,"numDocsScanned":76,"numEntriesScannedInFilter":0,"numEntriesScannedPostFilter":76,"numGroupsLimitReached":false,"maxRowsInJoinReached":false,"totalDocs":76,"timeUsedMs":35,"offlineThreadCpuTimeNs":0,"realtimeThreadCpuTimeNs":0,"offlineSystemActivitiesCpuTimeNs":0,"realtimeSystemActivitiesCpuTimeNs":0,"offlineResponseSerializationCpuTimeNs":0,"realtimeResponseSerializationCpuTimeNs":0,"offlineTotalCpuTimeNs":0,"realtimeTotalCpuTimeNs":0,"brokerReduceTimeMs":2,"partialResult":false,"numSegmentsPrunedByBroker":0,"minConsumingFreshnessTimeMs":1716276920346,"numSegmentsPrunedByServer":0,"numSegmentsPrunedInvalid":0,"numSegmentsPrunedByLimit":0,"numSegmentsPrunedByValue":0,"explainPlanNumEmptyFilterSegments":0,"explainPlanNumMatchAllFilterSegments":0,"numRowsResultSet":6,"exceptionsSize":0,"maxRowsInOperator":66} +{"resultTable":{"dataSchema":{"columnNames":["stringDimSV1"],"columnDataTypes":["STRING"]},"rows":[["s1-4"],["s1-6"],["s1-2"],["s1-0"],["s1-5"],["s1-7"]]},"requestId":"11345778000000006","stageStats":{"type":"MAILBOX_RECEIVE","emittedRows":6,"fanIn":1,"rawMessages":2,"deserializedBytes":744,"upstreamWaitMs":1,"children":[{"type":"MAILBOX_SEND","executionTimeMs":2,"emittedRows":6,"stage":1,"parallelism":1,"fanOut":1,"rawMessages":2,"serializedBytes":174,"children":[{"type":"AGGREGATE","executionTimeMs":2,"emittedRows":6,"children":[{"type":"MAILBOX_RECEIVE","executionTimeMs":2,"emittedRows":6,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":3,"children":[{"type":"MAILBOX_SEND","executionTimeMs":1,"emittedRows":6,"stage":2,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"AGGREGATE","executionTimeMs":1,"emittedRows":6,"children":[{"type":"UNION","executionTimeMs":1,"emittedRows":76,"children":[{"type":"MAILBOX_RECEIVE","executionTimeMs":1,"emittedRows":10,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":1,"children":[{"type":"MAILBOX_SEND","emittedRows":10,"stage":3,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"LEAF","table":"FeatureTest1","emittedRows":10,"numDocsScanned":10,"numEntriesScannedPostFilter":10,"numSegmentsQueried":1,"numSegmentsProcessed":1,"numSegmentsMatched":1,"totalDocs":10}]}]},{"type":"MAILBOX_RECEIVE","emittedRows":66,"fanIn":1,"inMemoryMessages":3,"upstreamWaitMs":1,"children":[{"type":"MAILBOX_SEND","emittedRows":66,"stage":4,"parallelism":1,"fanOut":1,"inMemoryMessages":3,"children":[{"type":"LEAF","table":"FeatureTest2","emittedRows":66,"numDocsScanned":66,"numEntriesScannedPostFilter":66,"numSegmentsQueried":2,"numSegmentsProcessed":2,"numSegmentsMatched":2,"numConsumingSegmentsQueried":1,"minConsumingFreshnessTimeMs":1716276920346,"totalDocs":66}]}]}]}]}]}]}]}]}]},"brokerId":"Broker_192.168.29.25_8099","exceptions":[],"numServersQueried":0,"numServersResponded":0,"numSegmentsQueried":3,"numSegmentsProcessed":3,"numSegmentsMatched":3,"numConsumingSegmentsQueried":1,"numConsumingSegmentsProcessed":0,"numConsumingSegmentsMatched":0,"numDocsScanned":76,"numEntriesScannedInFilter":0,"numEntriesScannedPostFilter":76,"numGroupsLimitReached":false,"maxRowsInJoinReached":false,"totalDocs":76,"timeUsedMs":13,"offlineThreadCpuTimeNs":0,"realtimeThreadCpuTimeNs":0,"offlineSystemActivitiesCpuTimeNs":0,"realtimeSystemActivitiesCpuTimeNs":0,"offlineResponseSerializationCpuTimeNs":0,"realtimeResponseSerializationCpuTimeNs":0,"offlineTotalCpuTimeNs":0,"realtimeTotalCpuTimeNs":0,"brokerReduceTimeMs":1,"partialResult":false,"numSegmentsPrunedByBroker":0,"minConsumingFreshnessTimeMs":1716276920346,"numSegmentsPrunedByServer":0,"numSegmentsPrunedInvalid":0,"numSegmentsPrunedByLimit":0,"numSegmentsPrunedByValue":0,"explainPlanNumEmptyFilterSegments":0,"explainPlanNumMatchAllFilterSegments":0,"numRowsResultSet":6,"exceptionsSize":0,"maxRowsInOperator":76} + +# Windows +{"resultTable":{"dataSchema":{"columnNames":["stringDimSV1","longMetric1","EXPR$2"],"columnDataTypes":["STRING","LONG","LONG"]},"rows":[["s1-5",24,24],["s1-4",24,24],["s1-7",251,251],["s1-6",251,753],["s1-6",251,753],["s1-6",251,753],["s1-0",11,22],["s1-0",11,22],["s1-2",21,42],["s1-2",21,42]]},"partialResult":false,"exceptions":[],"numGroupsLimitReached":false,"maxRowsInJoinReached":false,"timeUsedMs":34,"stageStats":{"type":"MAILBOX_RECEIVE","executionTimeMs":3,"emittedRows":10,"fanIn":1,"rawMessages":2,"deserializedBytes":717,"upstreamWaitMs":1,"children":[{"type":"MAILBOX_SEND","executionTimeMs":1,"emittedRows":10,"stage":1,"parallelism":1,"fanOut":1,"rawMessages":2,"serializedBytes":407,"children":[{"type":"TRANSFORM","executionTimeMs":1,"emittedRows":10,"children":[{"type":"WINDOW","executionTimeMs":1,"emittedRows":10,"children":[{"type":"MAILBOX_RECEIVE","emittedRows":10,"fanIn":1,"inMemoryMessages":2,"upstreamWaitMs":2,"children":[{"type":"MAILBOX_SEND","emittedRows":10,"stage":2,"parallelism":1,"fanOut":1,"inMemoryMessages":2,"children":[{"type":"LEAF","table":"FeatureTest1","executionTimeMs":1,"emittedRows":10,"numDocsScanned":10,"totalDocs":10,"numEntriesScannedPostFilter":20,"numSegmentsQueried":1,"numSegmentsProcessed":1,"numSegmentsMatched":1}]}]}]}]}]}]},"maxRowsInOperator":10,"requestId":"11345778000000006","brokerId":"Broker_192.168.29.25_8099","numDocsScanned":10,"totalDocs":10,"numEntriesScannedInFilter":0,"numEntriesScannedPostFilter":20,"numServersQueried":0,"numServersResponded":0,"numSegmentsQueried":1,"numSegmentsProcessed":1,"numSegmentsMatched":1,"numConsumingSegmentsQueried":0,"numConsumingSegmentsProcessed":0,"numConsumingSegmentsMatched":0,"minConsumingFreshnessTimeMs":0,"numSegmentsPrunedByBroker":0,"numSegmentsPrunedByServer":0,"numSegmentsPrunedInvalid":0,"numSegmentsPrunedByLimit":0,"numSegmentsPrunedByValue":0,"brokerReduceTimeMs":3,"offlineThreadCpuTimeNs":0,"realtimeThreadCpuTimeNs":0,"offlineSystemActivitiesCpuTimeNs":0,"realtimeSystemActivitiesCpuTimeNs":0,"offlineResponseSerializationCpuTimeNs":0,"realtimeResponseSerializationCpuTimeNs":0,"offlineTotalCpuTimeNs":0,"realtimeTotalCpuTimeNs":0,"explainPlanNumEmptyFilterSegments":0,"explainPlanNumMatchAllFilterSegments":0,"traceInfo":{}} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/post-broker-rollback.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/post-broker-rollback.yaml new file mode 100644 index 000000000000..3e927386b107 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/post-broker-rollback.yaml @@ -0,0 +1,43 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Operations to be done. +description: Operations to be run after broker rollback +operations: + - type: segmentOp + description: Add segment FeatureTest1_Segment6 to table FeatureTest1 + op: UPLOAD + inputDataFileName: data/FeatureTest1-data-00.csv + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json + recordReaderConfigFileName: data/recordReaderConfig.json + segmentName: FeatureTest1_Segment6 + - type: streamOp + description: publish rows to PinotRealtimeFeatureTest2Event + op: PRODUCE + streamConfigFileName: feature-test-2-realtime-stream-config.json + numRows: 66 + inputDataFileName: data/FeatureTest2-data-realtime-00.csv + recordReaderConfigFileName: data/recordReaderConfig.json + tableConfigFileName: feature-test-2-realtime.json + - type: queryOp + description: Run multi-stage queries on FeatureTest1 and FeatureTest2 using SQL + useMultiStageQueryEngine: true + queryFileName: queries/feature-test-multi-stage.queries + expectedResultsFileName: query-results/feature-test-multi-stage.results diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/post-controller-rollback.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/post-controller-rollback.yaml new file mode 100644 index 000000000000..c3bcbf631c0d --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/post-controller-rollback.yaml @@ -0,0 +1,53 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Operations to be done. +description: Operations to be run after controller rollback +operations: + - type: segmentOp + description: Add segment FeatureTest1_Segment7 to table FeatureTest1 + op: UPLOAD + inputDataFileName: data/FeatureTest1-data-00.csv + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json + recordReaderConfigFileName: data/recordReaderConfig.json + segmentName: FeatureTest1_Segment7 + - type: streamOp + description: publish rows to PinotRealtimeFeatureTest2Event + op: PRODUCE + streamConfigFileName: feature-test-2-realtime-stream-config.json + numRows: 66 + inputDataFileName: data/FeatureTest2-data-realtime-00.csv + recordReaderConfigFileName: data/recordReaderConfig.json + tableConfigFileName: feature-test-2-realtime.json + - type: queryOp + description: Run multi-stage queries on FeatureTest1 and FeatureTest2 using SQL + useMultiStageQueryEngine: true + queryFileName: queries/feature-test-multi-stage.queries + expectedResultsFileName: query-results/feature-test-multi-stage.results + - type: segmentOp + description: Delete segment FeatureTest1_Segment + op: DELETE + tableConfigFileName: feature-test-1.json + segmentName: FeatureTest1_Segment + - type: tableOp + description: Delete table feature-test-1.json + op: DELETE + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-rollback.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-rollback.yaml new file mode 100644 index 000000000000..c31a41294bb4 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-rollback.yaml @@ -0,0 +1,43 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Operations to be done. +description: Operations to be run after server rollback +operations: + - type: segmentOp + description: Add segment FeatureTest1_Segment5 to table FeatureTest1 + op: UPLOAD + inputDataFileName: data/FeatureTest1-data-00.csv + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json + recordReaderConfigFileName: data/recordReaderConfig.json + segmentName: FeatureTest1_Segment5 + - type: streamOp + description: publish rows to PinotRealtimeFeatureTest2Event + op: PRODUCE + streamConfigFileName: feature-test-2-realtime-stream-config.json + numRows: 66 + inputDataFileName: data/FeatureTest2-data-realtime-00.csv + recordReaderConfigFileName: data/recordReaderConfig.json + tableConfigFileName: feature-test-2-realtime.json + - type: queryOp + description: Run multi-stage queries on FeatureTest1 and FeatureTest2 using SQL + useMultiStageQueryEngine: true + queryFileName: queries/feature-test-multi-stage.queries + expectedResultsFileName: query-results/feature-test-multi-stage.results diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-upgrade.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-upgrade.yaml new file mode 100644 index 000000000000..786cb3fb3981 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-upgrade.yaml @@ -0,0 +1,43 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Operations to be done. +description: Operations to be run after server upgrade +operations: + - type: segmentOp + description: Add segment FeatureTest1_Segment4 to table FeatureTest1 + op: UPLOAD + inputDataFileName: data/FeatureTest1-data-00.csv + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json + recordReaderConfigFileName: data/recordReaderConfig.json + segmentName: FeatureTest1_Segment4 + - type: streamOp + description: publish rows to PinotRealtimeFeatureTest2Event + op: PRODUCE + streamConfigFileName: feature-test-2-realtime-stream-config.json + numRows: 66 + inputDataFileName: data/FeatureTest2-data-realtime-00.csv + recordReaderConfigFileName: data/recordReaderConfig.json + tableConfigFileName: feature-test-2-realtime.json + - type: queryOp + description: Run multi-stage queries on FeatureTest1 and FeatureTest2 using SQL + useMultiStageQueryEngine: true + queryFileName: queries/feature-test-multi-stage.queries + expectedResultsFileName: query-results/feature-test-multi-stage.results diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/pre-broker-upgrade.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/pre-broker-upgrade.yaml new file mode 100644 index 000000000000..4e255b8cf8a8 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/pre-broker-upgrade.yaml @@ -0,0 +1,43 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Operations to be done. +description: Operations to be run before Broker upgrade +operations: + - type: segmentOp + description: Add segment FeatureTest1_Segment2 to table FeatureTest1 + op: UPLOAD + inputDataFileName: data/FeatureTest1-data-00.csv + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json + recordReaderConfigFileName: data/recordReaderConfig.json + segmentName: FeatureTest1_Segment2 + - type: streamOp + description: publish rows to PinotRealtimeFeatureTest2Event + op: PRODUCE + streamConfigFileName: feature-test-2-realtime-stream-config.json + numRows: 66 + inputDataFileName: data/FeatureTest2-data-realtime-00.csv + recordReaderConfigFileName: data/recordReaderConfig.json + tableConfigFileName: feature-test-2-realtime.json + - type: queryOp + description: Run multi-stage queries on FeatureTest1 and FeatureTest2 using SQL + useMultiStageQueryEngine: true + queryFileName: queries/feature-test-multi-stage.queries + expectedResultsFileName: query-results/feature-test-multi-stage.results diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/pre-controller-upgrade.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/pre-controller-upgrade.yaml new file mode 100644 index 000000000000..39f8bcacdf09 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/pre-controller-upgrade.yaml @@ -0,0 +1,58 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Operations to be done. +description: Operations to be run before Controller upgrade +operations: + - type: streamOp + description: create Kafka topic PinotRealtimeFeatureTest2Event + op: CREATE + streamConfigFileName: feature-test-2-realtime-stream-config.json + - type: tableOp + description: Create realtime table FeatureTest2 + op: CREATE + schemaFileName: FeatureTest2-schema.json + tableConfigFileName: feature-test-2-realtime.json + recordReaderConfigFileName: data/recordReaderConfig.json + - type: tableOp + description: Create offline table FeatureTest1 + op: CREATE + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json + - type: segmentOp + description: Add segment FeatureTest1_Segment to table FeatureTest1 + op: UPLOAD + inputDataFileName: data/FeatureTest1-data-00.csv + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json + recordReaderConfigFileName: data/recordReaderConfig.json + segmentName: FeatureTest1_Segment + - type: streamOp + description: publish rows to PinotRealtimeFeatureTest2Event + op: PRODUCE + streamConfigFileName: feature-test-2-realtime-stream-config.json + numRows: 66 + inputDataFileName: data/FeatureTest2-data-realtime-00.csv + recordReaderConfigFileName: data/recordReaderConfig.json + tableConfigFileName: feature-test-2-realtime.json + - type: queryOp + description: Run multi-stage queries on FeatureTest1 and FeatureTest2 using SQL + useMultiStageQueryEngine: true + queryFileName: queries/feature-test-multi-stage.queries + expectedResultsFileName: query-results/feature-test-multi-stage.results diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/pre-server-upgrade.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/pre-server-upgrade.yaml new file mode 100644 index 000000000000..167f71fadf2f --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/pre-server-upgrade.yaml @@ -0,0 +1,43 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Operations to be done. +description: Operations to be run before server upgrade +operations: + - type: segmentOp + description: Add segment FeatureTest1_Segment3 to table FeatureTest1 + op: UPLOAD + inputDataFileName: data/FeatureTest1-data-00.csv + schemaFileName: FeatureTest1-schema.json + tableConfigFileName: feature-test-1.json + recordReaderConfigFileName: data/recordReaderConfig.json + segmentName: FeatureTest1_Segment3 + - type: streamOp + description: publish rows to PinotRealtimeFeatureTest2Event + op: PRODUCE + streamConfigFileName: feature-test-2-realtime-stream-config.json + numRows: 66 + inputDataFileName: data/FeatureTest2-data-realtime-00.csv + recordReaderConfigFileName: data/recordReaderConfig.json + tableConfigFileName: feature-test-2-realtime.json + - type: queryOp + description: Run multi-stage queries on FeatureTest1 and FeatureTest2 using SQL + useMultiStageQueryEngine: true + queryFileName: queries/feature-test-multi-stage.queries + expectedResultsFileName: query-results/feature-test-multi-stage.results diff --git a/docker/images/pinot-base/pinot-base-build/amazoncorretto.dockerfile b/docker/images/pinot-base/pinot-base-build/amazoncorretto.dockerfile index 05a5e73e46fe..79126ed713ce 100644 --- a/docker/images/pinot-base/pinot-base-build/amazoncorretto.dockerfile +++ b/docker/images/pinot-base/pinot-base-build/amazoncorretto.dockerfile @@ -46,7 +46,7 @@ ENV JAVA_HOME=/usr/lib/jvm/java-11-amazon-corretto # install maven RUN mkdir -p /usr/share/maven /usr/share/maven/ref \ - && wget https://dlcdn.apache.org/maven/maven-3/3.9.6/binaries/apache-maven-3.9.6-bin.tar.gz -P /tmp \ + && wget https://dlcdn.apache.org/maven/maven-3/3.9.7/binaries/apache-maven-3.9.7-bin.tar.gz -P /tmp \ && tar -xzf /tmp/apache-maven-*.tar.gz -C /usr/share/maven --strip-components=1 \ && rm -f /tmp/apache-maven-*.tar.gz \ && ln -s /usr/share/maven/bin/mvn /usr/bin/mvn diff --git a/docker/images/pinot-base/pinot-base-build/ms-openjdk.dockerfile b/docker/images/pinot-base/pinot-base-build/ms-openjdk.dockerfile index 14ea0cd84fca..8255d84a1608 100644 --- a/docker/images/pinot-base/pinot-base-build/ms-openjdk.dockerfile +++ b/docker/images/pinot-base/pinot-base-build/ms-openjdk.dockerfile @@ -30,7 +30,7 @@ RUN apt-get update && \ # install maven RUN mkdir -p /usr/share/maven /usr/share/maven/ref \ - && wget https://dlcdn.apache.org/maven/maven-3/3.9.6/binaries/apache-maven-3.9.6-bin.tar.gz -P /tmp \ + && wget https://dlcdn.apache.org/maven/maven-3/3.9.7/binaries/apache-maven-3.9.7-bin.tar.gz -P /tmp \ && tar -xzf /tmp/apache-maven-*.tar.gz -C /usr/share/maven --strip-components=1 \ && rm -f /tmp/apache-maven-*.tar.gz \ && ln -s /usr/share/maven/bin/mvn /usr/bin/mvn diff --git a/docker/images/pinot-base/pinot-base-build/openjdk.dockerfile b/docker/images/pinot-base/pinot-base-build/openjdk.dockerfile index ab62df775d15..3e0eb30adc3a 100644 --- a/docker/images/pinot-base/pinot-base-build/openjdk.dockerfile +++ b/docker/images/pinot-base/pinot-base-build/openjdk.dockerfile @@ -34,7 +34,7 @@ RUN apt-get update && \ # install maven RUN mkdir -p /usr/share/maven /usr/share/maven/ref \ - && wget https://dlcdn.apache.org/maven/maven-3/3.9.6/binaries/apache-maven-3.9.6-bin.tar.gz -P /tmp \ + && wget https://dlcdn.apache.org/maven/maven-3/3.9.7/binaries/apache-maven-3.9.7-bin.tar.gz -P /tmp \ && tar -xzf /tmp/apache-maven-*.tar.gz -C /usr/share/maven --strip-components=1 \ && rm -f /tmp/apache-maven-*.tar.gz \ && ln -s /usr/share/maven/bin/mvn /usr/bin/mvn \ diff --git a/docker/images/pinot-presto/Dockerfile b/docker/images/pinot-presto/Dockerfile index 7fdf2c6a0f23..16ec15385db8 100644 --- a/docker/images/pinot-presto/Dockerfile +++ b/docker/images/pinot-presto/Dockerfile @@ -43,7 +43,7 @@ USER presto RUN git clone ${PRESTO_GIT_URL} ${PRESTO_BUILD_DIR} && \ cd ${PRESTO_BUILD_DIR} && \ git checkout ${PRESTO_BRANCH} && \ - echo "distributionUrl=https://repo1.maven.org/maven2/org/apache/maven/apache-maven/3.9.6/apache-maven-3.9.6-bin.zip" > .mvn/wrapper/maven-wrapper.properties && \ + echo "distributionUrl=https://repo1.maven.org/maven2/org/apache/maven/apache-maven/3.9.7/apache-maven-3.9.7-bin.zip" > .mvn/wrapper/maven-wrapper.properties && \ ./mvnw clean install -DskipTests && \ mkdir -p ${PRESTO_HOME}/data && \ cp -r presto-server/target/presto-server-*/presto-server-*/* ${PRESTO_HOME} && \ diff --git a/helm/index.yaml b/helm/index.yaml index dc8a9446e1de..f08bd52f4c48 100644 --- a/helm/index.yaml +++ b/helm/index.yaml @@ -1,6 +1,34 @@ apiVersion: v1 entries: pinot: + - apiVersion: v1 + appVersion: 1.0.0 + created: "2024-05-30T09:29:30.175007-04:00" + dependencies: + - condition: pinot.zookeeper.enabled,zookeeper.enabled + name: zookeeper + repository: https://charts.bitnami.com/bitnami + version: 13.x.x + description: Apache Pinot is a realtime distributed OLAP datastore, which is used + to deliver scalable real time analytics with low latency. It can ingest data + from offline data sources (such as Hadoop and flat files) as well as online + sources (such as Kafka). Pinot is designed to scale horizontally. + digest: a39ab0af4e8b3c179516b071a79232d52d987b547cfb358bd46140dd27ce2530 + home: https://pinot.apache.org/ + keywords: + - olap + - analytics + - database + - pinot + maintainers: + - email: dev@pinot.apache.org + name: pinot-dev + name: pinot + sources: + - https://github.com/apache/pinot/tree/master/helm + urls: + - pinot-0.2.9.tgz + version: 0.2.9 - apiVersion: v1 appVersion: 1.0.0 created: "2024-02-22T11:12:34.734773+08:00" diff --git a/helm/pinot-0.2.9.tgz b/helm/pinot-0.2.9.tgz new file mode 100644 index 000000000000..20c345d6247f Binary files /dev/null and b/helm/pinot-0.2.9.tgz differ diff --git a/helm/pinot/Chart.yaml b/helm/pinot/Chart.yaml index d8cad372dec4..8f3518e86313 100644 --- a/helm/pinot/Chart.yaml +++ b/helm/pinot/Chart.yaml @@ -21,7 +21,7 @@ apiVersion: v1 appVersion: 1.0.0 name: pinot description: Apache Pinot is a realtime distributed OLAP datastore, which is used to deliver scalable real time analytics with low latency. It can ingest data from offline data sources (such as Hadoop and flat files) as well as online sources (such as Kafka). Pinot is designed to scale horizontally. -version: 0.2.9-SNAPSHOT +version: 0.3.0-SNAPSHOT keywords: - olap - analytics diff --git a/helm/pinot/charts/zookeeper-13.2.0.tgz b/helm/pinot/charts/zookeeper-13.2.0.tgz new file mode 100644 index 000000000000..d044e68b8f0e Binary files /dev/null and b/helm/pinot/charts/zookeeper-13.2.0.tgz differ diff --git a/helm/pinot/charts/zookeeper-7.0.0.tgz b/helm/pinot/charts/zookeeper-7.0.0.tgz deleted file mode 100644 index 1925c8f9959a..000000000000 Binary files a/helm/pinot/charts/zookeeper-7.0.0.tgz and /dev/null differ diff --git a/helm/pinot/requirements.lock b/helm/pinot/requirements.lock index 4a3dfaaea73e..3269c660f02d 100644 --- a/helm/pinot/requirements.lock +++ b/helm/pinot/requirements.lock @@ -20,6 +20,6 @@ dependencies: - name: zookeeper repository: https://charts.bitnami.com/bitnami - version: 9.2.7 -digest: sha256:a5da7ddd352d63b0a0e1e5bd85e90e304ae5d0fa7759d5cb7ffb39f61adef1e9 -generated: "2022-06-20T14:57:34.981883-07:00" + version: 13.2.0 +digest: sha256:d9252aa40c5511fa1124c0db142b3222f8aff999ee0add757bbb6f19b6f45826 +generated: "2024-05-05T19:02:58.334678-04:00" diff --git a/helm/pinot/requirements.yaml b/helm/pinot/requirements.yaml index ece3625593cb..23448f616cbb 100644 --- a/helm/pinot/requirements.yaml +++ b/helm/pinot/requirements.yaml @@ -19,6 +19,6 @@ dependencies: - name: zookeeper - version: 9.x.x + version: 13.x.x repository: https://charts.bitnami.com/bitnami condition: pinot.zookeeper.enabled,zookeeper.enabled diff --git a/helm/pinot/templates/broker/poddisruptionbudget.yaml b/helm/pinot/templates/broker/poddisruptionbudget.yaml new file mode 100644 index 000000000000..1c22db404937 --- /dev/null +++ b/helm/pinot/templates/broker/poddisruptionbudget.yaml @@ -0,0 +1,35 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +{{- if .Values.broker.pdb.enabled }} +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: {{ include "pinot.broker.fullname" . }} +spec: + {{- if .Values.broker.pdb.minAvailable }} + minAvailable: {{ .Values.broker.pdb.minAvailable }} + {{- end }} + {{- if .Values.broker.pdb.maxUnavailable }} + maxUnavailable: {{ .Values.broker.pdb.maxUnavailable }} + {{- end }} + selector: + matchLabels: + {{- include "pinot.brokerMatchLabels" . | nindent 6 }} +{{- end }} diff --git a/helm/pinot/templates/broker/statefulset.yaml b/helm/pinot/templates/broker/statefulset.yaml index de69c0408f42..342a368d2820 100644 --- a/helm/pinot/templates/broker/statefulset.yaml +++ b/helm/pinot/templates/broker/statefulset.yaml @@ -37,7 +37,12 @@ spec: labels: {{- include "pinot.brokerLabels" . | nindent 8 }} annotations: -{{ toYaml .Values.broker.podAnnotations | indent 8 }} + {{- if .Values.broker.automaticReload.enabled }} + checksum/config: {{ include (print $.Template.BasePath "/broker/configmap.yaml") . | sha256sum }} + {{- end }} + {{- with .Values.broker.podAnnotations }} + {{- toYaml . | nindent 8 }} + {{- end }} spec: terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds }} serviceAccountName: {{ include "pinot.serviceAccountName" . }} @@ -88,16 +93,33 @@ spec: {{- end }} {{- if .Values.broker.probes.livenessEnabled }} livenessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.broker.probes.liveness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.broker.probes.liveness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.broker.probes.liveness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.broker.probes.liveness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.broker.probes.liveness.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: path: {{ .Values.broker.probes.endpoint }} port: {{ .Values.broker.service.port }} {{- end }} {{- if .Values.broker.probes.readinessEnabled }} readinessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.broker.probes.readiness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.broker.probes.readiness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.broker.probes.readiness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.broker.probes.readiness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.broker.probes.readiness.timeoutSeconds | default .Values.probes.timeoutSeconds }} + httpGet: + path: {{ .Values.broker.probes.endpoint }} + port: {{ .Values.broker.service.port }} + {{- end }} + {{- if .Values.broker.probes.startupEnabled }} + startupProbe: + initialDelaySeconds: {{ .Values.broker.probes.startup.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.broker.probes.startup.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.broker.probes.startup.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.broker.probes.startup.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.broker.probes.startup.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: path: {{ .Values.broker.probes.endpoint }} port: {{ .Values.broker.service.port }} diff --git a/helm/pinot/templates/controller/poddisruptionbudget.yaml b/helm/pinot/templates/controller/poddisruptionbudget.yaml new file mode 100644 index 000000000000..d705c165c290 --- /dev/null +++ b/helm/pinot/templates/controller/poddisruptionbudget.yaml @@ -0,0 +1,35 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +{{- if .Values.controller.pdb.enabled }} +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: {{ include "pinot.controller.fullname" . }} +spec: + {{- if .Values.controller.pdb.minAvailable }} + minAvailable: {{ .Values.controller.pdb.minAvailable }} + {{- end }} + {{- if .Values.controller.pdb.maxUnavailable }} + maxUnavailable: {{ .Values.controller.pdb.maxUnavailable }} + {{- end }} + selector: + matchLabels: + {{- include "pinot.controllerMatchLabels" . | nindent 6 }} +{{- end }} diff --git a/helm/pinot/templates/controller/statefulset.yaml b/helm/pinot/templates/controller/statefulset.yaml index 8b652b301e57..0e81484f9fc6 100644 --- a/helm/pinot/templates/controller/statefulset.yaml +++ b/helm/pinot/templates/controller/statefulset.yaml @@ -37,7 +37,12 @@ spec: labels: {{- include "pinot.controllerLabels" . | nindent 8 }} annotations: -{{ toYaml .Values.controller.podAnnotations | indent 8 }} + {{- if .Values.controller.automaticReload.enabled }} + checksum/config: {{ include (print $.Template.BasePath "/controller/configmap.yaml") . | sha256sum }} + {{- end }} + {{- with .Values.controller.podAnnotations }} + {{- toYaml . | nindent 8 }} + {{- end }} spec: terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds }} serviceAccountName: {{ include "pinot.serviceAccountName" . }} @@ -77,16 +82,33 @@ spec: {{- end }} {{- if .Values.controller.probes.livenessEnabled }} livenessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.controller.probes.liveness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.controller.probes.liveness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.controller.probes.liveness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.controller.probes.liveness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.controller.probes.liveness.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: path: {{ .Values.controller.probes.endpoint }} port: {{ .Values.controller.service.port }} {{- end }} {{- if .Values.controller.probes.readinessEnabled }} readinessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.controller.probes.readiness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.controller.probes.readiness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.controller.probes.readiness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.controller.probes.readiness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.controller.probes.readiness.timeoutSeconds | default .Values.probes.timeoutSeconds }} + httpGet: + path: {{ .Values.controller.probes.endpoint }} + port: {{ .Values.controller.service.port }} + {{- end }} + {{- if .Values.controller.probes.startupEnabled }} + startupProbe: + initialDelaySeconds: {{ .Values.controller.probes.startup.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.controller.probes.startup.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.controller.probes.startup.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.controller.probes.startup.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.controller.probes.startup.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: path: {{ .Values.controller.probes.endpoint }} port: {{ .Values.controller.service.port }} diff --git a/helm/pinot/templates/minion-stateless/deployment.yaml b/helm/pinot/templates/minion-stateless/deployment.yaml index d68fd40275c7..4c4a5f4c7105 100644 --- a/helm/pinot/templates/minion-stateless/deployment.yaml +++ b/helm/pinot/templates/minion-stateless/deployment.yaml @@ -79,20 +79,37 @@ spec: {{- end }} {{- if .Values.minionStateless.probes.livenessEnabled }} livenessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.minionStateless.probes.liveness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.minionStateless.probes.liveness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.minionStateless.probes.liveness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.minionStateless.probes.liveness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.minionStateless.probes.liveness.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: path: {{ .Values.minionStateless.probes.endpoint }} port: {{ .Values.minionStateless.service.port }} {{- end }} {{- if .Values.minionStateless.probes.readinessEnabled }} readinessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.minionStateless.probes.readiness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.minionStateless.probes.readiness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.minionStateless.probes.readiness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.minionStateless.probes.readiness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.minionStateless.probes.readiness.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: path: {{ .Values.minionStateless.probes.endpoint }} port: {{ .Values.minionStateless.service.port }} {{- end }} + {{- if .Values.minionStateless.probes.startupEnabled }} + startupProbe: + initialDelaySeconds: {{ .Values.minionStateless.probes.startup.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.minionStateless.probes.startup.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.minionStateless.probes.startup.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.minionStateless.probes.startup.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.minionStateless.probes.startup.timeoutSeconds | default .Values.probes.timeoutSeconds }} + httpGet: + path: {{ .Values.minionStateless.probes.endpoint }} + port: {{ .Values.minionStateless.service.port }} + {{- end }} volumeMounts: - name: config mountPath: /var/pinot/minion/config diff --git a/helm/pinot/templates/minion/statefulset.yaml b/helm/pinot/templates/minion/statefulset.yaml index 443056dacfb5..6286395fd700 100644 --- a/helm/pinot/templates/minion/statefulset.yaml +++ b/helm/pinot/templates/minion/statefulset.yaml @@ -38,7 +38,12 @@ spec: labels: {{- include "pinot.minionLabels" . | nindent 8 }} annotations: -{{ toYaml .Values.minion.podAnnotations | indent 8 }} + {{- if .Values.minion.automaticReload.enabled }} + checksum/config: {{ include (print $.Template.BasePath "/minion/configmap.yaml") . | sha256sum}} + {{- end }} + {{- with .Values.minion.podAnnotations }} + {{- toYaml . | nindent 8 }} + {{- end }} spec: terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds }} serviceAccountName: {{ include "pinot.serviceAccountName" . }} @@ -83,16 +88,33 @@ spec: {{- end }} {{- if .Values.minion.probes.livenessEnabled }} livenessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.minion.probes.liveness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.minion.probes.liveness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.minion.probes.liveness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.minion.probes.liveness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.minion.probes.liveness.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: path: {{ .Values.minion.probes.endpoint }} port: {{ .Values.minion.service.port }} {{- end }} {{- if .Values.minion.probes.readinessEnabled }} readinessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.minion.probes.readiness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.minion.probes.readiness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.minion.probes.readiness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.minion.probes.readiness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.minion.probes.readiness.timeoutSeconds | default .Values.probes.timeoutSeconds }} + httpGet: + path: {{ .Values.minion.probes.endpoint }} + port: {{ .Values.minion.service.port }} + {{- end }} + {{- if .Values.minion.probes.startupEnabled }} + startupProbe: + initialDelaySeconds: {{ .Values.minion.probes.startup.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.minion.probes.startup.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.minion.probes.startup.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.minion.probes.startup.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.minion.probes.startup.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: path: {{ .Values.minion.probes.endpoint }} port: {{ .Values.minion.service.port }} diff --git a/helm/pinot/templates/server/poddisruptionbudget.yaml b/helm/pinot/templates/server/poddisruptionbudget.yaml new file mode 100644 index 000000000000..090f6b57bdd2 --- /dev/null +++ b/helm/pinot/templates/server/poddisruptionbudget.yaml @@ -0,0 +1,35 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +{{- if .Values.server.pdb.enabled }} +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: {{ include "pinot.server.fullname" . }} +spec: + {{- if .Values.server.pdb.minAvailable }} + minAvailable: {{ .Values.server.pdb.minAvailable }} + {{- end }} + {{- if .Values.server.pdb.maxUnavailable }} + maxUnavailable: {{ .Values.server.pdb.maxUnavailable }} + {{- end }} + selector: + matchLabels: + {{- include "pinot.serverMatchLabels" . | nindent 6 }} +{{- end }} diff --git a/helm/pinot/templates/server/statefulset.yaml b/helm/pinot/templates/server/statefulset.yaml index 071330f6f5a8..9811fd2bf48a 100644 --- a/helm/pinot/templates/server/statefulset.yaml +++ b/helm/pinot/templates/server/statefulset.yaml @@ -37,7 +37,12 @@ spec: labels: {{- include "pinot.serverLabels" . | nindent 8 }} annotations: -{{ toYaml .Values.server.podAnnotations | indent 8 }} + {{- if .Values.server.automaticReload.enabled }} + checksum/config: {{ include (print $.Template.BasePath "/server/configmap.yaml") . | sha256sum }} + {{- end }} + {{- with .Values.server.podAnnotations }} + {{- toYaml . | nindent 8 }} + {{- end }} spec: terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds }} serviceAccountName: {{ include "pinot.serviceAccountName" . }} @@ -85,26 +90,35 @@ spec: {{- end }} {{- if .Values.server.probes.livenessEnabled }} livenessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.server.probes.liveness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.server.probes.liveness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.server.probes.liveness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.server.probes.liveness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.server.probes.liveness.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: - {{- if and .Values.server.probes.livenessProbe .Values.server.probes.livenessProbe.endpoint}} - path: {{ .Values.server.probes.livenessProbe.endpoint }} - {{- else }} - path: {{ .Values.server.probes.endpoint }} - {{- end }} + path: {{ .Values.server.probes.liveness.endpoint | default .Values.server.probes.endpoint }} port: {{ .Values.server.service.adminPort }} {{- end }} {{- if .Values.server.probes.readinessEnabled }} readinessProbe: - initialDelaySeconds: {{ .Values.probes.initialDelaySeconds }} - periodSeconds: {{ .Values.probes.periodSeconds }} + initialDelaySeconds: {{ .Values.server.probes.readiness.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.server.probes.readiness.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.server.probes.readiness.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.server.probes.readiness.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.server.probes.readiness.timeoutSeconds | default .Values.probes.timeoutSeconds }} + httpGet: + path: {{ .Values.server.probes.readiness.endpoint | default .Values.server.probes.endpoint }} + port: {{ .Values.server.service.adminPort }} + {{- end }} + {{- if .Values.server.probes.startupEnabled }} + startupProbe: + initialDelaySeconds: {{ .Values.server.probes.startup.initialDelaySeconds | default .Values.probes.initialDelaySeconds }} + periodSeconds: {{ .Values.server.probes.startup.periodSeconds | default .Values.probes.periodSeconds }} + failureThreshold: {{ .Values.server.probes.startup.failureThreshold | default .Values.probes.failureThreshold }} + successThreshold: {{ .Values.server.probes.startup.successThreshold | default .Values.probes.successThreshold }} + timeoutSeconds: {{ .Values.server.probes.startup.timeoutSeconds | default .Values.probes.timeoutSeconds }} httpGet: - {{- if and .Values.server.probes.readinessProbe .Values.server.probes.readinessProbe.endpoint}} - path: {{ .Values.server.probes.readinessProbe.endpoint }} - {{- else }} - path: {{ .Values.server.probes.endpoint }} - {{- end }} + path: {{ .Values.server.probes.liveness.endpoint | default .Values.server.probes.endpoint }} port: {{ .Values.server.service.adminPort }} {{- end }} volumeMounts: diff --git a/helm/pinot/values.yaml b/helm/pinot/values.yaml index ef93913a173c..235e963e4f9d 100644 --- a/helm/pinot/values.yaml +++ b/helm/pinot/values.yaml @@ -49,9 +49,15 @@ securityContext: {} # runAsNonRoot: true # runAsUser: 1000 +# default values of the probes i.e. liveness and readiness. +# customization of values is present at the component level. probes: initialDelaySeconds: 60 periodSeconds: 10 + failureThreshold: 10 + # should be 1 for liveness and startup probe, as per K8s doc. + successThreshold: 1 + timeoutSeconds: 10 serviceAccount: # Specifies whether a service account should be created @@ -92,6 +98,27 @@ controller: endpoint: "/health" livenessEnabled: false readinessEnabled: false + startupEnabled: false + liveness: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + readiness: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + startup: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 persistence: enabled: true @@ -115,6 +142,11 @@ controller: log4j2ConfFile: /opt/pinot/etc/conf/pinot-controller-log4j2.xml pluginsDir: /opt/pinot/plugins + pdb: + enabled: false + minAvailable: "" + maxUnavailable: 50% + service: annotations: {} clusterIP: "None" @@ -164,6 +196,10 @@ controller: podAnnotations: {} + # set enabled as true, to automatically roll controller stateful set for configmap change + automaticReload: + enabled: false + updateStrategy: type: RollingUpdate @@ -214,11 +250,37 @@ broker: endpoint: "/health" livenessEnabled: true readinessEnabled: true + startupEnabled: false + liveness: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + readiness: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + startup: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 persistence: extraVolumes: [] extraVolumeMounts: [] + pdb: + enabled: false + minAvailable: "" + maxUnavailable: 50% + service: annotations: {} clusterIP: "None" @@ -269,6 +331,10 @@ broker: podAnnotations: {} + # set enabled as true, to automatically roll broker stateful set for configmap change + automaticReload: + enabled: false + updateStrategy: type: RollingUpdate @@ -312,6 +378,30 @@ server: endpoint: "/health" livenessEnabled: false readinessEnabled: false + startupEnabled: false + liveness: + endpoint: "/health/liveness" + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + readiness: + endpoint: "/health/readiness" + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + startup: + endpoint: "/health/liveness" + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 dataDir: /var/pinot/server/data/index segmentTarDir: /var/pinot/server/data/segment @@ -331,6 +421,11 @@ server: log4j2ConfFile: /opt/pinot/etc/conf/pinot-server-log4j2.xml pluginsDir: /opt/pinot/plugins + pdb: + enabled: false + minAvailable: "" + maxUnavailable: 1 + service: annotations: {} clusterIP: "" @@ -362,6 +457,10 @@ server: podAnnotations: {} + # set enabled as true, to automatically roll server stateful set for configmap change + automaticReload: + enabled: false + updateStrategy: type: RollingUpdate @@ -407,6 +506,27 @@ minion: endpoint: "/health" livenessEnabled: true readinessEnabled: true + startupEnabled: false + liveness: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + readiness: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + startup: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 dataDir: /var/pinot/minion/data jvmOpts: "-XX:ActiveProcessorCount=2 -Xms256M -Xmx1G -XX:+UseG1GC -XX:MaxGCPauseMillis=200 -Xlog:gc*:file=/opt/pinot/gc-pinot-minion.log" @@ -452,6 +572,9 @@ minion: podAnnotations: {} + automaticReload: + enabled: false + updateStrategy: type: RollingUpdate @@ -494,6 +617,27 @@ minionStateless: endpoint: "/health" livenessEnabled: true readinessEnabled: true + startupEnabled: true + liveness: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + readiness: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 + + startup: + initialDelaySeconds: 60 + failureThreshold: 10 + timeoutSeconds: 10 + successThreshold: 1 + periodSeconds: 10 dataDir: /var/pinot/minion/data jvmOpts: "-XX:ActiveProcessorCount=2 -Xms256M -Xmx1G -XX:+UseG1GC -XX:MaxGCPauseMillis=200 -Xlog:gc*:file=/opt/pinot/gc-pinot-minion.log" diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/api/AccessControl.java b/pinot-broker/src/main/java/org/apache/pinot/broker/api/AccessControl.java index 85c1bf74ae7e..f18fcc7f4129 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/api/AccessControl.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/api/AccessControl.java @@ -23,6 +23,9 @@ import org.apache.pinot.core.auth.FineGrainedAccessControl; import org.apache.pinot.spi.annotations.InterfaceAudience; import org.apache.pinot.spi.annotations.InterfaceStability; +import org.apache.pinot.spi.auth.AuthorizationResult; +import org.apache.pinot.spi.auth.BasicAuthorizationResultImpl; +import org.apache.pinot.spi.auth.TableAuthorizationResult; @InterfaceAudience.Public @@ -31,32 +34,89 @@ public interface AccessControl extends FineGrainedAccessControl { /** * First-step access control when processing broker requests. Decides whether request is allowed to acquire resources * for further processing. Request may still be rejected at table-level later on. - * + * The default implementation is kept to have backward compatibility with the existing implementations * @param requesterIdentity requester identity * * @return {@code true} if authorized, {@code false} otherwise */ + @Deprecated default boolean hasAccess(RequesterIdentity requesterIdentity) { return true; } /** - * Fine-grained access control on parsed broker request. May check table, column, permissions, etc. + * First-step access control when processing broker requests. Decides whether request is allowed to acquire resources + * for further processing. Request may still be rejected at table-level later on. + * The default implementation returns a {@link BasicAuthorizationResultImpl} with the result of the hasAccess() of + * the implementation * * @param requesterIdentity requester identity + * + * @return {@code AuthorizationResult} with the result of the access control check + */ + default AuthorizationResult authorize(RequesterIdentity requesterIdentity) { + return new BasicAuthorizationResultImpl(hasAccess(requesterIdentity)); + } + + /** + * Fine-grained access control on parsed broker request. May check table, column, permissions, etc. + * The default implementation is kept to have backward compatibility with the existing implementations + * @param requesterIdentity requester identity * @param brokerRequest broker request (incl query) * * @return {@code true} if authorized, {@code false} otherwise */ - boolean hasAccess(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest); + @Deprecated + default boolean hasAccess(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { + throw new UnsupportedOperationException( + "Both hasAccess() and authorize() are not implemented . Do implement authorize() method for new " + + "implementations."); + } + + /** + * Verify access control on parsed broker request. May check table, column, permissions, etc. + * The default implementation returns a {@link BasicAuthorizationResultImpl} with the result of the hasAccess() of + * the implementation + * + * @param requesterIdentity requester identity + * @param brokerRequest broker request (incl query) + * + * @return {@code AuthorizationResult} with the result of the access control check + */ + default AuthorizationResult authorize(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { + return new BasicAuthorizationResultImpl(hasAccess(requesterIdentity, brokerRequest)); + } /** * Fine-grained access control on pinot tables. + * The default implementation is kept to have backward compatibility with the existing implementations * * @param requesterIdentity requester identity * @param tables Set of pinot tables used in the query. Table name can be with or without tableType. * * @return {@code true} if authorized, {@code false} otherwise */ - boolean hasAccess(RequesterIdentity requesterIdentity, Set tables); + @Deprecated + default boolean hasAccess(RequesterIdentity requesterIdentity, Set tables) { + throw new UnsupportedOperationException( + "Both hasAccess() and authorize() are not implemented . Do implement authorize() method for new " + + "implementations."); + } + + /** + * Verify access control on pinot tables. + * The default implementation returns a {@link TableAuthorizationResult} with the result of the hasAccess() of the + * implementation + * + * @param requesterIdentity requester identity + * @param tables Set of pinot tables used in the query. Table name can be with or without tableType. + * + * @return {@code TableAuthorizationResult} with the result of the access control check + */ + default TableAuthorizationResult authorize(RequesterIdentity requesterIdentity, Set tables) { + // Taking all tables when hasAccess Failed , to not break existing implementations + // It will say all tables names failed AuthZ even only some failed AuthZ - which is same as just boolean output + return hasAccess(requesterIdentity, tables) ? TableAuthorizationResult.success() + : new TableAuthorizationResult(tables); + } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerLogger.java b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerLogger.java index e2a8dc71ebed..593e05d8daca 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerLogger.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerLogger.java @@ -68,7 +68,7 @@ public class PinotBrokerLogger { @Produces(MediaType.APPLICATION_JSON) @ApiOperation(value = "Get all the loggers", notes = "Return all the logger names") public List getLoggers() { - return LoggerUtils.getAllLoggers(); + return LoggerUtils.getAllConfiguredLoggers(); } @GET diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java index 0bdec44a090e..91a35d61302c 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java @@ -34,6 +34,7 @@ import io.swagger.annotations.SecurityDefinition; import io.swagger.annotations.SwaggerDefinition; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.concurrent.Executor; import javax.inject.Inject; @@ -60,6 +61,7 @@ import org.apache.pinot.common.metrics.BrokerMetrics; import org.apache.pinot.common.response.BrokerResponse; import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.response.broker.QueryProcessingException; import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.core.auth.Actions; import org.apache.pinot.core.auth.Authorize; @@ -114,7 +116,6 @@ public class PinotClientRequest { @ManualAuthorization public void processSqlQueryGet(@ApiParam(value = "Query", required = true) @QueryParam("sql") String query, @ApiParam(value = "Trace enabled") @QueryParam(Request.TRACE) String traceEnabled, - @ApiParam(value = "Debug options") @QueryParam(Request.DEBUG_OPTIONS) String debugOptions, @Suspended AsyncResponse asyncResponse, @Context org.glassfish.grizzly.http.server.Request requestContext, @Context HttpHeaders httpHeaders) { try { @@ -123,9 +124,6 @@ public void processSqlQueryGet(@ApiParam(value = "Query", required = true) @Quer if (traceEnabled != null) { requestJson.put(Request.TRACE, traceEnabled); } - if (debugOptions != null) { - requestJson.put(Request.DEBUG_OPTIONS, debugOptions); - } BrokerResponse brokerResponse = executeSqlQuery(requestJson, makeHttpIdentity(requestContext), true, httpHeaders); asyncResponse.resume(getPinotQueryResponse(brokerResponse)); } catch (WebApplicationException wae) { @@ -299,6 +297,7 @@ private BrokerResponse executeSqlQuery(ObjectNode sqlRequestJson, HttpRequesterI private BrokerResponse executeSqlQuery(ObjectNode sqlRequestJson, HttpRequesterIdentity httpRequesterIdentity, boolean onlyDql, HttpHeaders httpHeaders, boolean forceUseMultiStage) throws Exception { + long requestArrivalTimeMs = System.currentTimeMillis(); SqlNodeAndOptions sqlNodeAndOptions; try { sqlNodeAndOptions = RequestUtils.parseQuery(sqlRequestJson.get(Request.SQL).asText(), sqlRequestJson); @@ -315,9 +314,10 @@ private BrokerResponse executeSqlQuery(ObjectNode sqlRequestJson, HttpRequesterI } switch (sqlType) { case DQL: - try (RequestScope requestStatistics = Tracing.getTracer().createRequestScope()) { - return _requestHandler.handleRequest(sqlRequestJson, sqlNodeAndOptions, httpRequesterIdentity, - requestStatistics, httpHeaders); + try (RequestScope requestContext = Tracing.getTracer().createRequestScope()) { + requestContext.setRequestArrivalTimeMillis(requestArrivalTimeMs); + return _requestHandler.handleRequest(sqlRequestJson, sqlNodeAndOptions, httpRequesterIdentity, requestContext, + httpHeaders); } catch (Exception e) { LOGGER.error("Error handling DQL request:\n{}\nException: {}", sqlRequestJson, QueryException.getTruncatedStackTrace(e)); @@ -365,10 +365,10 @@ private static HttpRequesterIdentity makeHttpIdentity(org.glassfish.grizzly.http static Response getPinotQueryResponse(BrokerResponse brokerResponse) throws Exception { int queryErrorCodeHeaderValue = -1; // default value of the header. - - if (brokerResponse.getExceptionsSize() != 0) { + List exceptions = brokerResponse.getExceptions(); + if (!exceptions.isEmpty()) { // set the header value as first exception error code value. - queryErrorCodeHeaderValue = brokerResponse.getProcessingExceptions().get(0).getErrorCode(); + queryErrorCodeHeaderValue = exceptions.get(0).getErrorCode(); } // returning the Response with OK status and header value. diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/AllowAllAccessControlFactory.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/AllowAllAccessControlFactory.java index 1e5a888a66f2..d0eb2ed65c1c 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/AllowAllAccessControlFactory.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/AllowAllAccessControlFactory.java @@ -22,6 +22,9 @@ import org.apache.pinot.broker.api.AccessControl; import org.apache.pinot.broker.api.RequesterIdentity; import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.spi.auth.AuthorizationResult; +import org.apache.pinot.spi.auth.BasicAuthorizationResultImpl; +import org.apache.pinot.spi.auth.TableAuthorizationResult; import org.apache.pinot.spi.env.PinotConfiguration; @@ -41,13 +44,13 @@ public AccessControl create() { private static class AllowAllAccessControl implements AccessControl { @Override - public boolean hasAccess(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { - return true; + public AuthorizationResult authorize(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { + return BasicAuthorizationResultImpl.success(); } @Override - public boolean hasAccess(RequesterIdentity requesterIdentity, Set tables) { - return true; + public TableAuthorizationResult authorize(RequesterIdentity requesterIdentity, Set tables) { + return TableAuthorizationResult.success(); } } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/AuthenticationFilter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/AuthenticationFilter.java index 787a9da685f0..a61a860356e7 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/AuthenticationFilter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/AuthenticationFilter.java @@ -34,10 +34,12 @@ import javax.ws.rs.core.HttpHeaders; import javax.ws.rs.core.Response; import javax.ws.rs.core.UriInfo; +import org.apache.commons.lang.StringUtils; import org.apache.pinot.broker.api.AccessControl; import org.apache.pinot.broker.api.HttpRequesterIdentity; import org.apache.pinot.core.auth.FineGrainedAuthUtils; import org.apache.pinot.core.auth.ManualAuthorization; +import org.apache.pinot.spi.auth.AuthorizationResult; import org.glassfish.grizzly.http.server.Request; /** @@ -82,9 +84,15 @@ public void filter(ContainerRequestContext requestContext) HttpRequesterIdentity httpRequestIdentity = HttpRequesterIdentity.fromRequest(request); - // default authorization handling - if (!accessControl.hasAccess(httpRequestIdentity)) { - throw new WebApplicationException("Failed access check for " + httpRequestIdentity.getEndpointUrl(), + AuthorizationResult authorizationResult = accessControl.authorize(httpRequestIdentity); + + if (!authorizationResult.hasAccess()) { + String failureMessage = authorizationResult.getFailureMessage(); + if (StringUtils.isNotBlank(failureMessage)) { + failureMessage = "Reason: " + failureMessage; + } + throw new WebApplicationException( + "Failed access check for " + httpRequestIdentity.getEndpointUrl() + "." + failureMessage, Response.Status.FORBIDDEN); } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BasicAuthAccessControlFactory.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BasicAuthAccessControlFactory.java index 2de2541d825e..18c1932a521e 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BasicAuthAccessControlFactory.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BasicAuthAccessControlFactory.java @@ -20,6 +20,7 @@ import com.google.common.base.Preconditions; import java.util.Collection; +import java.util.HashSet; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -32,6 +33,8 @@ import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.core.auth.BasicAuthPrincipal; import org.apache.pinot.core.auth.BasicAuthUtils; +import org.apache.pinot.spi.auth.AuthorizationResult; +import org.apache.pinot.spi.auth.TableAuthorizationResult; import org.apache.pinot.spi.env.PinotConfiguration; @@ -78,12 +81,12 @@ public BasicAuthAccessControl(Collection principals) { } @Override - public boolean hasAccess(RequesterIdentity requesterIdentity) { - return hasAccess(requesterIdentity, (BrokerRequest) null); + public AuthorizationResult authorize(RequesterIdentity requesterIdentity) { + return authorize(requesterIdentity, (BrokerRequest) null); } @Override - public boolean hasAccess(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { + public AuthorizationResult authorize(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { Optional principalOpt = getPrincipalOpt(requesterIdentity); if (!principalOpt.isPresent()) { @@ -94,14 +97,22 @@ public boolean hasAccess(RequesterIdentity requesterIdentity, BrokerRequest brok if (brokerRequest == null || !brokerRequest.isSetQuerySource() || !brokerRequest.getQuerySource() .isSetTableName()) { // no table restrictions? accept - return true; + return TableAuthorizationResult.success(); } - return principal.hasTable(brokerRequest.getQuerySource().getTableName()); + Set failedTables = new HashSet<>(); + + if (!principal.hasTable(brokerRequest.getQuerySource().getTableName())) { + failedTables.add(brokerRequest.getQuerySource().getTableName()); + } + if (failedTables.isEmpty()) { + return TableAuthorizationResult.success(); + } + return new TableAuthorizationResult(failedTables); } @Override - public boolean hasAccess(RequesterIdentity requesterIdentity, Set tables) { + public TableAuthorizationResult authorize(RequesterIdentity requesterIdentity, Set tables) { Optional principalOpt = getPrincipalOpt(requesterIdentity); if (!principalOpt.isPresent()) { @@ -109,17 +120,19 @@ public boolean hasAccess(RequesterIdentity requesterIdentity, Set tables } if (tables == null || tables.isEmpty()) { - return true; + return TableAuthorizationResult.success(); } - BasicAuthPrincipal principal = principalOpt.get(); + Set failedTables = new HashSet<>(); for (String table : tables) { if (!principal.hasTable(table)) { - return false; + failedTables.add(table); } } - - return true; + if (failedTables.isEmpty()) { + return TableAuthorizationResult.success(); + } + return new TableAuthorizationResult(failedTables); } private Optional getPrincipalOpt(RequesterIdentity requesterIdentity) { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerAdminApiApplication.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerAdminApiApplication.java index 628e6956f256..8dedc8c72b3b 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerAdminApiApplication.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerAdminApiApplication.java @@ -19,10 +19,8 @@ package org.apache.pinot.broker.broker; import com.google.common.util.concurrent.ThreadFactoryBuilder; -import io.swagger.jaxrs.config.BeanConfig; +import io.swagger.jaxrs.listing.SwaggerSerializers; import java.io.IOException; -import java.net.URL; -import java.net.URLClassLoader; import java.time.Instant; import java.util.List; import java.util.concurrent.Executor; @@ -35,6 +33,8 @@ import org.apache.pinot.broker.requesthandler.BrokerRequestHandler; import org.apache.pinot.broker.routing.BrokerRoutingManager; import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.swagger.SwaggerApiListingResource; +import org.apache.pinot.common.swagger.SwaggerSetupUtils; import org.apache.pinot.common.utils.log.DummyLogFileServer; import org.apache.pinot.common.utils.log.LocalLogFileServer; import org.apache.pinot.common.utils.log.LogFileServer; @@ -46,8 +46,6 @@ import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.PinotReflectionUtils; -import org.glassfish.grizzly.http.server.CLStaticHttpHandler; -import org.glassfish.grizzly.http.server.HttpHandler; import org.glassfish.grizzly.http.server.HttpServer; import org.glassfish.hk2.utilities.binding.AbstractBinder; import org.glassfish.jersey.jackson.JacksonFeature; @@ -88,9 +86,8 @@ public BrokerAdminApiApplication(BrokerRoutingManager routingManager, BrokerRequ _executorService = Executors.newCachedThreadPool(new ThreadFactoryBuilder().setNameFormat("async-task-thread-%d").build()); PoolingHttpClientConnectionManager connMgr = new PoolingHttpClientConnectionManager(); - int timeoutMs = (int) brokerConf - .getProperty(CommonConstants.Broker.CONFIG_OF_BROKER_TIMEOUT_MS, - CommonConstants.Broker.DEFAULT_BROKER_TIMEOUT_MS); + int timeoutMs = (int) brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_BROKER_TIMEOUT_MS, + CommonConstants.Broker.DEFAULT_BROKER_TIMEOUT_MS); connMgr.setDefaultSocketConfig(SocketConfig.custom().setSoTimeout(timeoutMs).build()); Instant startTime = Instant.now(); register(new AbstractBinder() { @@ -115,15 +112,15 @@ protected void configure() { bind(startTime).named(BrokerAdminApiApplication.START_TIME); } }); - boolean enableBoundedJerseyThreadPoolExecutor = brokerConf - .getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_BOUNDED_JERSEY_THREADPOOL_EXECUTOR, + boolean enableBoundedJerseyThreadPoolExecutor = + brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_BOUNDED_JERSEY_THREADPOOL_EXECUTOR, CommonConstants.Broker.DEFAULT_ENABLE_BOUNDED_JERSEY_THREADPOOL_EXECUTOR); if (enableBoundedJerseyThreadPoolExecutor) { register(buildBrokerManagedAsyncExecutorProvider(brokerConf, brokerMetrics)); } register(JacksonFeature.class); - registerClasses(io.swagger.jaxrs.listing.ApiListingResource.class); - registerClasses(io.swagger.jaxrs.listing.SwaggerSerializers.class); + register(SwaggerApiListingResource.class); + register(SwaggerSerializers.class); register(AuthenticationFilter.class); } @@ -137,45 +134,20 @@ public void start(List listenerConfigs) { } if (_swaggerBrokerEnabled) { - PinotReflectionUtils.runWithLock(this::setupSwagger); + PinotReflectionUtils.runWithLock(() -> + SwaggerSetupUtils.setupSwagger("Broker", _brokerResourcePackages, _useHttps, "/", _httpServer)); } else { LOGGER.info("Hiding Swagger UI for Broker, by {}", CommonConstants.Broker.CONFIG_OF_SWAGGER_BROKER_ENABLED); } } - private void setupSwagger() { - BeanConfig beanConfig = new BeanConfig(); - beanConfig.setTitle("Pinot Broker API"); - beanConfig.setDescription("APIs for accessing Pinot broker information"); - beanConfig.setContact("https://github.com/apache/pinot"); - beanConfig.setVersion("1.0"); - beanConfig.setExpandSuperTypes(false); - if (_useHttps) { - beanConfig.setSchemes(new String[]{CommonConstants.HTTPS_PROTOCOL}); - } else { - beanConfig.setSchemes(new String[]{CommonConstants.HTTP_PROTOCOL, CommonConstants.HTTPS_PROTOCOL}); - } - beanConfig.setBasePath("/"); - beanConfig.setResourcePackage(_brokerResourcePackages); - beanConfig.setScan(true); - - HttpHandler httpHandler = new CLStaticHttpHandler(BrokerAdminApiApplication.class.getClassLoader(), "/api/"); - // map both /api and /help to swagger docs. /api because it looks nice. /help for backward compatibility - _httpServer.getServerConfiguration().addHttpHandler(httpHandler, "/api/", "/help/"); - - URL swaggerDistLocation = - BrokerAdminApiApplication.class.getClassLoader().getResource(CommonConstants.CONFIG_OF_SWAGGER_RESOURCES_PATH); - CLStaticHttpHandler swaggerDist = new CLStaticHttpHandler(new URLClassLoader(new URL[]{swaggerDistLocation})); - _httpServer.getServerConfiguration().addHttpHandler(swaggerDist, "/swaggerui-dist/"); - } - private BrokerManagedAsyncExecutorProvider buildBrokerManagedAsyncExecutorProvider(PinotConfiguration brokerConf, BrokerMetrics brokerMetrics) { - int corePoolSize = brokerConf - .getProperty(CommonConstants.Broker.CONFIG_OF_JERSEY_THREADPOOL_EXECUTOR_CORE_POOL_SIZE, + int corePoolSize = + brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_JERSEY_THREADPOOL_EXECUTOR_CORE_POOL_SIZE, CommonConstants.Broker.DEFAULT_JERSEY_THREADPOOL_EXECUTOR_CORE_POOL_SIZE); - int maximumPoolSize = brokerConf - .getProperty(CommonConstants.Broker.CONFIG_OF_JERSEY_THREADPOOL_EXECUTOR_MAX_POOL_SIZE, + int maximumPoolSize = + brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_JERSEY_THREADPOOL_EXECUTOR_MAX_POOL_SIZE, CommonConstants.Broker.DEFAULT_JERSEY_THREADPOOL_EXECUTOR_MAX_POOL_SIZE); int queueSize = brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_JERSEY_THREADPOOL_EXECUTOR_QUEUE_SIZE, CommonConstants.Broker.DEFAULT_JERSEY_THREADPOOL_EXECUTOR_QUEUE_SIZE); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerManagedAsyncExecutorProvider.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerManagedAsyncExecutorProvider.java index c85f421f4c15..8d6823be84e5 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerManagedAsyncExecutorProvider.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerManagedAsyncExecutorProvider.java @@ -98,7 +98,7 @@ public BrokerThreadPoolRejectExecutionHandler(BrokerMetrics brokerMetrics) { @Override public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) { _brokerMetrics.addMeteredGlobalValue(BrokerMeter.QUERY_REJECTED_EXCEPTIONS, 1L); - LOGGER.error("Task " + r + " rejected from " + executor); + LOGGER.error("Task {} rejected from {}", r, executor); throw new ServiceUnavailableException(Response.status(Response.Status.SERVICE_UNAVAILABLE).entity( "Pinot Broker thread pool can not accommodate more requests now. " + "Request is rejected from " + executor) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/ZkBasicAuthAccessControlFactory.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/ZkBasicAuthAccessControlFactory.java index 557484cc65f5..d8434a23a4bb 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/ZkBasicAuthAccessControlFactory.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/ZkBasicAuthAccessControlFactory.java @@ -21,6 +21,7 @@ import com.google.common.base.Preconditions; import java.util.Collection; import java.util.Collections; +import java.util.HashSet; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -38,6 +39,8 @@ import org.apache.pinot.core.auth.BasicAuthPrincipal; import org.apache.pinot.core.auth.BasicAuthUtils; import org.apache.pinot.core.auth.ZkBasicAuthPrincipal; +import org.apache.pinot.spi.auth.AuthorizationResult; +import org.apache.pinot.spi.auth.TableAuthorizationResult; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.utils.builder.TableNameBuilder; @@ -82,39 +85,42 @@ public BasicAuthAccessControl(AccessControlUserCache userCache) { } @Override - public boolean hasAccess(RequesterIdentity requesterIdentity) { - return hasAccess(requesterIdentity, (BrokerRequest) null); + public AuthorizationResult authorize(RequesterIdentity requesterIdentity) { + return authorize(requesterIdentity, (BrokerRequest) null); } @Override - public boolean hasAccess(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { + public AuthorizationResult authorize(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { if (brokerRequest == null || !brokerRequest.isSetQuerySource() || !brokerRequest.getQuerySource() .isSetTableName()) { // no table restrictions? accept - return true; + return TableAuthorizationResult.success(); } - return hasAccess(requesterIdentity, Collections.singleton(brokerRequest.getQuerySource().getTableName())); + return authorize(requesterIdentity, Collections.singleton(brokerRequest.getQuerySource().getTableName())); } @Override - public boolean hasAccess(RequesterIdentity requesterIdentity, Set tables) { + public TableAuthorizationResult authorize(RequesterIdentity requesterIdentity, Set tables) { Optional principalOpt = getPrincipalAuth(requesterIdentity); if (!principalOpt.isPresent()) { throw new NotAuthorizedException("Basic"); } if (tables == null || tables.isEmpty()) { - return true; + return TableAuthorizationResult.success(); } ZkBasicAuthPrincipal principal = principalOpt.get(); + Set failedTables = new HashSet<>(); for (String table : tables) { if (!principal.hasTable(TableNameBuilder.extractRawTableName(table))) { - return false; + failedTables.add(table); } } - - return true; + if (failedTables.isEmpty()) { + return TableAuthorizationResult.success(); + } + return new TableAuthorizationResult(failedTables); } private Optional getPrincipalAuth(RequesterIdentity requesterIdentity) { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java index 3bf78a56d64d..b16489b35bb3 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java @@ -42,6 +42,7 @@ import org.apache.pinot.broker.broker.AccessControlFactory; import org.apache.pinot.broker.broker.BrokerAdminApiApplication; import org.apache.pinot.broker.queryquota.HelixExternalViewBasedQueryQuotaManager; +import org.apache.pinot.broker.requesthandler.BaseSingleStageBrokerRequestHandler; import org.apache.pinot.broker.requesthandler.BrokerRequestHandler; import org.apache.pinot.broker.requesthandler.BrokerRequestHandlerDelegate; import org.apache.pinot.broker.requesthandler.GrpcBrokerRequestHandler; @@ -72,8 +73,7 @@ import org.apache.pinot.core.util.ListenerConfigUtil; import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; import org.apache.pinot.spi.env.PinotConfiguration; -import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListener; -import org.apache.pinot.spi.eventlistener.query.PinotBrokerQueryEventListenerFactory; +import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; import org.apache.pinot.spi.metrics.PinotMetricUtils; import org.apache.pinot.spi.metrics.PinotMetricsRegistry; import org.apache.pinot.spi.services.ServiceRole; @@ -128,7 +128,6 @@ public abstract class BaseBrokerStarter implements ServiceStartable { protected HelixManager _participantHelixManager; // Handles the server routing stats. protected ServerRoutingStatsManager _serverRoutingStatsManager; - protected BrokerQueryEventListener _brokerQueryEventListener; @Override public void init(PinotConfiguration brokerConf) @@ -139,8 +138,8 @@ public void init(PinotConfiguration brokerConf) _clusterName = brokerConf.getProperty(Helix.CONFIG_OF_CLUSTER_NAME); ServiceStartableUtils.applyClusterConfig(_brokerConf, _zkServers, _clusterName, ServiceRole.BROKER); - PinotInsecureMode.setPinotInInsecureMode( - Boolean.valueOf(_brokerConf.getProperty(CommonConstants.CONFIG_OF_PINOT_INSECURE_MODE, + PinotInsecureMode.setPinotInInsecureMode(Boolean.valueOf( + _brokerConf.getProperty(CommonConstants.CONFIG_OF_PINOT_INSECURE_MODE, CommonConstants.DEFAULT_PINOT_INSECURE_MODE))); if (_brokerConf.getProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, @@ -277,8 +276,7 @@ public void start() final PinotConfiguration factoryConf = _brokerConf.subset(Broker.ACCESS_CONTROL_CONFIG_PREFIX); // Adding cluster name to the config so that it can be used by the AccessControlFactory factoryConf.setProperty(Helix.CONFIG_OF_CLUSTER_NAME, _brokerConf.getProperty(Helix.CONFIG_OF_CLUSTER_NAME)); - _accessControlFactory = - AccessControlFactory.loadFactory(factoryConf, _propertyStore); + _accessControlFactory = AccessControlFactory.loadFactory(factoryConf, _propertyStore); HelixExternalViewBasedQueryQuotaManager queryQuotaManager = new HelixExternalViewBasedQueryQuotaManager(_brokerMetrics, _instanceId); queryQuotaManager.init(_spectatorHelixManager); @@ -292,49 +290,42 @@ public void start() boolean caseInsensitive = _brokerConf.getProperty(Helix.ENABLE_CASE_INSENSITIVE_KEY, Helix.DEFAULT_ENABLE_CASE_INSENSITIVE); TableCache tableCache = new TableCache(_propertyStore, caseInsensitive); - // Configure TLS for netty connection to server - TlsConfig tlsDefaults = TlsUtils.extractTlsConfig(_brokerConf, Broker.BROKER_TLS_PREFIX); - NettyConfig nettyDefaults = NettyConfig.extractNettyConfig(_brokerConf, Broker.BROKER_NETTY_PREFIX); LOGGER.info("Initializing Broker Event Listener Factory"); - _brokerQueryEventListener = PinotBrokerQueryEventListenerFactory.getBrokerQueryEventListener( - _brokerConf.subset(Broker.EVENT_LISTENER_CONFIG_PREFIX)); + BrokerQueryEventListenerFactory.init(_brokerConf.subset(Broker.EVENT_LISTENER_CONFIG_PREFIX)); // Create Broker request handler. String brokerId = _brokerConf.getProperty(Broker.CONFIG_OF_BROKER_ID, getDefaultBrokerId()); String brokerRequestHandlerType = _brokerConf.getProperty(Broker.BROKER_REQUEST_HANDLER_TYPE, Broker.DEFAULT_BROKER_REQUEST_HANDLER_TYPE); - BrokerRequestHandler singleStageBrokerRequestHandler = null; + BaseSingleStageBrokerRequestHandler singleStageBrokerRequestHandler; if (brokerRequestHandlerType.equalsIgnoreCase(Broker.GRPC_BROKER_REQUEST_HANDLER_TYPE)) { singleStageBrokerRequestHandler = new GrpcBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory, queryQuotaManager, - tableCache, _brokerMetrics, null, _brokerQueryEventListener); - } else { // default request handler type, e.g. netty + tableCache); + } else { + // Default request handler type, i.e. netty + NettyConfig nettyDefaults = NettyConfig.extractNettyConfig(_brokerConf, Broker.BROKER_NETTY_PREFIX); + // Configure TLS for netty connection to server + TlsConfig tlsDefaults = null; if (_brokerConf.getProperty(Broker.BROKER_NETTYTLS_ENABLED, false)) { - singleStageBrokerRequestHandler = - new SingleConnectionBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory, - queryQuotaManager, tableCache, _brokerMetrics, nettyDefaults, tlsDefaults, _serverRoutingStatsManager, - _brokerQueryEventListener); - } else { - singleStageBrokerRequestHandler = - new SingleConnectionBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory, - queryQuotaManager, tableCache, _brokerMetrics, nettyDefaults, null, _serverRoutingStatsManager, - _brokerQueryEventListener); + tlsDefaults = TlsUtils.extractTlsConfig(_brokerConf, Broker.BROKER_TLS_PREFIX); } + singleStageBrokerRequestHandler = + new SingleConnectionBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory, + queryQuotaManager, tableCache, nettyDefaults, tlsDefaults, _serverRoutingStatsManager); } - - BrokerRequestHandler multiStageBrokerRequestHandler = null; + MultiStageBrokerRequestHandler multiStageBrokerRequestHandler = null; if (_brokerConf.getProperty(Helix.CONFIG_OF_MULTI_STAGE_ENGINE_ENABLED, Helix.DEFAULT_MULTI_STAGE_ENGINE_ENABLED)) { // multi-stage request handler uses both Netty and GRPC ports. // worker requires both the "Netty port" for protocol transport; and "GRPC port" for mailbox transport. // TODO: decouple protocol and engine selection. multiStageBrokerRequestHandler = new MultiStageBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory, - queryQuotaManager, tableCache, _brokerMetrics, _brokerQueryEventListener); + queryQuotaManager, tableCache); } - - _brokerRequestHandler = new BrokerRequestHandlerDelegate(brokerId, singleStageBrokerRequestHandler, - multiStageBrokerRequestHandler, _brokerMetrics); + _brokerRequestHandler = + new BrokerRequestHandlerDelegate(singleStageBrokerRequestHandler, multiStageBrokerRequestHandler); _brokerRequestHandler.start(); // Enable/disable thread CPU time measurement through instance config. @@ -345,8 +336,8 @@ public void start() ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled( _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT, CommonConstants.Broker.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT)); - Tracing.ThreadAccountantOps - .initializeThreadAccountant(_brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId); + Tracing.ThreadAccountantOps.initializeThreadAccountant( + _brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId); String controllerUrl = _brokerConf.getProperty(Broker.CONTROLLER_URL); if (controllerUrl != null) { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/querylog/QueryLogger.java b/pinot-broker/src/main/java/org/apache/pinot/broker/querylog/QueryLogger.java index 28564cc1c67a..6b34b46ae9d6 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/querylog/QueryLogger.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/querylog/QueryLogger.java @@ -25,7 +25,7 @@ import javax.annotation.Nullable; import org.apache.commons.lang3.StringUtils; import org.apache.pinot.broker.api.RequesterIdentity; -import org.apache.pinot.broker.requesthandler.BaseBrokerRequestHandler; +import org.apache.pinot.broker.requesthandler.BaseSingleStageBrokerRequestHandler.ServerStats; import org.apache.pinot.common.response.BrokerResponse; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.trace.RequestContext; @@ -89,7 +89,8 @@ public void log(QueryLogParams params) { } // always log the query last - don't add this to the QueryLogEntry enum - queryLogBuilder.append("query=").append(StringUtils.substring(params._query, 0, _maxQueryLengthToLog)); + queryLogBuilder.append("query=") + .append(StringUtils.substring(params._requestContext.getQuery(), 0, _maxQueryLengthToLog)); _logger.info(queryLogBuilder.toString()); if (_droppedLogRateLimiter.tryAcquire()) { @@ -112,35 +113,26 @@ public double getLogRateLimit() { } private boolean shouldForceLog(QueryLogParams params) { - return params._response.isNumGroupsLimitReached() || params._response.getExceptionsSize() > 0 - || params._timeUsedMs > TimeUnit.SECONDS.toMillis(1); + return params._response.isPartialResult() || params._response.getTimeUsedMs() > TimeUnit.SECONDS.toMillis(1); } public static class QueryLogParams { - final long _requestId; - final String _query; - final RequestContext _requestContext; - final String _table; - final int _numUnavailableSegments; + private final RequestContext _requestContext; + private final String _table; + private final BrokerResponse _response; @Nullable - final BaseBrokerRequestHandler.ServerStats _serverStats; - final BrokerResponse _response; - final long _timeUsedMs; + private final RequesterIdentity _identity; @Nullable - final RequesterIdentity _requester; + private final ServerStats _serverStats; - public QueryLogParams(long requestId, String query, RequestContext requestContext, String table, - int numUnavailableSegments, @Nullable BaseBrokerRequestHandler.ServerStats serverStats, BrokerResponse response, - long timeUsedMs, @Nullable RequesterIdentity requester) { - _requestId = requestId; - _query = query; - _table = table; - _timeUsedMs = timeUsedMs; + public QueryLogParams(RequestContext requestContext, String table, BrokerResponse response, + @Nullable RequesterIdentity identity, @Nullable ServerStats serverStats) { _requestContext = requestContext; - _requester = requester; + // NOTE: Passing table name separately because table name within request context is always raw table name. + _table = table; _response = response; + _identity = identity; _serverStats = serverStats; - _numUnavailableSegments = numUnavailableSegments; } } @@ -152,7 +144,8 @@ private enum QueryLogEntry { REQUEST_ID("requestId") { @Override void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) { - builder.append(params._requestId); + // NOTE: At this moment, request ID is not available at response yet. + builder.append(params._requestContext.getRequestId()); } }, TABLE("table") { @@ -164,7 +157,7 @@ void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) TIME_MS("timeMs") { @Override void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) { - builder.append(params._timeUsedMs); + builder.append(params._response.getTimeUsedMs()); } }, DOCS("docs") { @@ -190,7 +183,8 @@ void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) .append(params._response.getNumConsumingSegmentsQueried()).append('/') .append(params._response.getNumConsumingSegmentsProcessed()).append('/') .append(params._response.getNumConsumingSegmentsMatched()).append('/') - .append(params._numUnavailableSegments); + // TODO: Consider adding the number of unavailable segments to the response + .append(params._requestContext.getNumUnavailableSegments()); } }, CONSUMING_FRESHNESS_MS("consumingFreshnessTimeMs") { @@ -215,7 +209,7 @@ void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) BROKER_REDUCE_TIME_MS("brokerReduceTimeMs") { @Override void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) { - builder.append(params._requestContext.getReduceTimeMillis()); + builder.append(params._response.getBrokerReduceTimeMs()); } }, EXCEPTIONS("exceptions") { @@ -255,8 +249,8 @@ void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) CLIENT_IP("clientIp") { @Override void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) { - if (logger._enableIpLogging && params._requester != null) { - builder.append(params._requester.getClientIp()); + if (logger._enableIpLogging && params._identity != null) { + builder.append(params._identity.getClientIp()); } else { builder.append(CommonConstants.UNKNOWN); } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManager.java index 04db0f6a4255..dabb95867b9b 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManager.java @@ -224,6 +224,7 @@ private void createOrUpdateRateLimiter(String tableNameWithType, ExternalView br tableNameWithType, overallRate, previousRate, perBrokerRate, onlineCount, stat.getVersion()); } addMaxBurstQPSCallbackTableGaugeIfNeeded(tableNameWithType, queryQuotaEntity); + addQueryQuotaCapacityUtilizationRateTableGaugeIfNeeded(tableNameWithType, queryQuotaEntity); if (isQueryRateLimitDisabled()) { LOGGER.info("Query rate limiting is currently disabled for this broker. So it won't take effect immediately."); } @@ -245,6 +246,7 @@ private void buildEmptyOrResetRateLimiterInQueryQuotaEntity(String tableNameWith queryQuotaEntity.setRateLimiter(null); } addMaxBurstQPSCallbackTableGaugeIfNeeded(tableNameWithType, queryQuotaEntity); + addQueryQuotaCapacityUtilizationRateTableGaugeIfNeeded(tableNameWithType, queryQuotaEntity); } /** @@ -256,6 +258,27 @@ private void addMaxBurstQPSCallbackTableGaugeIfNeeded(String tableNameWithType, () -> (long) finalQueryQuotaEntity.getMaxQpsTracker().getMaxCountPerBucket()); } + /** + * Add the query quota capacity utilization rate table gauge to the metric system if the qps quota is specified. + */ + private void addQueryQuotaCapacityUtilizationRateTableGaugeIfNeeded(String tableNameWithType, + QueryQuotaEntity queryQuotaEntity) { + if (queryQuotaEntity.getRateLimiter() != null) { + final QueryQuotaEntity finalQueryQuotaEntity = queryQuotaEntity; + _brokerMetrics.setOrUpdateTableGauge(tableNameWithType, BrokerGauge.QUERY_QUOTA_CAPACITY_UTILIZATION_RATE, () -> { + double perBrokerRate = finalQueryQuotaEntity.getRateLimiter().getRate(); + int actualHitCountWithinTimeRange = finalQueryQuotaEntity.getMaxQpsTracker().getHitCount(); + long hitCountAllowedWithinTimeRage = + (long) (perBrokerRate * finalQueryQuotaEntity.getMaxQpsTracker().getDefaultTimeRangeMs() / 1000L); + // Since the MaxQpsTracker specifies 1-min window as valid time range, we can get the query quota capacity + // utilization by using the actual hit count within 1 min divided by the expected hit count within 1 min. + long percentageOfCapacityUtilization = actualHitCountWithinTimeRange * 100L / hitCountAllowedWithinTimeRage; + LOGGER.debug("The percentage of rate limit capacity utilization is {}", percentageOfCapacityUtilization); + return percentageOfCapacityUtilization; + }); + } + } + /** * {@inheritDoc} *

Acquires a token from rate limiter based on the table name. @@ -316,13 +339,6 @@ private boolean tryAcquireToken(String tableNameWithType, QueryQuotaEntity query // Emit the qps capacity utilization rate. int numHits = queryQuotaEntity.getQpsTracker().getHitCount(); - if (_brokerMetrics != null) { - int percentageOfCapacityUtilization = (int) (numHits * 100 / perBrokerRate); - LOGGER.debug("The percentage of rate limit capacity utilization is {}", percentageOfCapacityUtilization); - _brokerMetrics.setValueOfTableGauge(tableNameWithType, BrokerGauge.QUERY_QUOTA_CAPACITY_UTILIZATION_RATE, - percentageOfCapacityUtilization); - } - if (!rateLimiter.tryAcquire()) { LOGGER.info("Quota is exceeded for table: {}. Per-broker rate: {}. Current qps: {}", tableNameWithType, perBrokerRate, numHits); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HitCounter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HitCounter.java index eedc53903d1f..b656c0234449 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HitCounter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/HitCounter.java @@ -83,10 +83,20 @@ public int getHitCount() { @VisibleForTesting int getHitCount(long timestamp) { + return getHitCount(timestamp, _bucketCount); + } + + /** + * Get the hit count within the valid number of buckets. + * @param timestamp the current timestamp + * @param validBucketCount the valid number of buckets + * @return the number of hits within the valid bucket count + */ + int getHitCount(long timestamp, int validBucketCount) { long numTimeUnits = timestamp / _timeBucketWidthMs; int count = 0; for (int i = 0; i < _bucketCount; i++) { - if (numTimeUnits - _bucketStartTime.get(i) < _bucketCount) { + if (numTimeUnits - _bucketStartTime.get(i) < validBucketCount) { count += _bucketHitCount.get(i); } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/MaxHitRateTracker.java b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/MaxHitRateTracker.java index bdb8dbc2149f..b0cbd88b0bee 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/MaxHitRateTracker.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/queryquota/MaxHitRateTracker.java @@ -34,6 +34,7 @@ public class MaxHitRateTracker extends HitCounter { private static final int ONE_SECOND_BUCKET_WIDTH_MS = 1000; private static final int MAX_TIME_RANGE_FACTOR = 2; + private final int _validBucketCount; private final long _maxTimeRangeMs; private final long _defaultTimeRangeMs; private volatile long _lastAccessTimestamp; @@ -44,6 +45,7 @@ public MaxHitRateTracker(int timeRangeInSeconds) { private MaxHitRateTracker(int defaultTimeRangeInSeconds, int maxTimeRangeInSeconds) { super(maxTimeRangeInSeconds, (int) (maxTimeRangeInSeconds * 1000L / ONE_SECOND_BUCKET_WIDTH_MS)); + _validBucketCount = (int) (defaultTimeRangeInSeconds * 1000L / ONE_SECOND_BUCKET_WIDTH_MS); _defaultTimeRangeMs = defaultTimeRangeInSeconds * 1000L; _maxTimeRangeMs = maxTimeRangeInSeconds * 1000L; } @@ -80,4 +82,14 @@ int getMaxCountPerBucket(long now) { _lastAccessTimestamp = now; return maxCount; } + + @VisibleForTesting + @Override + int getHitCount(long now) { + return super.getHitCount(now, _validBucketCount); + } + + public long getDefaultTimeRangeMs() { + return _defaultTimeRangeMs; + } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java index 2fdc36e1ea3f..b43dcd9763aa 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java @@ -19,33 +19,18 @@ package org.apache.pinot.broker.requesthandler; import com.fasterxml.jackson.databind.JsonNode; -import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableMap; -import java.net.URI; +import com.google.common.collect.Maps; import java.util.ArrayList; -import java.util.Collections; -import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.concurrent.CompletionService; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Executor; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.stream.Collectors; import javax.annotation.Nullable; import javax.annotation.concurrent.ThreadSafe; import javax.ws.rs.WebApplicationException; import javax.ws.rs.core.HttpHeaders; +import javax.ws.rs.core.MultivaluedMap; import javax.ws.rs.core.Response; -import org.apache.commons.collections.CollectionUtils; import org.apache.commons.lang3.StringUtils; -import org.apache.commons.lang3.tuple.Pair; -import org.apache.http.client.methods.HttpDelete; -import org.apache.http.conn.HttpClientConnectionManager; -import org.apache.http.util.EntityUtils; import org.apache.pinot.broker.api.AccessControl; import org.apache.pinot.broker.api.RequesterIdentity; import org.apache.pinot.broker.broker.AccessControlFactory; @@ -54,1853 +39,133 @@ import org.apache.pinot.broker.routing.BrokerRoutingManager; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.exception.QueryException; -import org.apache.pinot.common.http.MultiHttpRequest; -import org.apache.pinot.common.http.MultiHttpRequestResponse; -import org.apache.pinot.common.metrics.BrokerGauge; import org.apache.pinot.common.metrics.BrokerMeter; import org.apache.pinot.common.metrics.BrokerMetrics; -import org.apache.pinot.common.metrics.BrokerQueryPhase; -import org.apache.pinot.common.metrics.BrokerTimer; -import org.apache.pinot.common.request.BrokerRequest; -import org.apache.pinot.common.request.DataSource; -import org.apache.pinot.common.request.Expression; -import org.apache.pinot.common.request.ExpressionType; -import org.apache.pinot.common.request.Function; -import org.apache.pinot.common.request.Identifier; -import org.apache.pinot.common.request.Literal; -import org.apache.pinot.common.request.PinotQuery; import org.apache.pinot.common.response.BrokerResponse; -import org.apache.pinot.common.response.ProcessingException; -import org.apache.pinot.common.response.broker.BrokerResponseNative; -import org.apache.pinot.common.response.broker.ResultTable; -import org.apache.pinot.common.utils.DataSchema; -import org.apache.pinot.common.utils.DatabaseUtils; -import org.apache.pinot.common.utils.config.QueryOptionsUtils; -import org.apache.pinot.common.utils.request.RequestUtils; -import org.apache.pinot.core.auth.Actions; -import org.apache.pinot.core.auth.TargetType; -import org.apache.pinot.core.query.optimizer.QueryOptimizer; -import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.TimeBoundaryInfo; -import org.apache.pinot.core.transport.ServerInstance; -import org.apache.pinot.core.util.GapfillUtils; -import org.apache.pinot.spi.config.table.FieldConfig; -import org.apache.pinot.spi.config.table.QueryConfig; -import org.apache.pinot.spi.config.table.TableConfig; -import org.apache.pinot.spi.config.table.TableType; -import org.apache.pinot.spi.data.DimensionFieldSpec; -import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.common.response.broker.QueryProcessingException; +import org.apache.pinot.spi.auth.AuthorizationResult; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListener; -import org.apache.pinot.spi.eventlistener.query.PinotBrokerQueryEventListenerFactory; +import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; import org.apache.pinot.spi.exception.BadQueryRequestException; -import org.apache.pinot.spi.exception.DatabaseConflictException; import org.apache.pinot.spi.trace.RequestContext; -import org.apache.pinot.spi.trace.Tracing; -import org.apache.pinot.spi.utils.BigDecimalUtils; -import org.apache.pinot.spi.utils.BytesUtils; -import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.CommonConstants.Broker; -import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; -import org.apache.pinot.spi.utils.DataSizeUtils; -import org.apache.pinot.spi.utils.TimestampIndexUtils; -import org.apache.pinot.spi.utils.builder.TableNameBuilder; -import org.apache.pinot.sql.FilterKind; -import org.apache.pinot.sql.parsers.CalciteSqlCompiler; -import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.apache.pinot.sql.parsers.SqlNodeAndOptions; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; @ThreadSafe public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler { - private static final Logger LOGGER = LoggerFactory.getLogger(BaseBrokerRequestHandler.class); - private static final String IN_SUBQUERY = "insubquery"; - private static final String IN_ID_SET = "inidset"; - private static final Expression FALSE = RequestUtils.getLiteralExpression(false); - private static final Expression TRUE = RequestUtils.getLiteralExpression(true); - private static final Expression STAR = RequestUtils.getIdentifierExpression("*"); - private static final int MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION = 10; - private static final Map DISTINCT_MV_COL_FUNCTION_OVERRIDE_MAP = - ImmutableMap.builder().put("distinctcount", "distinctcountmv") - .put("distinctcountbitmap", "distinctcountbitmapmv").put("distinctcounthll", "distinctcounthllmv") - .put("distinctcountrawhll", "distinctcountrawhllmv").put("distinctsum", "distinctsummv") - .put("distinctavg", "distinctavgmv").put("count", "countmv").put("min", "minmv").put("max", "maxmv") - .put("avg", "avgmv").put("sum", "summv").put("minmaxrange", "minmaxrangemv") - .put("distinctcounthllplus", "distinctcounthllplusmv") - .put("distinctcountrawhllplus", "distinctcountrawhllplusmv").build(); - protected final PinotConfiguration _config; + protected final String _brokerId; protected final BrokerRoutingManager _routingManager; protected final AccessControlFactory _accessControlFactory; protected final QueryQuotaManager _queryQuotaManager; protected final TableCache _tableCache; protected final BrokerMetrics _brokerMetrics; - - protected final BrokerRequestIdGenerator _brokerIdGenerator; - protected final QueryOptimizer _queryOptimizer = new QueryOptimizer(); - - protected final String _brokerId; + protected final BrokerQueryEventListener _brokerQueryEventListener; + protected final Set _trackedHeaders; + protected final BrokerRequestIdGenerator _requestIdGenerator; protected final long _brokerTimeoutMs; - protected final int _queryResponseLimit; protected final QueryLogger _queryLogger; - protected final BrokerQueryEventListener _brokerQueryEventListener; - - private final boolean _disableGroovy; - private final boolean _useApproximateFunction; - private final int _defaultHllLog2m; - private final boolean _enableQueryLimitOverride; - private final boolean _enableDistinctCountBitmapOverride; - private final Map _queriesById; public BaseBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, - AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, - BrokerMetrics brokerMetrics, BrokerQueryEventListener brokerQueryEventListener) { - _brokerId = brokerId; - _brokerIdGenerator = new BrokerRequestIdGenerator(brokerId); + AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache) { _config = config; + _brokerId = brokerId; _routingManager = routingManager; _accessControlFactory = accessControlFactory; _queryQuotaManager = queryQuotaManager; _tableCache = tableCache; - _brokerMetrics = brokerMetrics; - _disableGroovy = _config.getProperty(Broker.DISABLE_GROOVY, Broker.DEFAULT_DISABLE_GROOVY); - _useApproximateFunction = _config.getProperty(Broker.USE_APPROXIMATE_FUNCTION, false); - _defaultHllLog2m = _config.getProperty(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M_KEY, - CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M); - _enableQueryLimitOverride = _config.getProperty(Broker.CONFIG_OF_ENABLE_QUERY_LIMIT_OVERRIDE, false); - _enableDistinctCountBitmapOverride = - _config.getProperty(CommonConstants.Helix.ENABLE_DISTINCT_COUNT_BITMAP_OVERRIDE_KEY, false); - + _brokerMetrics = BrokerMetrics.get(); + _brokerQueryEventListener = BrokerQueryEventListenerFactory.getBrokerQueryEventListener(); + _trackedHeaders = BrokerQueryEventListenerFactory.getTrackedHeaders(); + _requestIdGenerator = new BrokerRequestIdGenerator(brokerId); _brokerTimeoutMs = config.getProperty(Broker.CONFIG_OF_BROKER_TIMEOUT_MS, Broker.DEFAULT_BROKER_TIMEOUT_MS); - _queryResponseLimit = - config.getProperty(Broker.CONFIG_OF_BROKER_QUERY_RESPONSE_LIMIT, Broker.DEFAULT_BROKER_QUERY_RESPONSE_LIMIT); _queryLogger = new QueryLogger(config); - boolean enableQueryCancellation = - Boolean.parseBoolean(config.getProperty(Broker.CONFIG_OF_BROKER_ENABLE_QUERY_CANCELLATION)); - _queriesById = enableQueryCancellation ? new ConcurrentHashMap<>() : null; - _brokerQueryEventListener = brokerQueryEventListener; - LOGGER.info( - "Broker Id: {}, timeout: {}ms, query response limit: {}, query log length: {}, query log max rate: {}qps, " - + "enabling query cancellation: {}", _brokerId, _brokerTimeoutMs, _queryResponseLimit, - _queryLogger.getMaxQueryLengthToLog(), _queryLogger.getLogRateLimit(), enableQueryCancellation); - } - - @Override - public Map getRunningQueries() { - Preconditions.checkState(_queriesById != null, "Query cancellation is not enabled on broker"); - return _queriesById.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue()._query)); - } - - @VisibleForTesting - Set getRunningServers(long requestId) { - Preconditions.checkState(_queriesById != null, "Query cancellation is not enabled on broker"); - QueryServers queryServers = _queriesById.get(requestId); - return queryServers != null ? queryServers._servers : Collections.emptySet(); - } - - @Override - public boolean cancelQuery(long requestId, int timeoutMs, Executor executor, HttpClientConnectionManager connMgr, - Map serverResponses) - throws Exception { - Preconditions.checkState(_queriesById != null, "Query cancellation is not enabled on broker"); - QueryServers queryServers = _queriesById.get(requestId); - if (queryServers == null) { - return false; - } - // TODO: Use different global query id for OFFLINE and REALTIME table after releasing 0.12.0. See QueryIdUtils for - // details - String globalQueryId = getGlobalQueryId(requestId); - List> serverUrls = new ArrayList<>(); - for (ServerInstance serverInstance : queryServers._servers) { - serverUrls.add(Pair.of(String.format("%s/query/%s", serverInstance.getAdminEndpoint(), globalQueryId), null)); - } - LOGGER.debug("Cancelling the query: {} via server urls: {}", queryServers._query, serverUrls); - CompletionService completionService = - new MultiHttpRequest(executor, connMgr).execute(serverUrls, null, timeoutMs, "DELETE", HttpDelete::new); - List errMsgs = new ArrayList<>(serverUrls.size()); - for (int i = 0; i < serverUrls.size(); i++) { - MultiHttpRequestResponse httpRequestResponse = null; - try { - // Wait for all requests to respond before returning to be sure that the servers have handled the cancel - // requests. The completion order is different from serverUrls, thus use uri in the response. - httpRequestResponse = completionService.take().get(); - URI uri = httpRequestResponse.getURI(); - int status = httpRequestResponse.getResponse().getStatusLine().getStatusCode(); - // Unexpected server responses are collected and returned as exception. - if (status != 200 && status != 404) { - String responseString = EntityUtils.toString(httpRequestResponse.getResponse().getEntity()); - throw new Exception( - String.format("Unexpected status=%d and response='%s' from uri='%s'", status, responseString, uri)); - } - if (serverResponses != null) { - serverResponses.put(uri.getHost() + ":" + uri.getPort(), status); - } - } catch (Exception e) { - LOGGER.error("Failed to cancel query: {}", queryServers._query, e); - // Can't just throw exception from here as there is a need to release the other connections. - // So just collect the error msg to throw them together after the for-loop. - errMsgs.add(e.getMessage()); - } finally { - if (httpRequestResponse != null) { - httpRequestResponse.close(); - } - } - } - if (errMsgs.size() > 0) { - throw new Exception("Unexpected responses from servers: " + StringUtils.join(errMsgs, ",")); - } - return true; } @Override public BrokerResponse handleRequest(JsonNode request, @Nullable SqlNodeAndOptions sqlNodeAndOptions, @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders) throws Exception { - requestContext.setRequestArrivalTimeMillis(System.currentTimeMillis()); - - long requestId = _brokerIdGenerator.get(); + requestContext.setBrokerId(_brokerId); + long requestId = _requestIdGenerator.get(); requestContext.setRequestId(requestId); - if (httpHeaders != null) { - requestContext.setRequestHttpHeaders(httpHeaders.getRequestHeaders().entrySet().stream().filter( - entry -> PinotBrokerQueryEventListenerFactory.getAllowlistQueryRequestHeaders().contains(entry.getKey())) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))); + + if (httpHeaders != null && !_trackedHeaders.isEmpty()) { + MultivaluedMap requestHeaders = httpHeaders.getRequestHeaders(); + Map> trackedHeadersMap = Maps.newHashMapWithExpectedSize(_trackedHeaders.size()); + for (Map.Entry> entry : requestHeaders.entrySet()) { + String key = entry.getKey().toLowerCase(); + if (_trackedHeaders.contains(key)) { + trackedHeadersMap.put(key, entry.getValue()); + } + } + requestContext.setRequestHttpHeaders(trackedHeadersMap); } // First-stage access control to prevent unauthenticated requests from using up resources. Secondary table-level // check comes later. - boolean hasAccess = _accessControlFactory.create().hasAccess(requesterIdentity); - if (!hasAccess) { + AccessControl accessControl = _accessControlFactory.create(); + AuthorizationResult authorizationResult = accessControl.authorize(requesterIdentity); + if (!authorizationResult.hasAccess()) { _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1); requestContext.setErrorCode(QueryException.ACCESS_DENIED_ERROR_CODE); _brokerQueryEventListener.onQueryCompletion(requestContext); - throw new WebApplicationException("Permission denied", Response.Status.FORBIDDEN); + String failureMessage = authorizationResult.getFailureMessage(); + if (StringUtils.isNotBlank(failureMessage)) { + failureMessage = "Reason: " + failureMessage; + } + throw new WebApplicationException("Permission denied." + failureMessage, + Response.Status.FORBIDDEN); } JsonNode sql = request.get(Broker.Request.SQL); - if (sql == null) { + if (sql == null || !sql.isTextual()) { requestContext.setErrorCode(QueryException.SQL_PARSING_ERROR_CODE); _brokerQueryEventListener.onQueryCompletion(requestContext); throw new BadQueryRequestException("Failed to find 'sql' in the request: " + request); } - String query = sql.asText(); + + String query = sql.textValue(); requestContext.setQuery(query); BrokerResponse brokerResponse = - handleRequest(requestId, query, sqlNodeAndOptions, request, requesterIdentity, requestContext, httpHeaders); - - if (request.has(Broker.Request.QUERY_OPTIONS)) { - String queryOptions = request.get(Broker.Request.QUERY_OPTIONS).asText(); - Map optionsFromString = RequestUtils.getOptionsFromString(queryOptions); - if (QueryOptionsUtils.shouldDropResults(optionsFromString)) { - brokerResponse.setResultTable(null); - } - } - - brokerResponse.setRequestId(String.valueOf(requestId)); + handleRequest(requestId, query, sqlNodeAndOptions, request, requesterIdentity, requestContext, httpHeaders, + accessControl); brokerResponse.setBrokerId(_brokerId); - brokerResponse.setBrokerReduceTimeMs(requestContext.getReduceTimeMillis()); + brokerResponse.setRequestId(Long.toString(requestId)); _brokerQueryEventListener.onQueryCompletion(requestContext); - return brokerResponse; - } - - protected BrokerResponse handleRequest(long requestId, String query, @Nullable SqlNodeAndOptions sqlNodeAndOptions, - JsonNode request, @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, - @Nullable HttpHeaders httpHeaders) - throws Exception { - LOGGER.debug("SQL query for request {}: {}", requestId, query); - - //Start instrumentation context. This must not be moved further below interspersed into the code. - Tracing.ThreadAccountantOps.setupRunner(String.valueOf(requestId)); - - try { - long compilationStartTimeNs; - PinotQuery pinotQuery; - try { - // Parse the request - sqlNodeAndOptions = sqlNodeAndOptions != null ? sqlNodeAndOptions : RequestUtils.parseQuery(query, request); - // Compile the request into PinotQuery - compilationStartTimeNs = System.nanoTime(); - pinotQuery = CalciteSqlParser.compileToPinotQuery(sqlNodeAndOptions); - } catch (Exception e) { - LOGGER.info("Caught exception while compiling SQL request {}: {}, {}", requestId, query, e.getMessage()); - _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_COMPILATION_EXCEPTIONS, 1); - requestContext.setErrorCode(QueryException.SQL_PARSING_ERROR_CODE); - return new BrokerResponseNative(QueryException.getException(QueryException.SQL_PARSING_ERROR, e)); - } - - if (isLiteralOnlyQuery(pinotQuery)) { - LOGGER.debug("Request {} contains only Literal, skipping server query: {}", requestId, query); - try { - if (pinotQuery.isExplain()) { - // EXPLAIN PLAN results to show that query is evaluated exclusively by Broker. - return BrokerResponseNative.BROKER_ONLY_EXPLAIN_PLAN_OUTPUT; - } - return processLiteralOnlyQuery(pinotQuery, compilationStartTimeNs, requestContext); - } catch (Exception e) { - // TODO: refine the exceptions here to early termination the queries won't requires to send to servers. - LOGGER.warn("Unable to execute literal request {}: {} at broker, fallback to server query. {}", requestId, - query, e.getMessage()); - } - } - - PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery); - DataSource dataSource = serverPinotQuery.getDataSource(); - if (dataSource == null) { - LOGGER.info("Data source (FROM clause) not found in request {}: {}", requestId, query); - requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); - return new BrokerResponseNative( - QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, "Data source (FROM clause) not found")); - } - if (dataSource.getJoin() != null) { - LOGGER.info("JOIN is not supported in request {}: {}", requestId, query); - requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); - return new BrokerResponseNative( - QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, "JOIN is not supported")); - } - if (dataSource.getTableName() == null) { - LOGGER.info("Table name not found in request {}: {}", requestId, query); - requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); - return new BrokerResponseNative( - QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, "Table name not found")); - } - - try { - handleSubquery(serverPinotQuery, requestId, request, requesterIdentity, requestContext, httpHeaders); - } catch (Exception e) { - LOGGER.info("Caught exception while handling the subquery in request {}: {}, {}", requestId, query, - e.getMessage()); - requestContext.setErrorCode(QueryException.QUERY_EXECUTION_ERROR_CODE); - return new BrokerResponseNative(QueryException.getException(QueryException.QUERY_EXECUTION_ERROR, e)); - } - - boolean ignoreCase = _tableCache.isIgnoreCase(); - String tableName; - try { - tableName = - getActualTableName(DatabaseUtils.translateTableName(dataSource.getTableName(), httpHeaders, ignoreCase), - _tableCache); - } catch (DatabaseConflictException e) { - LOGGER.info("{}. Request {}: {}", e.getMessage(), requestId, query); - _brokerMetrics.addMeteredGlobalValue(BrokerMeter.QUERY_VALIDATION_EXCEPTIONS, 1); - requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); - return new BrokerResponseNative(QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, e)); - } - dataSource.setTableName(tableName); - String rawTableName = TableNameBuilder.extractRawTableName(tableName); - requestContext.setTableName(rawTableName); - - try { - Map columnNameMap = _tableCache.getColumnNameMap(rawTableName); - if (columnNameMap != null) { - updateColumnNames(rawTableName, serverPinotQuery, ignoreCase, columnNameMap); - } - } catch (Exception e) { - // Throw exceptions with column in-existence error. - if (e instanceof BadQueryRequestException) { - LOGGER.info("Caught exception while checking column names in request {}: {}, {}", requestId, query, - e.getMessage()); - requestContext.setErrorCode(QueryException.UNKNOWN_COLUMN_ERROR_CODE); - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.UNKNOWN_COLUMN_EXCEPTIONS, 1); - return new BrokerResponseNative(QueryException.getException(QueryException.UNKNOWN_COLUMN_ERROR, e)); - } - LOGGER.warn("Caught exception while updating column names in request {}: {}, {}", requestId, query, - e.getMessage()); - } - if (_defaultHllLog2m > 0) { - handleHLLLog2mOverride(serverPinotQuery, _defaultHllLog2m); - } - if (_enableQueryLimitOverride) { - handleQueryLimitOverride(serverPinotQuery, _queryResponseLimit); - } - handleSegmentPartitionedDistinctCountOverride(serverPinotQuery, - getSegmentPartitionedColumns(_tableCache, tableName)); - if (_enableDistinctCountBitmapOverride) { - handleDistinctCountBitmapOverride(serverPinotQuery); - } - - Schema schema = _tableCache.getSchema(rawTableName); - if (schema != null) { - handleDistinctMultiValuedOverride(serverPinotQuery, schema); - } - - long compilationEndTimeNs = System.nanoTime(); - // full request compile time = compilationTimeNs + parserTimeNs - _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.REQUEST_COMPILATION, - (compilationEndTimeNs - compilationStartTimeNs) + sqlNodeAndOptions.getParseTimeNs()); - - // Second-stage table-level access control - // TODO: Modify AccessControl interface to directly take PinotQuery - BrokerRequest brokerRequest = CalciteSqlCompiler.convertToBrokerRequest(pinotQuery); - BrokerRequest serverBrokerRequest = - serverPinotQuery == pinotQuery ? brokerRequest : CalciteSqlCompiler.convertToBrokerRequest(serverPinotQuery); - AccessControl accessControl = _accessControlFactory.create(); - boolean hasTableAccess = - accessControl.hasAccess(requesterIdentity, serverBrokerRequest) && accessControl.hasAccess(httpHeaders, - TargetType.TABLE, tableName, Actions.Table.QUERY); - - _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.AUTHORIZATION, - System.nanoTime() - compilationEndTimeNs); - - if (!hasTableAccess) { - _brokerMetrics.addMeteredTableValue(tableName, BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1); - LOGGER.info("Access denied for request {}: {}, table: {}", requestId, query, tableName); - requestContext.setErrorCode(QueryException.ACCESS_DENIED_ERROR_CODE); - throw new WebApplicationException("Permission denied", Response.Status.FORBIDDEN); - } - - // Get the tables hit by the request - String offlineTableName = null; - String realtimeTableName = null; - TableType tableType = TableNameBuilder.getTableTypeFromTableName(tableName); - if (tableType == TableType.OFFLINE) { - // Offline table - if (_routingManager.routingExists(tableName)) { - offlineTableName = tableName; - } - } else if (tableType == TableType.REALTIME) { - // Realtime table - if (_routingManager.routingExists(tableName)) { - realtimeTableName = tableName; - } - } else { - // Hybrid table (check both OFFLINE and REALTIME) - String offlineTableNameToCheck = TableNameBuilder.OFFLINE.tableNameWithType(tableName); - if (_routingManager.routingExists(offlineTableNameToCheck)) { - offlineTableName = offlineTableNameToCheck; - } - String realtimeTableNameToCheck = TableNameBuilder.REALTIME.tableNameWithType(tableName); - if (_routingManager.routingExists(realtimeTableNameToCheck)) { - realtimeTableName = realtimeTableNameToCheck; - } - } - - TableConfig offlineTableConfig = - _tableCache.getTableConfig(TableNameBuilder.OFFLINE.tableNameWithType(rawTableName)); - TableConfig realtimeTableConfig = - _tableCache.getTableConfig(TableNameBuilder.REALTIME.tableNameWithType(rawTableName)); - - if (offlineTableName == null && realtimeTableName == null) { - // No table matches the request - if (realtimeTableConfig == null && offlineTableConfig == null) { - LOGGER.info("Table not found for request {}: {}", requestId, query); - requestContext.setErrorCode(QueryException.TABLE_DOES_NOT_EXIST_ERROR_CODE); - return BrokerResponseNative.TABLE_DOES_NOT_EXIST; - } - LOGGER.info("No table matches for request {}: {}", requestId, query); - requestContext.setErrorCode(QueryException.BROKER_RESOURCE_MISSING_ERROR_CODE); - _brokerMetrics.addMeteredGlobalValue(BrokerMeter.RESOURCE_MISSING_EXCEPTIONS, 1); - return BrokerResponseNative.NO_TABLE_RESULT; - } - - // Handle query rewrite that can be overridden by the table configs - if (offlineTableName == null) { - offlineTableConfig = null; - } - if (realtimeTableName == null) { - realtimeTableConfig = null; - } - HandlerContext handlerContext = getHandlerContext(offlineTableConfig, realtimeTableConfig); - if (handlerContext._disableGroovy) { - rejectGroovyQuery(serverPinotQuery); - } - if (handlerContext._useApproximateFunction) { - handleApproximateFunctionOverride(serverPinotQuery); - } - - // Validate QPS quota - if (!_queryQuotaManager.acquire(tableName)) { - String errorMessage = - String.format("Request %d: %s exceeds query quota for table: %s", requestId, query, tableName); - LOGGER.info(errorMessage); - requestContext.setErrorCode(QueryException.TOO_MANY_REQUESTS_ERROR_CODE); - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.QUERY_QUOTA_EXCEEDED, 1); - return new BrokerResponseNative(QueryException.getException(QueryException.QUOTA_EXCEEDED_ERROR, errorMessage)); - } - - // Validate the request - try { - validateRequest(serverPinotQuery, _queryResponseLimit); - } catch (Exception e) { - LOGGER.info("Caught exception while validating request {}: {}, {}", requestId, query, e.getMessage()); - requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.QUERY_VALIDATION_EXCEPTIONS, 1); - return new BrokerResponseNative(QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, e)); - } - - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.QUERIES, 1); - _brokerMetrics.addValueToTableGauge(rawTableName, BrokerGauge.REQUEST_SIZE, query.length()); - - // Prepare OFFLINE and REALTIME requests - BrokerRequest offlineBrokerRequest = null; - BrokerRequest realtimeBrokerRequest = null; - TimeBoundaryInfo timeBoundaryInfo = null; - if (offlineTableName != null && realtimeTableName != null) { - // Time boundary info might be null when there is no segment in the offline table, query real-time side only - timeBoundaryInfo = _routingManager.getTimeBoundaryInfo(offlineTableName); - if (timeBoundaryInfo == null) { - LOGGER.debug("No time boundary info found for hybrid table: {}", rawTableName); - offlineTableName = null; - } - } - if (offlineTableName != null && realtimeTableName != null) { - // Hybrid - PinotQuery offlinePinotQuery = serverPinotQuery.deepCopy(); - offlinePinotQuery.getDataSource().setTableName(offlineTableName); - attachTimeBoundary(offlinePinotQuery, timeBoundaryInfo, true); - handleExpressionOverride(offlinePinotQuery, _tableCache.getExpressionOverrideMap(offlineTableName)); - handleTimestampIndexOverride(offlinePinotQuery, offlineTableConfig); - _queryOptimizer.optimize(offlinePinotQuery, offlineTableConfig, schema); - offlineBrokerRequest = CalciteSqlCompiler.convertToBrokerRequest(offlinePinotQuery); - - PinotQuery realtimePinotQuery = serverPinotQuery.deepCopy(); - realtimePinotQuery.getDataSource().setTableName(realtimeTableName); - attachTimeBoundary(realtimePinotQuery, timeBoundaryInfo, false); - handleExpressionOverride(realtimePinotQuery, _tableCache.getExpressionOverrideMap(realtimeTableName)); - handleTimestampIndexOverride(realtimePinotQuery, realtimeTableConfig); - _queryOptimizer.optimize(realtimePinotQuery, realtimeTableConfig, schema); - realtimeBrokerRequest = CalciteSqlCompiler.convertToBrokerRequest(realtimePinotQuery); - - requestContext.setFanoutType(RequestContext.FanoutType.HYBRID); - requestContext.setOfflineServerTenant(getServerTenant(offlineTableName)); - requestContext.setRealtimeServerTenant(getServerTenant(realtimeTableName)); - } else if (offlineTableName != null) { - // OFFLINE only - setTableName(serverBrokerRequest, offlineTableName); - handleExpressionOverride(serverPinotQuery, _tableCache.getExpressionOverrideMap(offlineTableName)); - handleTimestampIndexOverride(serverPinotQuery, offlineTableConfig); - _queryOptimizer.optimize(serverPinotQuery, offlineTableConfig, schema); - offlineBrokerRequest = serverBrokerRequest; - - requestContext.setFanoutType(RequestContext.FanoutType.OFFLINE); - requestContext.setOfflineServerTenant(getServerTenant(offlineTableName)); - } else { - // REALTIME only - setTableName(serverBrokerRequest, realtimeTableName); - handleExpressionOverride(serverPinotQuery, _tableCache.getExpressionOverrideMap(realtimeTableName)); - handleTimestampIndexOverride(serverPinotQuery, realtimeTableConfig); - _queryOptimizer.optimize(serverPinotQuery, realtimeTableConfig, schema); - realtimeBrokerRequest = serverBrokerRequest; - - requestContext.setFanoutType(RequestContext.FanoutType.REALTIME); - requestContext.setRealtimeServerTenant(getServerTenant(realtimeTableName)); - } - - // Check if response can be sent without server query evaluation. - if (offlineBrokerRequest != null && isFilterAlwaysFalse(offlineBrokerRequest.getPinotQuery())) { - // We don't need to evaluate offline request - offlineBrokerRequest = null; - } - if (realtimeBrokerRequest != null && isFilterAlwaysFalse(realtimeBrokerRequest.getPinotQuery())) { - // We don't need to evaluate realtime request - realtimeBrokerRequest = null; - } - - if (offlineBrokerRequest == null && realtimeBrokerRequest == null) { - return getEmptyBrokerOnlyResponse(requestId, query, requesterIdentity, requestContext, pinotQuery, tableName); - } - - if (offlineBrokerRequest != null && isFilterAlwaysTrue(offlineBrokerRequest.getPinotQuery())) { - // Drop offline request filter since it is always true - offlineBrokerRequest.getPinotQuery().setFilterExpression(null); - } - if (realtimeBrokerRequest != null && isFilterAlwaysTrue(realtimeBrokerRequest.getPinotQuery())) { - // Drop realtime request filter since it is always true - realtimeBrokerRequest.getPinotQuery().setFilterExpression(null); - } - - // Calculate routing table for the query - // TODO: Modify RoutingManager interface to directly take PinotQuery - long routingStartTimeNs = System.nanoTime(); - Map, List>> offlineRoutingTable = null; - Map, List>> realtimeRoutingTable = null; - List unavailableSegments = new ArrayList<>(); - int numPrunedSegmentsTotal = 0; - if (offlineBrokerRequest != null) { - // NOTE: Routing table might be null if table is just removed - RoutingTable routingTable = _routingManager.getRoutingTable(offlineBrokerRequest, requestId); - if (routingTable != null) { - unavailableSegments.addAll(routingTable.getUnavailableSegments()); - Map, List>> serverInstanceToSegmentsMap = - routingTable.getServerInstanceToSegmentsMap(); - if (!serverInstanceToSegmentsMap.isEmpty()) { - offlineRoutingTable = serverInstanceToSegmentsMap; - } else { - offlineBrokerRequest = null; - } - numPrunedSegmentsTotal += routingTable.getNumPrunedSegments(); - } else { - offlineBrokerRequest = null; - } - } - if (realtimeBrokerRequest != null) { - // NOTE: Routing table might be null if table is just removed - RoutingTable routingTable = _routingManager.getRoutingTable(realtimeBrokerRequest, requestId); - if (routingTable != null) { - unavailableSegments.addAll(routingTable.getUnavailableSegments()); - Map, List>> serverInstanceToSegmentsMap = - routingTable.getServerInstanceToSegmentsMap(); - if (!serverInstanceToSegmentsMap.isEmpty()) { - realtimeRoutingTable = serverInstanceToSegmentsMap; - } else { - realtimeBrokerRequest = null; - } - numPrunedSegmentsTotal += routingTable.getNumPrunedSegments(); - } else { - realtimeBrokerRequest = null; - } - } - int numUnavailableSegments = unavailableSegments.size(); - requestContext.setNumUnavailableSegments(numUnavailableSegments); - - List exceptions = new ArrayList<>(); - if (numUnavailableSegments > 0) { - String errorMessage; - if (numUnavailableSegments > MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION) { - errorMessage = String.format("%d segments unavailable, sampling %d: %s", numUnavailableSegments, - MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION, - unavailableSegments.subList(0, MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION)); - } else { - errorMessage = String.format("%d segments unavailable: %s", numUnavailableSegments, unavailableSegments); - } - exceptions.add(QueryException.getException(QueryException.BROKER_SEGMENT_UNAVAILABLE_ERROR, errorMessage)); - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_UNAVAILABLE_SEGMENTS, 1); - } - - if (offlineBrokerRequest == null && realtimeBrokerRequest == null) { - if (!exceptions.isEmpty()) { - LOGGER.info("No server found for request {}: {}", requestId, query); - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.NO_SERVER_FOUND_EXCEPTIONS, 1); - return new BrokerResponseNative(exceptions); - } else { - // When all segments have been pruned, we can just return an empty response. - return getEmptyBrokerOnlyResponse(requestId, query, requesterIdentity, requestContext, pinotQuery, tableName); - } - } - long routingEndTimeNs = System.nanoTime(); - _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.QUERY_ROUTING, - routingEndTimeNs - routingStartTimeNs); - - // Set timeout in the requests - long timeSpentMs = TimeUnit.NANOSECONDS.toMillis(routingEndTimeNs - compilationStartTimeNs); - // Remaining time in milliseconds for the server query execution - // NOTE: For hybrid use case, in most cases offline table and real-time table should have the same query timeout - // configured, but if necessary, we also allow different timeout for them. - // If the timeout is not the same for offline table and real-time table, use the max of offline table - // remaining time and realtime table remaining time. Server side will have different remaining time set for - // each table type, and broker should wait for both types to return. - long remainingTimeMs = 0; - try { - if (offlineBrokerRequest != null) { - remainingTimeMs = - setQueryTimeout(offlineTableName, offlineBrokerRequest.getPinotQuery().getQueryOptions(), timeSpentMs); - } - if (realtimeBrokerRequest != null) { - remainingTimeMs = Math.max(remainingTimeMs, - setQueryTimeout(realtimeTableName, realtimeBrokerRequest.getPinotQuery().getQueryOptions(), timeSpentMs)); - } - } catch (TimeoutException e) { - String errorMessage = e.getMessage(); - LOGGER.info("{} {}: {}", errorMessage, requestId, query); - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.REQUEST_TIMEOUT_BEFORE_SCATTERED_EXCEPTIONS, 1); - exceptions.add(QueryException.getException(QueryException.BROKER_TIMEOUT_ERROR, errorMessage)); - return new BrokerResponseNative(exceptions); - } - - // Set the maximum serialized response size per server, and ask server to directly return final response when only - // one server is queried - int numServers = 0; - if (offlineRoutingTable != null) { - numServers += offlineRoutingTable.size(); - } - if (realtimeRoutingTable != null) { - numServers += realtimeRoutingTable.size(); - } - if (offlineBrokerRequest != null) { - Map queryOptions = offlineBrokerRequest.getPinotQuery().getQueryOptions(); - setMaxServerResponseSizeBytes(numServers, queryOptions, offlineTableConfig); - // Set the query option to directly return final result for single server query unless it is explicitly disabled - if (numServers == 1) { - // Set the same flag in the original server request to be used in the reduce phase for hybrid table - if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null - && offlineBrokerRequest != serverBrokerRequest) { - serverBrokerRequest.getPinotQuery().getQueryOptions() - .put(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true"); - } - } - } - if (realtimeBrokerRequest != null) { - Map queryOptions = realtimeBrokerRequest.getPinotQuery().getQueryOptions(); - setMaxServerResponseSizeBytes(numServers, queryOptions, realtimeTableConfig); - // Set the query option to directly return final result for single server query unless it is explicitly disabled - if (numServers == 1) { - // Set the same flag in the original server request to be used in the reduce phase for hybrid table - if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null - && realtimeBrokerRequest != serverBrokerRequest) { - serverBrokerRequest.getPinotQuery().getQueryOptions() - .put(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true"); - } - } - } - - // Execute the query - // TODO: Replace ServerStats with ServerRoutingStatsEntry. - ServerStats serverStats = new ServerStats(); - // TODO: Handle broker specific operations for explain plan queries such as: - // - Alias handling - // - Compile time function invocation - // - Literal only queries - // - Any rewrites - if (pinotQuery.isExplain()) { - // Update routing tables to only send request to offline servers for OFFLINE and HYBRID tables. - // TODO: Assess if the Explain Plan Query should also be routed to REALTIME servers for HYBRID tables - if (offlineRoutingTable != null) { - // For OFFLINE and HYBRID tables, don't send EXPLAIN query to realtime servers. - realtimeBrokerRequest = null; - realtimeRoutingTable = null; - } - } - BrokerResponseNative brokerResponse; - if (_queriesById != null) { - // Start to track the running query for cancellation just before sending it out to servers to avoid any - // potential failures that could happen before sending it out, like failures to calculate the routing table etc. - // TODO: Even tracking the query as late as here, a potential race condition between calling cancel API and - // query being sent out to servers can still happen. If cancel request arrives earlier than query being - // sent out to servers, the servers miss the cancel request and continue to run the queries. The users - // can always list the running queries and cancel query again until it ends. Just that such race - // condition makes cancel API less reliable. This should be rare as it assumes sending queries out to - // servers takes time, but will address later if needed. - _queriesById.put(requestId, new QueryServers(query, offlineRoutingTable, realtimeRoutingTable)); - LOGGER.debug("Keep track of running query: {}", requestId); - try { - brokerResponse = processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, offlineBrokerRequest, - offlineRoutingTable, realtimeBrokerRequest, realtimeRoutingTable, remainingTimeMs, serverStats, - requestContext); - } finally { - _queriesById.remove(requestId); - LOGGER.debug("Remove track of running query: {}", requestId); - } - } else { - brokerResponse = processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, offlineBrokerRequest, - offlineRoutingTable, realtimeBrokerRequest, realtimeRoutingTable, remainingTimeMs, serverStats, - requestContext); - } - - brokerResponse.setExceptions(exceptions); - brokerResponse.setNumSegmentsPrunedByBroker(numPrunedSegmentsTotal); - long executionEndTimeNs = System.nanoTime(); - _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.QUERY_EXECUTION, - executionEndTimeNs - routingEndTimeNs); - - // Track number of queries with number of groups limit reached - if (brokerResponse.isNumGroupsLimitReached()) { - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_NUM_GROUPS_LIMIT_REACHED, - 1); - } - brokerResponse.setPartialResult(isPartialResult(brokerResponse)); - - // Set total query processing time - long totalTimeMs = TimeUnit.NANOSECONDS.toMillis(executionEndTimeNs - compilationStartTimeNs); - brokerResponse.setTimeUsedMs(totalTimeMs); - requestContext.setQueryProcessingTime(totalTimeMs); - augmentStatistics(requestContext, brokerResponse); - _brokerMetrics.addTimedTableValue(rawTableName, BrokerTimer.QUERY_TOTAL_TIME_MS, totalTimeMs, - TimeUnit.MILLISECONDS); - - // Log query and stats - _queryLogger.log( - new QueryLogger.QueryLogParams(requestId, query, requestContext, tableName, numUnavailableSegments, - serverStats, brokerResponse, totalTimeMs, requesterIdentity)); - - return brokerResponse; - } finally { - Tracing.ThreadAccountantOps.clear(); - } - } - - private BrokerResponseNative getEmptyBrokerOnlyResponse(long requestId, String query, - RequesterIdentity requesterIdentity, RequestContext requestContext, PinotQuery pinotQuery, String tableName) { - if (pinotQuery.isExplain()) { - // EXPLAIN PLAN results to show that query is evaluated exclusively by Broker. - return BrokerResponseNative.BROKER_ONLY_EXPLAIN_PLAN_OUTPUT; - } - // Send empty response since we don't need to evaluate either offline or realtime request. - BrokerResponseNative brokerResponse = BrokerResponseNative.empty(); - // Extract source info from incoming request - _queryLogger.log( - new QueryLogger.QueryLogParams(requestId, query, requestContext, tableName, 0, null, brokerResponse, - System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis(), requesterIdentity)); return brokerResponse; } - private void handleTimestampIndexOverride(PinotQuery pinotQuery, @Nullable TableConfig tableConfig) { - if (tableConfig == null || tableConfig.getFieldConfigList() == null) { - return; - } - - Set timestampIndexColumns = _tableCache.getTimestampIndexColumns(tableConfig.getTableName()); - if (CollectionUtils.isEmpty(timestampIndexColumns)) { - return; - } - for (Expression expression : pinotQuery.getSelectList()) { - setTimestampIndexExpressionOverrideHints(expression, timestampIndexColumns, pinotQuery); - } - setTimestampIndexExpressionOverrideHints(pinotQuery.getFilterExpression(), timestampIndexColumns, pinotQuery); - setTimestampIndexExpressionOverrideHints(pinotQuery.getHavingExpression(), timestampIndexColumns, pinotQuery); - List groupByList = pinotQuery.getGroupByList(); - if (CollectionUtils.isNotEmpty(groupByList)) { - groupByList.forEach( - expression -> setTimestampIndexExpressionOverrideHints(expression, timestampIndexColumns, pinotQuery)); - } - List orderByList = pinotQuery.getOrderByList(); - if (CollectionUtils.isNotEmpty(orderByList)) { - orderByList.forEach( - expression -> setTimestampIndexExpressionOverrideHints(expression, timestampIndexColumns, pinotQuery)); - } - } - - private void setTimestampIndexExpressionOverrideHints(@Nullable Expression expression, - Set timestampIndexColumns, PinotQuery pinotQuery) { - if (expression == null || expression.getFunctionCall() == null) { - return; - } - Function function = expression.getFunctionCall(); - switch (function.getOperator()) { - case "datetrunc": - String granularString = function.getOperands().get(0).getLiteral().getStringValue().toUpperCase(); - Expression timeExpression = function.getOperands().get(1); - if (((function.getOperandsSize() == 2) || (function.getOperandsSize() == 3 && "MILLISECONDS".equalsIgnoreCase( - function.getOperands().get(2).getLiteral().getStringValue()))) && TimestampIndexUtils.isValidGranularity( - granularString) && timeExpression.getIdentifier() != null) { - String timeColumn = timeExpression.getIdentifier().getName(); - String timeColumnWithGranularity = TimestampIndexUtils.getColumnWithGranularity(timeColumn, granularString); - if (timestampIndexColumns.contains(timeColumnWithGranularity)) { - pinotQuery.putToExpressionOverrideHints(expression, - RequestUtils.getIdentifierExpression(timeColumnWithGranularity)); - } - } - break; - default: - break; - } - function.getOperands() - .forEach(operand -> setTimestampIndexExpressionOverrideHints(operand, timestampIndexColumns, pinotQuery)); - } - - /** Given a {@link PinotQuery}, check if the WHERE clause will always evaluate to false. */ - private boolean isFilterAlwaysFalse(PinotQuery pinotQuery) { - return FALSE.equals(pinotQuery.getFilterExpression()); - } - - /** Given a {@link PinotQuery}, check if the WHERE clause will always evaluate to true. */ - private boolean isFilterAlwaysTrue(PinotQuery pinotQuery) { - return TRUE.equals(pinotQuery.getFilterExpression()); - } - - private String getServerTenant(String tableNameWithType) { - TableConfig tableConfig = _tableCache.getTableConfig(tableNameWithType); - if (tableConfig == null) { - LOGGER.debug("Table config is not available for table {}", tableNameWithType); - return "unknownTenant"; - } - return tableConfig.getTenantConfig().getServer(); - } - - /** - * Handles the subquery in the given query. - *

Currently only supports subquery within the filter. - */ - private void handleSubquery(PinotQuery pinotQuery, long requestId, JsonNode jsonRequest, - @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders) - throws Exception { - Expression filterExpression = pinotQuery.getFilterExpression(); - if (filterExpression != null) { - handleSubquery(filterExpression, requestId, jsonRequest, requesterIdentity, requestContext, httpHeaders); - } - } - - /** - * Handles the subquery in the given expression. - *

When subquery is detected, first executes the subquery and gets the response, then rewrites the expression with - * the subquery response. - *

Currently only supports ID_SET subquery within the IN_SUBQUERY transform function, which will be rewritten to an - * IN_ID_SET transform function. - */ - private void handleSubquery(Expression expression, long requestId, JsonNode jsonRequest, - @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders) - throws Exception { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - List operands = function.getOperands(); - if (function.getOperator().equals(IN_SUBQUERY)) { - Preconditions.checkState(operands.size() == 2, "IN_SUBQUERY requires 2 arguments: expression, subquery"); - Literal subqueryLiteral = operands.get(1).getLiteral(); - Preconditions.checkState(subqueryLiteral != null, "Second argument of IN_SUBQUERY must be a literal (subquery)"); - String subquery = subqueryLiteral.getStringValue(); - BrokerResponse response = - handleRequest(requestId, subquery, null, jsonRequest, requesterIdentity, requestContext, httpHeaders); - if (response.getExceptionsSize() != 0) { - throw new RuntimeException("Caught exception while executing subquery: " + subquery); - } - String serializedIdSet = (String) response.getResultTable().getRows().get(0)[0]; - function.setOperator(IN_ID_SET); - operands.set(1, RequestUtils.getLiteralExpression(serializedIdSet)); - } else { - for (Expression operand : operands) { - handleSubquery(operand, requestId, jsonRequest, requesterIdentity, requestContext, httpHeaders); - } - } - } - - /** - * Resolves the actual table name for: - * - Case-insensitive cluster - * - * @param tableName the table name in the query - * @param tableCache the table case-sensitive cache - * @return table name if the table name is found in Pinot registry. - */ - @VisibleForTesting - static String getActualTableName(String tableName, TableCache tableCache) { - String actualTableName = tableCache.getActualTableName(tableName); - if (actualTableName != null) { - return actualTableName; - } - return tableName; - } - - /** - * Retrieve segment partitioned columns for a table. - * For a hybrid table, a segment partitioned column has to be the intersection of both offline and realtime tables. - * - * @param tableCache - * @param tableName - * @return segment partitioned columns belong to both offline and realtime tables. - */ - private static Set getSegmentPartitionedColumns(TableCache tableCache, String tableName) { - final TableConfig offlineTableConfig = - tableCache.getTableConfig(TableNameBuilder.OFFLINE.tableNameWithType(tableName)); - final TableConfig realtimeTableConfig = - tableCache.getTableConfig(TableNameBuilder.REALTIME.tableNameWithType(tableName)); - if (offlineTableConfig == null) { - return getSegmentPartitionedColumns(realtimeTableConfig); - } - if (realtimeTableConfig == null) { - return getSegmentPartitionedColumns(offlineTableConfig); - } - Set segmentPartitionedColumns = getSegmentPartitionedColumns(offlineTableConfig); - segmentPartitionedColumns.retainAll(getSegmentPartitionedColumns(realtimeTableConfig)); - return segmentPartitionedColumns; - } - - private static Set getSegmentPartitionedColumns(@Nullable TableConfig tableConfig) { - Set segmentPartitionedColumns = new HashSet<>(); - if (tableConfig == null) { - return segmentPartitionedColumns; - } - List fieldConfigs = tableConfig.getFieldConfigList(); - if (fieldConfigs != null) { - for (FieldConfig fieldConfig : fieldConfigs) { - if (fieldConfig.getProperties() != null && Boolean.parseBoolean( - fieldConfig.getProperties().get(FieldConfig.IS_SEGMENT_PARTITIONED_COLUMN_KEY))) { - segmentPartitionedColumns.add(fieldConfig.getName()); - } - } - } - return segmentPartitionedColumns; - } - - /** - * Retrieve multivalued columns for a table. - * From the table Schema , we get the multi valued columns of dimension fields. - * - * @param tableSchema - * @param columnName - * @return multivalued columns of the table . - */ - private static boolean isMultiValueColumn(Schema tableSchema, String columnName) { - - DimensionFieldSpec dimensionFieldSpec = tableSchema.getDimensionSpec(columnName); - return dimensionFieldSpec != null && !dimensionFieldSpec.isSingleValueField(); - } - - /** - * Sets the table name in the given broker request. - * NOTE: Set table name in broker request because it is used for access control, query routing etc. - */ - private void setTableName(BrokerRequest brokerRequest, String tableName) { - brokerRequest.getQuerySource().setTableName(tableName); - brokerRequest.getPinotQuery().getDataSource().setTableName(tableName); - } - - /** - * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given query. - */ - private static void handleHLLLog2mOverride(PinotQuery pinotQuery, int hllLog2mOverride) { - List selectList = pinotQuery.getSelectList(); - for (Expression expression : selectList) { - handleHLLLog2mOverride(expression, hllLog2mOverride); - } - List orderByList = pinotQuery.getOrderByList(); - if (orderByList != null) { - for (Expression expression : orderByList) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleHLLLog2mOverride(expression.getFunctionCall().getOperands().get(0), hllLog2mOverride); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleHLLLog2mOverride(havingExpression, hllLog2mOverride); - } - } - - /** - * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given expression. - */ - private static void handleHLLLog2mOverride(Expression expression, int hllLog2mOverride) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - switch (function.getOperator()) { - case "distinctcounthll": - case "distinctcounthllmv": - case "distinctcountrawhll": - case "distinctcountrawhllmv": - if (function.getOperandsSize() == 1) { - function.addToOperands(RequestUtils.getLiteralExpression(hllLog2mOverride)); - } - return; - default: - break; - } - for (Expression operand : function.getOperands()) { - handleHLLLog2mOverride(operand, hllLog2mOverride); - } - } - - /** - * Overrides the LIMIT of the given query if it exceeds the query limit. - */ - @VisibleForTesting - static void handleQueryLimitOverride(PinotQuery pinotQuery, int queryLimit) { - if (queryLimit > 0 && pinotQuery.getLimit() > queryLimit) { - pinotQuery.setLimit(queryLimit); - } - } - - /** - * Rewrites 'DistinctCount' to 'SegmentPartitionDistinctCount' for the given query. - */ - @VisibleForTesting - static void handleSegmentPartitionedDistinctCountOverride(PinotQuery pinotQuery, - Set segmentPartitionedColumns) { - if (segmentPartitionedColumns.isEmpty()) { - return; - } - for (Expression expression : pinotQuery.getSelectList()) { - handleSegmentPartitionedDistinctCountOverride(expression, segmentPartitionedColumns); - } - List orderByExpressions = pinotQuery.getOrderByList(); - if (orderByExpressions != null) { - for (Expression expression : orderByExpressions) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleSegmentPartitionedDistinctCountOverride(expression.getFunctionCall().getOperands().get(0), - segmentPartitionedColumns); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleSegmentPartitionedDistinctCountOverride(havingExpression, segmentPartitionedColumns); - } - } - - /** - * Rewrites 'DistinctCount' to 'SegmentPartitionDistinctCount' for the given expression. - */ - private static void handleSegmentPartitionedDistinctCountOverride(Expression expression, - Set segmentPartitionedColumns) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - if (function.getOperator().equals("distinctcount")) { - List operands = function.getOperands(); - if (operands.size() == 1 && operands.get(0).isSetIdentifier() && segmentPartitionedColumns.contains( - operands.get(0).getIdentifier().getName())) { - function.setOperator("segmentpartitioneddistinctcount"); - } - } else { - for (Expression operand : function.getOperands()) { - handleSegmentPartitionedDistinctCountOverride(operand, segmentPartitionedColumns); - } - } - } - - /** - * Rewrites 'DistinctCount' to 'DistinctCountBitmap' for the given query. - */ - private static void handleDistinctCountBitmapOverride(PinotQuery pinotQuery) { - for (Expression expression : pinotQuery.getSelectList()) { - handleDistinctCountBitmapOverride(expression); - } - List orderByExpressions = pinotQuery.getOrderByList(); - if (orderByExpressions != null) { - for (Expression expression : orderByExpressions) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleDistinctCountBitmapOverride(expression.getFunctionCall().getOperands().get(0)); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleDistinctCountBitmapOverride(havingExpression); - } - } - - /** - * Rewrites selected 'Distinct' prefixed function to 'Distinct----MV' function for the field of multivalued type. - */ - @VisibleForTesting - static void handleDistinctMultiValuedOverride(PinotQuery pinotQuery, Schema tableSchema) { - for (Expression expression : pinotQuery.getSelectList()) { - handleDistinctMultiValuedOverride(expression, tableSchema); - } - List orderByExpressions = pinotQuery.getOrderByList(); - if (orderByExpressions != null) { - for (Expression expression : orderByExpressions) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleDistinctMultiValuedOverride(expression.getFunctionCall().getOperands().get(0), tableSchema); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleDistinctMultiValuedOverride(havingExpression, tableSchema); - } - } - - /** - * Rewrites selected 'Distinct' prefixed function to 'Distinct----MV' function for the field of multivalued type. - */ - private static void handleDistinctMultiValuedOverride(Expression expression, Schema tableSchema) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - - String overrideOperator = DISTINCT_MV_COL_FUNCTION_OVERRIDE_MAP.get(function.getOperator()); - if (overrideOperator != null) { - List operands = function.getOperands(); - if (operands.size() >= 1 && operands.get(0).isSetIdentifier() && isMultiValueColumn(tableSchema, - operands.get(0).getIdentifier().getName())) { - // we are only checking the first operand that if its a MV column as all the overriding agg. fn.'s have - // first operator is column name - function.setOperator(overrideOperator); - } - } else { - for (Expression operand : function.getOperands()) { - handleDistinctMultiValuedOverride(operand, tableSchema); - } - } - } - - /** - * Rewrites 'DistinctCount' to 'DistinctCountBitmap' for the given expression. - */ - private static void handleDistinctCountBitmapOverride(Expression expression) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - if (function.getOperator().equals("distinctcount")) { - function.setOperator("distinctcountbitmap"); - } else { - for (Expression operand : function.getOperands()) { - handleDistinctCountBitmapOverride(operand); - } - } - } - - private HandlerContext getHandlerContext(@Nullable TableConfig offlineTableConfig, - @Nullable TableConfig realtimeTableConfig) { - Boolean disableGroovyOverride = null; - Boolean useApproximateFunctionOverride = null; - if (offlineTableConfig != null && offlineTableConfig.getQueryConfig() != null) { - QueryConfig offlineTableQueryConfig = offlineTableConfig.getQueryConfig(); - Boolean disableGroovyOfflineTableOverride = offlineTableQueryConfig.getDisableGroovy(); - if (disableGroovyOfflineTableOverride != null) { - disableGroovyOverride = disableGroovyOfflineTableOverride; - } - Boolean useApproximateFunctionOfflineTableOverride = offlineTableQueryConfig.getUseApproximateFunction(); - if (useApproximateFunctionOfflineTableOverride != null) { - useApproximateFunctionOverride = useApproximateFunctionOfflineTableOverride; - } - } - if (realtimeTableConfig != null && realtimeTableConfig.getQueryConfig() != null) { - QueryConfig realtimeTableQueryConfig = realtimeTableConfig.getQueryConfig(); - Boolean disableGroovyRealtimeTableOverride = realtimeTableQueryConfig.getDisableGroovy(); - if (disableGroovyRealtimeTableOverride != null) { - if (disableGroovyOverride == null) { - disableGroovyOverride = disableGroovyRealtimeTableOverride; - } else { - // Disable Groovy if either offline or realtime table config disables Groovy - disableGroovyOverride |= disableGroovyRealtimeTableOverride; - } - } - Boolean useApproximateFunctionRealtimeTableOverride = realtimeTableQueryConfig.getUseApproximateFunction(); - if (useApproximateFunctionRealtimeTableOverride != null) { - if (useApproximateFunctionOverride == null) { - useApproximateFunctionOverride = useApproximateFunctionRealtimeTableOverride; - } else { - // Use approximate function if both offline and realtime table config uses approximate function - useApproximateFunctionOverride &= useApproximateFunctionRealtimeTableOverride; - } - } - } - - boolean disableGroovy = disableGroovyOverride != null ? disableGroovyOverride : _disableGroovy; - boolean useApproximateFunction = - useApproximateFunctionOverride != null ? useApproximateFunctionOverride : _useApproximateFunction; - return new HandlerContext(disableGroovy, useApproximateFunction); - } - - private static class HandlerContext { - final boolean _disableGroovy; - final boolean _useApproximateFunction; - - HandlerContext(boolean disableGroovy, boolean useApproximateFunction) { - _disableGroovy = disableGroovy; - _useApproximateFunction = useApproximateFunction; - } - } - - /** - * Verifies that no groovy is present in the PinotQuery when disabled. - */ - @VisibleForTesting - static void rejectGroovyQuery(PinotQuery pinotQuery) { - List selectList = pinotQuery.getSelectList(); - for (Expression expression : selectList) { - rejectGroovyQuery(expression); - } - List orderByList = pinotQuery.getOrderByList(); - if (orderByList != null) { - for (Expression expression : orderByList) { - // NOTE: Order-by is always a Function with the ordering of the Expression - rejectGroovyQuery(expression.getFunctionCall().getOperands().get(0)); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - rejectGroovyQuery(havingExpression); - } - Expression filterExpression = pinotQuery.getFilterExpression(); - if (filterExpression != null) { - rejectGroovyQuery(filterExpression); - } - List groupByList = pinotQuery.getGroupByList(); - if (groupByList != null) { - for (Expression expression : groupByList) { - rejectGroovyQuery(expression); - } - } - } - - private static void rejectGroovyQuery(Expression expression) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - if (function.getOperator().equals("groovy")) { - throw new BadQueryRequestException("Groovy transform functions are disabled for queries"); - } - for (Expression operandExpression : function.getOperands()) { - rejectGroovyQuery(operandExpression); - } - } - - /** - * Rewrites potential expensive functions to their approximation counterparts. - * - DISTINCT_COUNT -> DISTINCT_COUNT_SMART_HLL - * - PERCENTILE -> PERCENTILE_SMART_TDIGEST - */ - @VisibleForTesting - static void handleApproximateFunctionOverride(PinotQuery pinotQuery) { - for (Expression expression : pinotQuery.getSelectList()) { - handleApproximateFunctionOverride(expression); - } - List orderByExpressions = pinotQuery.getOrderByList(); - if (orderByExpressions != null) { - for (Expression expression : orderByExpressions) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleApproximateFunctionOverride(expression.getFunctionCall().getOperands().get(0)); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleApproximateFunctionOverride(havingExpression); - } - } - - private static void handleApproximateFunctionOverride(Expression expression) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - String functionName = function.getOperator(); - if (functionName.equals("distinctcount") || functionName.equals("distinctcountmv")) { - function.setOperator("distinctcountsmarthll"); - } else if (functionName.startsWith("percentile")) { - String remainingFunctionName = functionName.substring(10); - if (remainingFunctionName.isEmpty() || remainingFunctionName.equals("mv")) { - function.setOperator("percentilesmarttdigest"); - } else if (remainingFunctionName.matches("\\d+")) { - try { - int percentile = Integer.parseInt(remainingFunctionName); - function.setOperator("percentilesmarttdigest"); - function.addToOperands(RequestUtils.getLiteralExpression(percentile)); - } catch (Exception e) { - throw new BadQueryRequestException("Illegal function name: " + functionName); - } - } else if (remainingFunctionName.matches("\\d+mv")) { - try { - int percentile = Integer.parseInt(remainingFunctionName.substring(0, remainingFunctionName.length() - 2)); - function.setOperator("percentilesmarttdigest"); - function.addToOperands(RequestUtils.getLiteralExpression(percentile)); - } catch (Exception e) { - throw new BadQueryRequestException("Illegal function name: " + functionName); - } - } - } else { - for (Expression operand : function.getOperands()) { - handleApproximateFunctionOverride(operand); - } - } - } - - private static void handleExpressionOverride(PinotQuery pinotQuery, - @Nullable Map expressionOverrideMap) { - if (expressionOverrideMap == null) { - return; - } - pinotQuery.getSelectList().replaceAll(o -> handleExpressionOverride(o, expressionOverrideMap)); - Expression filterExpression = pinotQuery.getFilterExpression(); - if (filterExpression != null) { - pinotQuery.setFilterExpression(handleExpressionOverride(filterExpression, expressionOverrideMap)); - } - List groupByList = pinotQuery.getGroupByList(); - if (groupByList != null) { - groupByList.replaceAll(o -> handleExpressionOverride(o, expressionOverrideMap)); - } - List orderByList = pinotQuery.getOrderByList(); - if (orderByList != null) { - for (Expression expression : orderByList) { - // NOTE: Order-by is always a Function with the ordering of the Expression - expression.getFunctionCall().getOperands().replaceAll(o -> handleExpressionOverride(o, expressionOverrideMap)); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - pinotQuery.setHavingExpression(handleExpressionOverride(havingExpression, expressionOverrideMap)); - } - } - - private static Expression handleExpressionOverride(Expression expression, - Map expressionOverrideMap) { - Expression override = expressionOverrideMap.get(expression); - if (override != null) { - return new Expression(override); - } - Function function = expression.getFunctionCall(); - if (function != null) { - function.getOperands().replaceAll(o -> handleExpressionOverride(o, expressionOverrideMap)); - } - return expression; - } - - /** - * Returns {@code true} if the given query only contains literals, {@code false} otherwise. - */ - @VisibleForTesting - static boolean isLiteralOnlyQuery(PinotQuery pinotQuery) { - for (Expression expression : pinotQuery.getSelectList()) { - if (!CalciteSqlParser.isLiteralOnlyExpression(expression)) { - return false; - } - } - return true; - } - - /** - * Processes the literal only query. - */ - private BrokerResponseNative processLiteralOnlyQuery(PinotQuery pinotQuery, long compilationStartTimeNs, - RequestContext requestContext) { - BrokerResponseNative brokerResponse = new BrokerResponseNative(); - List columnNames = new ArrayList<>(); - List columnTypes = new ArrayList<>(); - List row = new ArrayList<>(); - for (Expression expression : pinotQuery.getSelectList()) { - computeResultsForExpression(expression, columnNames, columnTypes, row); - } - DataSchema dataSchema = - new DataSchema(columnNames.toArray(new String[0]), columnTypes.toArray(new DataSchema.ColumnDataType[0])); - List rows = new ArrayList<>(); - rows.add(row.toArray()); - ResultTable resultTable = new ResultTable(dataSchema, rows); - brokerResponse.setResultTable(resultTable); - - long totalTimeMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - compilationStartTimeNs); - brokerResponse.setTimeUsedMs(totalTimeMs); - requestContext.setQueryProcessingTime(totalTimeMs); - augmentStatistics(requestContext, brokerResponse); - return brokerResponse; - } - - // TODO(xiangfu): Move Literal function computation here from Calcite Parser. - private void computeResultsForExpression(Expression e, List columnNames, - List columnTypes, List row) { - if (e.getType() == ExpressionType.LITERAL) { - computeResultsForLiteral(e.getLiteral(), columnNames, columnTypes, row); - } - if (e.getType() == ExpressionType.FUNCTION) { - if (e.getFunctionCall().getOperator().equals("as")) { - String columnName = e.getFunctionCall().getOperands().get(1).getIdentifier().getName(); - computeResultsForExpression(e.getFunctionCall().getOperands().get(0), columnNames, columnTypes, row); - columnNames.set(columnNames.size() - 1, columnName); - } else { - throw new IllegalStateException( - "No able to compute results for function - " + e.getFunctionCall().getOperator()); - } - } - } - - private void computeResultsForLiteral(Literal literal, List columnNames, - List columnTypes, List row) { - Object fieldValue = literal.getFieldValue(); - columnNames.add(fieldValue.toString()); - switch (literal.getSetField()) { - case BOOL_VALUE: - columnTypes.add(DataSchema.ColumnDataType.BOOLEAN); - row.add(literal.getBoolValue()); - break; - case BYTE_VALUE: - columnTypes.add(DataSchema.ColumnDataType.INT); - row.add((int) literal.getByteValue()); - break; - case SHORT_VALUE: - columnTypes.add(DataSchema.ColumnDataType.INT); - row.add((int) literal.getShortValue()); - break; - case INT_VALUE: - columnTypes.add(DataSchema.ColumnDataType.INT); - row.add(literal.getIntValue()); - break; - case LONG_VALUE: - columnTypes.add(DataSchema.ColumnDataType.LONG); - row.add(literal.getLongValue()); - break; - case FLOAT_VALUE: - columnTypes.add(DataSchema.ColumnDataType.FLOAT); - row.add(Float.intBitsToFloat(literal.getFloatValue())); - break; - case DOUBLE_VALUE: - columnTypes.add(DataSchema.ColumnDataType.DOUBLE); - row.add(literal.getDoubleValue()); - break; - case BIG_DECIMAL_VALUE: - columnTypes.add(DataSchema.ColumnDataType.BIG_DECIMAL); - row.add(BigDecimalUtils.deserialize(literal.getBigDecimalValue())); - break; - case STRING_VALUE: - columnTypes.add(DataSchema.ColumnDataType.STRING); - row.add(literal.getStringValue()); - break; - case BINARY_VALUE: - columnTypes.add(DataSchema.ColumnDataType.BYTES); - row.add(BytesUtils.toHexString(literal.getBinaryValue())); - break; - case NULL_VALUE: - columnTypes.add(DataSchema.ColumnDataType.UNKNOWN); - row.add(null); - break; - case INT_ARRAY_VALUE: - columnTypes.add(DataSchema.ColumnDataType.INT_ARRAY); - row.add(literal.getIntArrayValue()); - break; - case LONG_ARRAY_VALUE: - columnTypes.add(DataSchema.ColumnDataType.LONG_ARRAY); - row.add(literal.getLongArrayValue()); - break; - case FLOAT_ARRAY_VALUE: - columnTypes.add(DataSchema.ColumnDataType.FLOAT_ARRAY); - row.add(literal.getFloatArrayValue().stream().map(Float::intBitsToFloat).collect(Collectors.toList())); - break; - case DOUBLE_ARRAY_VALUE: - columnTypes.add(DataSchema.ColumnDataType.DOUBLE_ARRAY); - row.add(literal.getDoubleArrayValue()); - break; - case STRING_ARRAY_VALUE: - columnTypes.add(DataSchema.ColumnDataType.STRING_ARRAY); - row.add(literal.getStringArrayValue()); - break; - default: - throw new IllegalStateException("Unsupported literal: " + literal); - } - } - - /** - * Fixes the column names to the actual column names in the given query. - */ - @VisibleForTesting - static void updateColumnNames(String rawTableName, PinotQuery pinotQuery, boolean isCaseInsensitive, - Map columnNameMap) { - if (pinotQuery != null) { - boolean hasStar = false; - for (Expression expression : pinotQuery.getSelectList()) { - fixColumnName(rawTableName, expression, columnNameMap, isCaseInsensitive); - //check if the select expression is '*' - if (!hasStar && expression.equals(STAR)) { - hasStar = true; - } - } - //if query has a '*' selection along with other columns - if (hasStar) { - expandStarExpressionsToActualColumns(pinotQuery, columnNameMap); - } - Expression filterExpression = pinotQuery.getFilterExpression(); - if (filterExpression != null) { - // We don't support alias in filter expression, so we don't need to pass aliasMap - fixColumnName(rawTableName, filterExpression, columnNameMap, isCaseInsensitive); - } - List groupByList = pinotQuery.getGroupByList(); - if (groupByList != null) { - for (Expression expression : groupByList) { - fixColumnName(rawTableName, expression, columnNameMap, isCaseInsensitive); - } - } - List orderByList = pinotQuery.getOrderByList(); - if (orderByList != null) { - for (Expression expression : orderByList) { - // NOTE: Order-by is always a Function with the ordering of the Expression - fixColumnName(rawTableName, expression.getFunctionCall().getOperands().get(0), columnNameMap, - isCaseInsensitive); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - fixColumnName(rawTableName, havingExpression, columnNameMap, isCaseInsensitive); - } - } - } - - private static void expandStarExpressionsToActualColumns(PinotQuery pinotQuery, Map columnNameMap) { - List originalSelections = pinotQuery.getSelectList(); - //expand '*' - List expandedSelections = new ArrayList<>(); - for (String tableCol : columnNameMap.values()) { - Expression newSelection = RequestUtils.getIdentifierExpression(tableCol); - //we exclude default virtual columns - if (tableCol.charAt(0) != '$') { - expandedSelections.add(newSelection); - } - } - //sort naturally - expandedSelections.sort(null); - List newSelections = new ArrayList<>(); - for (Expression originalSelection : originalSelections) { - if (originalSelection.equals(STAR)) { - newSelections.addAll(expandedSelections); - } else { - newSelections.add(originalSelection); - } - } - pinotQuery.setSelectList(newSelections); - } - - /** - * Fixes the column names to the actual column names in the given expression. - */ - private static void fixColumnName(String rawTableName, Expression expression, Map columnNameMap, - boolean ignoreCase) { - ExpressionType expressionType = expression.getType(); - if (expressionType == ExpressionType.IDENTIFIER) { - Identifier identifier = expression.getIdentifier(); - identifier.setName(getActualColumnName(rawTableName, identifier.getName(), columnNameMap, ignoreCase)); - } else if (expressionType == ExpressionType.FUNCTION) { - final Function functionCall = expression.getFunctionCall(); - switch (functionCall.getOperator()) { - case "as": - fixColumnName(rawTableName, functionCall.getOperands().get(0), columnNameMap, ignoreCase); - break; - case "lookup": - // LOOKUP function looks up another table's schema, skip the check for now. - break; - default: - for (Expression operand : functionCall.getOperands()) { - fixColumnName(rawTableName, operand, columnNameMap, ignoreCase); - } - break; - } - } - } - - /** - * Returns the actual column name for the given column name for: - * - Case-insensitive cluster - * - Column name in the format of [{@code rawTableName}].[column_name] - * - Column name in the format of [logical_table_name].[column_name] while {@code rawTableName} is a translated name - */ - @VisibleForTesting - static String getActualColumnName(String rawTableName, String columnName, @Nullable Map columnNameMap, - boolean ignoreCase) { - if ("*".equals(columnName)) { - return columnName; - } - String columnNameToCheck = trimTableName(rawTableName, columnName, ignoreCase); - if (ignoreCase) { - columnNameToCheck = columnNameToCheck.toLowerCase(); - } - if (columnNameMap != null) { - String actualColumnName = columnNameMap.get(columnNameToCheck); - if (actualColumnName != null) { - return actualColumnName; - } - } - if (columnName.charAt(0) == '$') { - return columnName; - } - throw new BadQueryRequestException("Unknown columnName '" + columnName + "' found in the query"); - } - - private static String trimTableName(String rawTableName, String columnName, boolean ignoreCase) { - int columnNameLength = columnName.length(); - int rawTableNameLength = rawTableName.length(); - if (columnNameLength > rawTableNameLength && columnName.charAt(rawTableNameLength) == '.' - && columnName.regionMatches(ignoreCase, 0, rawTableName, 0, rawTableNameLength)) { - return columnName.substring(rawTableNameLength + 1); - } - // Check if raw table name is translated name ([database_name].[logical_table_name]]) - String[] split = StringUtils.split(rawTableName, '.'); - if (split.length == 2) { - String logicalTableName = split[1]; - int logicalTableNameLength = logicalTableName.length(); - if (columnNameLength > logicalTableNameLength && columnName.charAt(logicalTableNameLength) == '.' - && columnName.regionMatches(ignoreCase, 0, logicalTableName, 0, logicalTableNameLength)) { - return columnName.substring(logicalTableNameLength + 1); - } - } - return columnName; - } - - /** - * Helper function to decide whether to force the log - * - * TODO: come up with other criteria for forcing a log and come up with better numbers - */ - private boolean forceLog(BrokerResponse brokerResponse, long totalTimeMs) { - if (brokerResponse.isNumGroupsLimitReached()) { - return true; - } - - if (brokerResponse.getExceptionsSize() > 0) { - return true; - } - - // If response time is more than 1 sec, force the log - return totalTimeMs > 1000L; - } - - /** - * Sets the query timeout (remaining time in milliseconds) into the query options, and returns the remaining time in - * milliseconds. - *

For the overall query timeout, use query-level timeout (in the query options) if exists, or use table-level - * timeout (in the table config) if exists, or use instance-level timeout (in the broker config). - */ - private long setQueryTimeout(String tableNameWithType, Map queryOptions, long timeSpentMs) - throws TimeoutException { - long queryTimeoutMs; - Long queryLevelTimeoutMs = QueryOptionsUtils.getTimeoutMs(queryOptions); - if (queryLevelTimeoutMs != null) { - // Use query-level timeout if exists - queryTimeoutMs = queryLevelTimeoutMs; - } else { - Long tableLevelTimeoutMs = _routingManager.getQueryTimeoutMs(tableNameWithType); - if (tableLevelTimeoutMs != null) { - // Use table-level timeout if exists - queryTimeoutMs = tableLevelTimeoutMs; - } else { - // Use instance-level timeout - queryTimeoutMs = _brokerTimeoutMs; - } - } - - long remainingTimeMs = queryTimeoutMs - timeSpentMs; - if (remainingTimeMs <= 0) { - String errorMessage = - String.format("Query timed out (time spent: %dms, timeout: %dms) for table: %s before scattering the request", - timeSpentMs, queryTimeoutMs, tableNameWithType); - throw new TimeoutException(errorMessage); - } - queryOptions.put(QueryOptionKey.TIMEOUT_MS, Long.toString(remainingTimeMs)); - return remainingTimeMs; - } - - /** - * Sets a query option indicating the maximum response size that can be sent from a server to the broker. This size - * is measured for the serialized response. - * - * The overriding order of priority is: - * 1. QueryOption -> maxServerResponseSizeBytes - * 2. QueryOption -> maxQueryResponseSizeBytes - * 3. TableConfig -> maxServerResponseSizeBytes - * 4. TableConfig -> maxQueryResponseSizeBytes - * 5. BrokerConfig -> maxServerResponseSizeBytes - * 6. BrokerConfig -> maxServerResponseSizeBytes - */ - private void setMaxServerResponseSizeBytes(int numServers, Map queryOptions, - @Nullable TableConfig tableConfig) { - // QueryOption - if (QueryOptionsUtils.getMaxServerResponseSizeBytes(queryOptions) != null) { - return; - } - Long maxQueryResponseSizeQueryOption = QueryOptionsUtils.getMaxQueryResponseSizeBytes(queryOptions); - if (maxQueryResponseSizeQueryOption != null) { - queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, - Long.toString(maxQueryResponseSizeQueryOption / numServers)); - return; - } - - // TableConfig - if (tableConfig != null && tableConfig.getQueryConfig() != null) { - QueryConfig queryConfig = tableConfig.getQueryConfig(); - if (queryConfig.getMaxServerResponseSizeBytes() != null) { - queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, - Long.toString(queryConfig.getMaxServerResponseSizeBytes())); - return; - } - if (queryConfig.getMaxQueryResponseSizeBytes() != null) { - queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, - Long.toString(queryConfig.getMaxQueryResponseSizeBytes() / numServers)); - return; - } - } - - // BrokerConfig - String maxServerResponseSizeBrokerConfig = _config.getProperty(Broker.CONFIG_OF_MAX_SERVER_RESPONSE_SIZE_BYTES); - if (maxServerResponseSizeBrokerConfig != null) { - queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, - Long.toString(DataSizeUtils.toBytes(maxServerResponseSizeBrokerConfig))); - return; - } - - String maxQueryResponseSizeBrokerConfig = _config.getProperty(Broker.CONFIG_OF_MAX_QUERY_RESPONSE_SIZE_BYTES); - if (maxQueryResponseSizeBrokerConfig != null) { - queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, - Long.toString(DataSizeUtils.toBytes(maxQueryResponseSizeBrokerConfig) / numServers)); - } - } - - /** - * Broker side validation on the query. - *

Throw exception if query does not pass validation. - *

Current validations are: - *

    - *
  • Value for 'LIMIT' <= configured value
    • - *
  • Query options must be set to SQL mode
    • - *
  • Check if numReplicaGroupsToQuery option provided is valid
    • - *
- */ - @VisibleForTesting - static void validateRequest(PinotQuery pinotQuery, int queryResponseLimit) { - // Verify LIMIT - int limit = pinotQuery.getLimit(); - if (limit > queryResponseLimit) { - throw new IllegalStateException( - "Value for 'LIMIT' (" + limit + ") exceeds maximum allowed value of " + queryResponseLimit); - } - - Map queryOptions = pinotQuery.getQueryOptions(); - try { - // throw errors if options is less than 1 or invalid - Integer numReplicaGroupsToQuery = QueryOptionsUtils.getNumReplicaGroupsToQuery(queryOptions); - if (numReplicaGroupsToQuery != null) { - Preconditions.checkState(numReplicaGroupsToQuery > 0, "numReplicaGroups must be " + "positive number, got: %d", - numReplicaGroupsToQuery); - } - } catch (NumberFormatException e) { - String numReplicaGroupsToQuery = queryOptions.get(QueryOptionKey.NUM_REPLICA_GROUPS_TO_QUERY); - throw new IllegalStateException( - String.format("numReplicaGroups must be a positive number, got: %s", numReplicaGroupsToQuery)); - } - - if (pinotQuery.getDataSource().getSubquery() != null) { - validateRequest(pinotQuery.getDataSource().getSubquery(), queryResponseLimit); - } - } - - /** - * Helper method to attach the time boundary to the given PinotQuery. - */ - private static void attachTimeBoundary(PinotQuery pinotQuery, TimeBoundaryInfo timeBoundaryInfo, - boolean isOfflineRequest) { - String functionName = isOfflineRequest ? FilterKind.LESS_THAN_OR_EQUAL.name() : FilterKind.GREATER_THAN.name(); - String timeColumn = timeBoundaryInfo.getTimeColumn(); - String timeValue = timeBoundaryInfo.getTimeValue(); - Expression timeFilterExpression = - RequestUtils.getFunctionExpression(functionName, RequestUtils.getIdentifierExpression(timeColumn), - RequestUtils.getLiteralExpression(timeValue)); - - Expression filterExpression = pinotQuery.getFilterExpression(); - if (filterExpression != null) { - pinotQuery.setFilterExpression( - RequestUtils.getFunctionExpression(FilterKind.AND.name(), filterExpression, timeFilterExpression)); - } else { - pinotQuery.setFilterExpression(timeFilterExpression); - } - } - - /** - * Processes the optimized broker requests for both OFFLINE and REALTIME table. - * TODO: Directly take PinotQuery - */ - protected abstract BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest, - BrokerRequest serverBrokerRequest, @Nullable BrokerRequest offlineBrokerRequest, - @Nullable Map, List>> offlineRoutingTable, - @Nullable BrokerRequest realtimeBrokerRequest, - @Nullable Map, List>> realtimeRoutingTable, long timeoutMs, - ServerStats serverStats, RequestContext requestContext) + protected abstract BrokerResponse handleRequest(long requestId, String query, + @Nullable SqlNodeAndOptions sqlNodeAndOptions, JsonNode request, @Nullable RequesterIdentity requesterIdentity, + RequestContext requestContext, @Nullable HttpHeaders httpHeaders, AccessControl accessControl) throws Exception; - protected static boolean isPartialResult(BrokerResponse brokerResponse) { - return brokerResponse.isNumGroupsLimitReached() || brokerResponse.isMaxRowsInJoinReached() - || brokerResponse.getExceptionsSize() > 0; - } - protected static void augmentStatistics(RequestContext statistics, BrokerResponse response) { - statistics.setTotalDocs(response.getTotalDocs()); + statistics.setNumRowsResultSet(response.getNumRowsResultSet()); + // TODO: Add partial result flag to RequestContext + List exceptions = response.getExceptions(); + int numExceptions = exceptions.size(); + List processingExceptions = new ArrayList<>(numExceptions); + for (QueryProcessingException exception : exceptions) { + processingExceptions.add(exception.toString()); + } + statistics.setProcessingExceptions(processingExceptions); + statistics.setNumExceptions(numExceptions); + statistics.setNumGroupsLimitReached(response.isNumGroupsLimitReached()); + statistics.setProcessingTimeMillis(response.getTimeUsedMs()); statistics.setNumDocsScanned(response.getNumDocsScanned()); + statistics.setTotalDocs(response.getTotalDocs()); statistics.setNumEntriesScannedInFilter(response.getNumEntriesScannedInFilter()); statistics.setNumEntriesScannedPostFilter(response.getNumEntriesScannedPostFilter()); + statistics.setNumServersQueried(response.getNumServersQueried()); + statistics.setNumServersResponded(response.getNumServersResponded()); statistics.setNumSegmentsQueried(response.getNumSegmentsQueried()); statistics.setNumSegmentsProcessed(response.getNumSegmentsProcessed()); statistics.setNumSegmentsMatched(response.getNumSegmentsMatched()); - statistics.setNumServersQueried(response.getNumServersQueried()); - statistics.setNumSegmentsProcessed(response.getNumSegmentsProcessed()); - statistics.setNumServersResponded(response.getNumServersResponded()); - statistics.setNumGroupsLimitReached(response.isNumGroupsLimitReached()); - statistics.setNumExceptions(response.getExceptionsSize()); - statistics.setOfflineThreadCpuTimeNs(response.getOfflineThreadCpuTimeNs()); - statistics.setRealtimeThreadCpuTimeNs(response.getRealtimeThreadCpuTimeNs()); - statistics.setOfflineSystemActivitiesCpuTimeNs(response.getOfflineSystemActivitiesCpuTimeNs()); - statistics.setRealtimeSystemActivitiesCpuTimeNs(response.getRealtimeSystemActivitiesCpuTimeNs()); - statistics.setOfflineResponseSerializationCpuTimeNs(response.getOfflineResponseSerializationCpuTimeNs()); - statistics.setRealtimeResponseSerializationCpuTimeNs(response.getRealtimeResponseSerializationCpuTimeNs()); - statistics.setOfflineTotalCpuTimeNs(response.getOfflineTotalCpuTimeNs()); - statistics.setRealtimeTotalCpuTimeNs(response.getRealtimeTotalCpuTimeNs()); - statistics.setNumRowsResultSet(response.getNumRowsResultSet()); statistics.setNumConsumingSegmentsQueried(response.getNumConsumingSegmentsQueried()); statistics.setNumConsumingSegmentsProcessed(response.getNumConsumingSegmentsProcessed()); statistics.setNumConsumingSegmentsMatched(response.getNumConsumingSegmentsMatched()); @@ -1910,47 +175,17 @@ protected static void augmentStatistics(RequestContext statistics, BrokerRespons statistics.setNumSegmentsPrunedInvalid(response.getNumSegmentsPrunedInvalid()); statistics.setNumSegmentsPrunedByLimit(response.getNumSegmentsPrunedByLimit()); statistics.setNumSegmentsPrunedByValue(response.getNumSegmentsPrunedByValue()); + statistics.setReduceTimeMillis(response.getBrokerReduceTimeMs()); + statistics.setOfflineThreadCpuTimeNs(response.getOfflineThreadCpuTimeNs()); + statistics.setRealtimeThreadCpuTimeNs(response.getRealtimeThreadCpuTimeNs()); + statistics.setOfflineSystemActivitiesCpuTimeNs(response.getOfflineSystemActivitiesCpuTimeNs()); + statistics.setRealtimeSystemActivitiesCpuTimeNs(response.getRealtimeSystemActivitiesCpuTimeNs()); + statistics.setOfflineResponseSerializationCpuTimeNs(response.getOfflineResponseSerializationCpuTimeNs()); + statistics.setRealtimeResponseSerializationCpuTimeNs(response.getRealtimeResponseSerializationCpuTimeNs()); + statistics.setOfflineTotalCpuTimeNs(response.getOfflineTotalCpuTimeNs()); + statistics.setRealtimeTotalCpuTimeNs(response.getRealtimeTotalCpuTimeNs()); statistics.setExplainPlanNumEmptyFilterSegments(response.getExplainPlanNumEmptyFilterSegments()); statistics.setExplainPlanNumMatchAllFilterSegments(response.getExplainPlanNumMatchAllFilterSegments()); - statistics.setProcessingExceptions( - response.getProcessingExceptions().stream().map(Object::toString).collect(Collectors.toList())); - } - - private String getGlobalQueryId(long requestId) { - return _brokerId + "_" + requestId; - } - - /** - * Helper class to pass the per server statistics. - */ - public static class ServerStats { - private String _serverStats; - - public String getServerStats() { - return _serverStats; - } - - public void setServerStats(String serverStats) { - _serverStats = serverStats; - } - } - - /** - * Helper class to track the query plaintext and the requested servers. - */ - private static class QueryServers { - final String _query; - final Set _servers = new HashSet<>(); - - QueryServers(String query, @Nullable Map, List>> offlineRoutingTable, - @Nullable Map, List>> realtimeRoutingTable) { - _query = query; - if (offlineRoutingTable != null) { - _servers.addAll(offlineRoutingTable.keySet()); - } - if (realtimeRoutingTable != null) { - _servers.addAll(realtimeRoutingTable.keySet()); - } - } + statistics.setTraceInfo(response.getTraceInfo()); } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java new file mode 100644 index 000000000000..5d6ee5bcc302 --- /dev/null +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -0,0 +1,1850 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.broker.requesthandler; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableMap; +import java.net.URI; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletionService; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import javax.annotation.concurrent.ThreadSafe; +import javax.ws.rs.WebApplicationException; +import javax.ws.rs.core.HttpHeaders; +import javax.ws.rs.core.Response; +import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.lang3.StringUtils; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.http.client.methods.HttpDelete; +import org.apache.http.conn.HttpClientConnectionManager; +import org.apache.http.util.EntityUtils; +import org.apache.pinot.broker.api.AccessControl; +import org.apache.pinot.broker.api.RequesterIdentity; +import org.apache.pinot.broker.broker.AccessControlFactory; +import org.apache.pinot.broker.querylog.QueryLogger; +import org.apache.pinot.broker.queryquota.QueryQuotaManager; +import org.apache.pinot.broker.routing.BrokerRoutingManager; +import org.apache.pinot.common.config.provider.TableCache; +import org.apache.pinot.common.exception.QueryException; +import org.apache.pinot.common.http.MultiHttpRequest; +import org.apache.pinot.common.http.MultiHttpRequestResponse; +import org.apache.pinot.common.metrics.BrokerGauge; +import org.apache.pinot.common.metrics.BrokerMeter; +import org.apache.pinot.common.metrics.BrokerQueryPhase; +import org.apache.pinot.common.metrics.BrokerTimer; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.common.request.DataSource; +import org.apache.pinot.common.request.Expression; +import org.apache.pinot.common.request.ExpressionType; +import org.apache.pinot.common.request.Function; +import org.apache.pinot.common.request.Identifier; +import org.apache.pinot.common.request.Literal; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.response.BrokerResponse; +import org.apache.pinot.common.response.ProcessingException; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DatabaseUtils; +import org.apache.pinot.common.utils.config.QueryOptionsUtils; +import org.apache.pinot.common.utils.request.RequestUtils; +import org.apache.pinot.core.auth.Actions; +import org.apache.pinot.core.auth.TargetType; +import org.apache.pinot.core.query.optimizer.QueryOptimizer; +import org.apache.pinot.core.routing.RoutingTable; +import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.transport.ServerInstance; +import org.apache.pinot.core.util.GapfillUtils; +import org.apache.pinot.spi.auth.AuthorizationResult; +import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.config.table.QueryConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.exception.BadQueryRequestException; +import org.apache.pinot.spi.exception.DatabaseConflictException; +import org.apache.pinot.spi.trace.RequestContext; +import org.apache.pinot.spi.trace.Tracing; +import org.apache.pinot.spi.utils.BigDecimalUtils; +import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.Broker; +import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; +import org.apache.pinot.spi.utils.DataSizeUtils; +import org.apache.pinot.spi.utils.TimestampIndexUtils; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.apache.pinot.sql.FilterKind; +import org.apache.pinot.sql.parsers.CalciteSqlCompiler; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.apache.pinot.sql.parsers.SqlNodeAndOptions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +@ThreadSafe +public abstract class BaseSingleStageBrokerRequestHandler extends BaseBrokerRequestHandler { + private static final Logger LOGGER = LoggerFactory.getLogger(BaseSingleStageBrokerRequestHandler.class); + private static final String IN_SUBQUERY = "insubquery"; + private static final String IN_ID_SET = "inidset"; + private static final Expression FALSE = RequestUtils.getLiteralExpression(false); + private static final Expression TRUE = RequestUtils.getLiteralExpression(true); + private static final Expression STAR = RequestUtils.getIdentifierExpression("*"); + private static final int MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION = 10; + private static final Map DISTINCT_MV_COL_FUNCTION_OVERRIDE_MAP = + ImmutableMap.builder().put("distinctcount", "distinctcountmv") + .put("distinctcountbitmap", "distinctcountbitmapmv").put("distinctcounthll", "distinctcounthllmv") + .put("distinctcountrawhll", "distinctcountrawhllmv").put("distinctsum", "distinctsummv") + .put("distinctavg", "distinctavgmv").put("count", "countmv").put("min", "minmv").put("max", "maxmv") + .put("avg", "avgmv").put("sum", "summv").put("minmaxrange", "minmaxrangemv") + .put("distinctcounthllplus", "distinctcounthllplusmv") + .put("distinctcountrawhllplus", "distinctcountrawhllplusmv").build(); + + protected final QueryOptimizer _queryOptimizer = new QueryOptimizer(); + protected final boolean _disableGroovy; + protected final boolean _useApproximateFunction; + protected final int _defaultHllLog2m; + protected final boolean _enableQueryLimitOverride; + protected final boolean _enableDistinctCountBitmapOverride; + protected final int _queryResponseLimit; + protected final Map _queriesById; + + public BaseSingleStageBrokerRequestHandler(PinotConfiguration config, String brokerId, + BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, + QueryQuotaManager queryQuotaManager, TableCache tableCache) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); + _disableGroovy = _config.getProperty(Broker.DISABLE_GROOVY, Broker.DEFAULT_DISABLE_GROOVY); + _useApproximateFunction = _config.getProperty(Broker.USE_APPROXIMATE_FUNCTION, false); + _defaultHllLog2m = _config.getProperty(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M_KEY, + CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M); + _enableQueryLimitOverride = _config.getProperty(Broker.CONFIG_OF_ENABLE_QUERY_LIMIT_OVERRIDE, false); + _enableDistinctCountBitmapOverride = + _config.getProperty(CommonConstants.Helix.ENABLE_DISTINCT_COUNT_BITMAP_OVERRIDE_KEY, false); + _queryResponseLimit = + config.getProperty(Broker.CONFIG_OF_BROKER_QUERY_RESPONSE_LIMIT, Broker.DEFAULT_BROKER_QUERY_RESPONSE_LIMIT); + boolean enableQueryCancellation = + Boolean.parseBoolean(config.getProperty(Broker.CONFIG_OF_BROKER_ENABLE_QUERY_CANCELLATION)); + _queriesById = enableQueryCancellation ? new ConcurrentHashMap<>() : null; + LOGGER.info("Initialized {} with broker id: {}, timeout: {}ms, query response limit: {}, query log max length: {}, " + + "query log max rate: {}, query cancellation enabled: {}", getClass().getSimpleName(), _brokerId, + _brokerTimeoutMs, _queryResponseLimit, _queryLogger.getMaxQueryLengthToLog(), _queryLogger.getLogRateLimit(), + enableQueryCancellation); + } + + @Override + public Map getRunningQueries() { + Preconditions.checkState(_queriesById != null, "Query cancellation is not enabled on broker"); + return _queriesById.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue()._query)); + } + + @VisibleForTesting + Set getRunningServers(long requestId) { + Preconditions.checkState(_queriesById != null, "Query cancellation is not enabled on broker"); + QueryServers queryServers = _queriesById.get(requestId); + return queryServers != null ? queryServers._servers : Collections.emptySet(); + } + + @Override + public boolean cancelQuery(long requestId, int timeoutMs, Executor executor, HttpClientConnectionManager connMgr, + Map serverResponses) + throws Exception { + Preconditions.checkState(_queriesById != null, "Query cancellation is not enabled on broker"); + QueryServers queryServers = _queriesById.get(requestId); + if (queryServers == null) { + return false; + } + // TODO: Use different global query id for OFFLINE and REALTIME table after releasing 0.12.0. See QueryIdUtils for + // details + String globalQueryId = getGlobalQueryId(requestId); + List> serverUrls = new ArrayList<>(); + for (ServerInstance serverInstance : queryServers._servers) { + serverUrls.add(Pair.of(String.format("%s/query/%s", serverInstance.getAdminEndpoint(), globalQueryId), null)); + } + LOGGER.debug("Cancelling the query: {} via server urls: {}", queryServers._query, serverUrls); + CompletionService completionService = + new MultiHttpRequest(executor, connMgr).execute(serverUrls, null, timeoutMs, "DELETE", HttpDelete::new); + List errMsgs = new ArrayList<>(serverUrls.size()); + for (int i = 0; i < serverUrls.size(); i++) { + MultiHttpRequestResponse httpRequestResponse = null; + try { + // Wait for all requests to respond before returning to be sure that the servers have handled the cancel + // requests. The completion order is different from serverUrls, thus use uri in the response. + httpRequestResponse = completionService.take().get(); + URI uri = httpRequestResponse.getURI(); + int status = httpRequestResponse.getResponse().getStatusLine().getStatusCode(); + // Unexpected server responses are collected and returned as exception. + if (status != 200 && status != 404) { + String responseString = EntityUtils.toString(httpRequestResponse.getResponse().getEntity()); + throw new Exception( + String.format("Unexpected status=%d and response='%s' from uri='%s'", status, responseString, uri)); + } + if (serverResponses != null) { + serverResponses.put(uri.getHost() + ":" + uri.getPort(), status); + } + } catch (Exception e) { + LOGGER.error("Failed to cancel query: {}", queryServers._query, e); + // Can't just throw exception from here as there is a need to release the other connections. + // So just collect the error msg to throw them together after the for-loop. + errMsgs.add(e.getMessage()); + } finally { + if (httpRequestResponse != null) { + httpRequestResponse.close(); + } + } + } + if (errMsgs.size() > 0) { + throw new Exception("Unexpected responses from servers: " + StringUtils.join(errMsgs, ",")); + } + return true; + } + + @Override + protected BrokerResponse handleRequest(long requestId, String query, @Nullable SqlNodeAndOptions sqlNodeAndOptions, + JsonNode request, @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, + @Nullable HttpHeaders httpHeaders, AccessControl accessControl) + throws Exception { + LOGGER.debug("SQL query for request {}: {}", requestId, query); + + //Start instrumentation context. This must not be moved further below interspersed into the code. + Tracing.ThreadAccountantOps.setupRunner(String.valueOf(requestId)); + + try { + // Parse the query if needed + if (sqlNodeAndOptions == null) { + try { + sqlNodeAndOptions = RequestUtils.parseQuery(query, request); + } catch (Exception e) { + // Do not log or emit metric here because it is pure user error + requestContext.setErrorCode(QueryException.SQL_PARSING_ERROR_CODE); + return new BrokerResponseNative(QueryException.getException(QueryException.SQL_PARSING_ERROR, e)); + } + } + + // Compile the request into PinotQuery + long compilationStartTimeNs = System.nanoTime(); + PinotQuery pinotQuery; + try { + pinotQuery = CalciteSqlParser.compileToPinotQuery(sqlNodeAndOptions); + } catch (Exception e) { + LOGGER.info("Caught exception while compiling SQL request {}: {}, {}", requestId, query, e.getMessage()); + _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_COMPILATION_EXCEPTIONS, 1); + requestContext.setErrorCode(QueryException.SQL_PARSING_ERROR_CODE); + return new BrokerResponseNative(QueryException.getException(QueryException.SQL_PARSING_ERROR, e)); + } + + if (isLiteralOnlyQuery(pinotQuery)) { + LOGGER.debug("Request {} contains only Literal, skipping server query: {}", requestId, query); + try { + if (pinotQuery.isExplain()) { + // EXPLAIN PLAN results to show that query is evaluated exclusively by Broker. + return BrokerResponseNative.BROKER_ONLY_EXPLAIN_PLAN_OUTPUT; + } + return processLiteralOnlyQuery(requestId, pinotQuery, requestContext); + } catch (Exception e) { + // TODO: refine the exceptions here to early termination the queries won't requires to send to servers. + LOGGER.warn("Unable to execute literal request {}: {} at broker, fallback to server query. {}", requestId, + query, e.getMessage()); + } + } + + PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery); + DataSource dataSource = serverPinotQuery.getDataSource(); + if (dataSource == null) { + LOGGER.info("Data source (FROM clause) not found in request {}: {}", requestId, query); + requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); + return new BrokerResponseNative( + QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, "Data source (FROM clause) not found")); + } + if (dataSource.getJoin() != null) { + LOGGER.info("JOIN is not supported in request {}: {}", requestId, query); + requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); + return new BrokerResponseNative( + QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, "JOIN is not supported")); + } + if (dataSource.getTableName() == null) { + LOGGER.info("Table name not found in request {}: {}", requestId, query); + requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); + return new BrokerResponseNative( + QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, "Table name not found")); + } + + try { + handleSubquery(serverPinotQuery, requestId, request, requesterIdentity, requestContext, httpHeaders, + accessControl); + } catch (Exception e) { + LOGGER.info("Caught exception while handling the subquery in request {}: {}, {}", requestId, query, + e.getMessage()); + requestContext.setErrorCode(QueryException.QUERY_EXECUTION_ERROR_CODE); + return new BrokerResponseNative(QueryException.getException(QueryException.QUERY_EXECUTION_ERROR, e)); + } + + boolean ignoreCase = _tableCache.isIgnoreCase(); + String tableName; + try { + tableName = + getActualTableName(DatabaseUtils.translateTableName(dataSource.getTableName(), httpHeaders, ignoreCase), + _tableCache); + } catch (DatabaseConflictException e) { + LOGGER.info("{}. Request {}: {}", e.getMessage(), requestId, query); + _brokerMetrics.addMeteredGlobalValue(BrokerMeter.QUERY_VALIDATION_EXCEPTIONS, 1); + requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); + return new BrokerResponseNative(QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, e)); + } + dataSource.setTableName(tableName); + String rawTableName = TableNameBuilder.extractRawTableName(tableName); + requestContext.setTableName(rawTableName); + + try { + Map columnNameMap = _tableCache.getColumnNameMap(rawTableName); + if (columnNameMap != null) { + updateColumnNames(rawTableName, serverPinotQuery, ignoreCase, columnNameMap); + } + } catch (Exception e) { + // Throw exceptions with column in-existence error. + if (e instanceof BadQueryRequestException) { + LOGGER.info("Caught exception while checking column names in request {}: {}, {}", requestId, query, + e.getMessage()); + requestContext.setErrorCode(QueryException.UNKNOWN_COLUMN_ERROR_CODE); + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.UNKNOWN_COLUMN_EXCEPTIONS, 1); + return new BrokerResponseNative(QueryException.getException(QueryException.UNKNOWN_COLUMN_ERROR, e)); + } + LOGGER.warn("Caught exception while updating column names in request {}: {}, {}", requestId, query, + e.getMessage()); + } + if (_defaultHllLog2m > 0) { + handleHLLLog2mOverride(serverPinotQuery, _defaultHllLog2m); + } + if (_enableQueryLimitOverride) { + handleQueryLimitOverride(serverPinotQuery, _queryResponseLimit); + } + handleSegmentPartitionedDistinctCountOverride(serverPinotQuery, + getSegmentPartitionedColumns(_tableCache, tableName)); + if (_enableDistinctCountBitmapOverride) { + handleDistinctCountBitmapOverride(serverPinotQuery); + } + + Schema schema = _tableCache.getSchema(rawTableName); + if (schema != null) { + handleDistinctMultiValuedOverride(serverPinotQuery, schema); + } + + long compilationEndTimeNs = System.nanoTime(); + // full request compile time = compilationTimeNs + parserTimeNs + _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.REQUEST_COMPILATION, + (compilationEndTimeNs - compilationStartTimeNs) + sqlNodeAndOptions.getParseTimeNs()); + + // Second-stage table-level access control + // TODO: Modify AccessControl interface to directly take PinotQuery + BrokerRequest brokerRequest = CalciteSqlCompiler.convertToBrokerRequest(pinotQuery); + BrokerRequest serverBrokerRequest = + serverPinotQuery == pinotQuery ? brokerRequest : CalciteSqlCompiler.convertToBrokerRequest(serverPinotQuery); + AuthorizationResult authorizationResult = accessControl.authorize(requesterIdentity, serverBrokerRequest); + if (authorizationResult.hasAccess()) { + authorizationResult = accessControl.authorize(httpHeaders, TargetType.TABLE, tableName, Actions.Table.QUERY); + } + + _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.AUTHORIZATION, + System.nanoTime() - compilationEndTimeNs); + + if (!authorizationResult.hasAccess()) { + _brokerMetrics.addMeteredTableValue(tableName, BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1); + LOGGER.info("Access denied for request {}: {}, table: {}, reason :{}", requestId, query, tableName, + authorizationResult.getFailureMessage()); + requestContext.setErrorCode(QueryException.ACCESS_DENIED_ERROR_CODE); + String failureMessage = authorizationResult.getFailureMessage(); + if (StringUtils.isNotBlank(failureMessage)) { + failureMessage = "Reason: " + failureMessage; + } + throw new WebApplicationException("Permission denied." + failureMessage, + Response.Status.FORBIDDEN); + } + + // Get the tables hit by the request + String offlineTableName = null; + String realtimeTableName = null; + TableType tableType = TableNameBuilder.getTableTypeFromTableName(tableName); + if (tableType == TableType.OFFLINE) { + // Offline table + if (_routingManager.routingExists(tableName)) { + offlineTableName = tableName; + } + } else if (tableType == TableType.REALTIME) { + // Realtime table + if (_routingManager.routingExists(tableName)) { + realtimeTableName = tableName; + } + } else { + // Hybrid table (check both OFFLINE and REALTIME) + String offlineTableNameToCheck = TableNameBuilder.OFFLINE.tableNameWithType(tableName); + if (_routingManager.routingExists(offlineTableNameToCheck)) { + offlineTableName = offlineTableNameToCheck; + } + String realtimeTableNameToCheck = TableNameBuilder.REALTIME.tableNameWithType(tableName); + if (_routingManager.routingExists(realtimeTableNameToCheck)) { + realtimeTableName = realtimeTableNameToCheck; + } + } + + TableConfig offlineTableConfig = + _tableCache.getTableConfig(TableNameBuilder.OFFLINE.tableNameWithType(rawTableName)); + TableConfig realtimeTableConfig = + _tableCache.getTableConfig(TableNameBuilder.REALTIME.tableNameWithType(rawTableName)); + + if (offlineTableName == null && realtimeTableName == null) { + // No table matches the request + if (realtimeTableConfig == null && offlineTableConfig == null) { + LOGGER.info("Table not found for request {}: {}", requestId, query); + requestContext.setErrorCode(QueryException.TABLE_DOES_NOT_EXIST_ERROR_CODE); + return BrokerResponseNative.TABLE_DOES_NOT_EXIST; + } + LOGGER.info("No table matches for request {}: {}", requestId, query); + requestContext.setErrorCode(QueryException.BROKER_RESOURCE_MISSING_ERROR_CODE); + _brokerMetrics.addMeteredGlobalValue(BrokerMeter.RESOURCE_MISSING_EXCEPTIONS, 1); + return BrokerResponseNative.NO_TABLE_RESULT; + } + + // Handle query rewrite that can be overridden by the table configs + if (offlineTableName == null) { + offlineTableConfig = null; + } + if (realtimeTableName == null) { + realtimeTableConfig = null; + } + HandlerContext handlerContext = getHandlerContext(offlineTableConfig, realtimeTableConfig); + if (handlerContext._disableGroovy) { + rejectGroovyQuery(serverPinotQuery); + } + if (handlerContext._useApproximateFunction) { + handleApproximateFunctionOverride(serverPinotQuery); + } + + // Validate QPS quota + if (!_queryQuotaManager.acquire(tableName)) { + String errorMessage = + String.format("Request %d: %s exceeds query quota for table: %s", requestId, query, tableName); + LOGGER.info(errorMessage); + requestContext.setErrorCode(QueryException.TOO_MANY_REQUESTS_ERROR_CODE); + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.QUERY_QUOTA_EXCEEDED, 1); + return new BrokerResponseNative(QueryException.getException(QueryException.QUOTA_EXCEEDED_ERROR, errorMessage)); + } + + // Validate the request + try { + validateRequest(serverPinotQuery, _queryResponseLimit); + } catch (Exception e) { + LOGGER.info("Caught exception while validating request {}: {}, {}", requestId, query, e.getMessage()); + requestContext.setErrorCode(QueryException.QUERY_VALIDATION_ERROR_CODE); + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.QUERY_VALIDATION_EXCEPTIONS, 1); + return new BrokerResponseNative(QueryException.getException(QueryException.QUERY_VALIDATION_ERROR, e)); + } + + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.QUERIES, 1); + _brokerMetrics.addValueToTableGauge(rawTableName, BrokerGauge.REQUEST_SIZE, query.length()); + + // Prepare OFFLINE and REALTIME requests + BrokerRequest offlineBrokerRequest = null; + BrokerRequest realtimeBrokerRequest = null; + TimeBoundaryInfo timeBoundaryInfo = null; + if (offlineTableName != null && realtimeTableName != null) { + // Time boundary info might be null when there is no segment in the offline table, query real-time side only + timeBoundaryInfo = _routingManager.getTimeBoundaryInfo(offlineTableName); + if (timeBoundaryInfo == null) { + LOGGER.debug("No time boundary info found for hybrid table: {}", rawTableName); + offlineTableName = null; + } + } + if (offlineTableName != null && realtimeTableName != null) { + // Hybrid + PinotQuery offlinePinotQuery = serverPinotQuery.deepCopy(); + offlinePinotQuery.getDataSource().setTableName(offlineTableName); + attachTimeBoundary(offlinePinotQuery, timeBoundaryInfo, true); + handleExpressionOverride(offlinePinotQuery, _tableCache.getExpressionOverrideMap(offlineTableName)); + handleTimestampIndexOverride(offlinePinotQuery, offlineTableConfig); + _queryOptimizer.optimize(offlinePinotQuery, offlineTableConfig, schema); + offlineBrokerRequest = CalciteSqlCompiler.convertToBrokerRequest(offlinePinotQuery); + + PinotQuery realtimePinotQuery = serverPinotQuery.deepCopy(); + realtimePinotQuery.getDataSource().setTableName(realtimeTableName); + attachTimeBoundary(realtimePinotQuery, timeBoundaryInfo, false); + handleExpressionOverride(realtimePinotQuery, _tableCache.getExpressionOverrideMap(realtimeTableName)); + handleTimestampIndexOverride(realtimePinotQuery, realtimeTableConfig); + _queryOptimizer.optimize(realtimePinotQuery, realtimeTableConfig, schema); + realtimeBrokerRequest = CalciteSqlCompiler.convertToBrokerRequest(realtimePinotQuery); + + requestContext.setFanoutType(RequestContext.FanoutType.HYBRID); + requestContext.setOfflineServerTenant(getServerTenant(offlineTableName)); + requestContext.setRealtimeServerTenant(getServerTenant(realtimeTableName)); + } else if (offlineTableName != null) { + // OFFLINE only + setTableName(serverBrokerRequest, offlineTableName); + handleExpressionOverride(serverPinotQuery, _tableCache.getExpressionOverrideMap(offlineTableName)); + handleTimestampIndexOverride(serverPinotQuery, offlineTableConfig); + _queryOptimizer.optimize(serverPinotQuery, offlineTableConfig, schema); + offlineBrokerRequest = serverBrokerRequest; + + requestContext.setFanoutType(RequestContext.FanoutType.OFFLINE); + requestContext.setOfflineServerTenant(getServerTenant(offlineTableName)); + } else { + // REALTIME only + setTableName(serverBrokerRequest, realtimeTableName); + handleExpressionOverride(serverPinotQuery, _tableCache.getExpressionOverrideMap(realtimeTableName)); + handleTimestampIndexOverride(serverPinotQuery, realtimeTableConfig); + _queryOptimizer.optimize(serverPinotQuery, realtimeTableConfig, schema); + realtimeBrokerRequest = serverBrokerRequest; + + requestContext.setFanoutType(RequestContext.FanoutType.REALTIME); + requestContext.setRealtimeServerTenant(getServerTenant(realtimeTableName)); + } + + // Check if response can be sent without server query evaluation. + if (offlineBrokerRequest != null && isFilterAlwaysFalse(offlineBrokerRequest.getPinotQuery())) { + // We don't need to evaluate offline request + offlineBrokerRequest = null; + } + if (realtimeBrokerRequest != null && isFilterAlwaysFalse(realtimeBrokerRequest.getPinotQuery())) { + // We don't need to evaluate realtime request + realtimeBrokerRequest = null; + } + + if (offlineBrokerRequest == null && realtimeBrokerRequest == null) { + return getEmptyBrokerOnlyResponse(pinotQuery, requestContext, tableName, requesterIdentity); + } + + if (offlineBrokerRequest != null && isFilterAlwaysTrue(offlineBrokerRequest.getPinotQuery())) { + // Drop offline request filter since it is always true + offlineBrokerRequest.getPinotQuery().setFilterExpression(null); + } + if (realtimeBrokerRequest != null && isFilterAlwaysTrue(realtimeBrokerRequest.getPinotQuery())) { + // Drop realtime request filter since it is always true + realtimeBrokerRequest.getPinotQuery().setFilterExpression(null); + } + + // Calculate routing table for the query + // TODO: Modify RoutingManager interface to directly take PinotQuery + long routingStartTimeNs = System.nanoTime(); + Map, List>> offlineRoutingTable = null; + Map, List>> realtimeRoutingTable = null; + List unavailableSegments = new ArrayList<>(); + int numPrunedSegmentsTotal = 0; + if (offlineBrokerRequest != null) { + // NOTE: Routing table might be null if table is just removed + RoutingTable routingTable = _routingManager.getRoutingTable(offlineBrokerRequest, requestId); + if (routingTable != null) { + unavailableSegments.addAll(routingTable.getUnavailableSegments()); + Map, List>> serverInstanceToSegmentsMap = + routingTable.getServerInstanceToSegmentsMap(); + if (!serverInstanceToSegmentsMap.isEmpty()) { + offlineRoutingTable = serverInstanceToSegmentsMap; + } else { + offlineBrokerRequest = null; + } + numPrunedSegmentsTotal += routingTable.getNumPrunedSegments(); + } else { + offlineBrokerRequest = null; + } + } + if (realtimeBrokerRequest != null) { + // NOTE: Routing table might be null if table is just removed + RoutingTable routingTable = _routingManager.getRoutingTable(realtimeBrokerRequest, requestId); + if (routingTable != null) { + unavailableSegments.addAll(routingTable.getUnavailableSegments()); + Map, List>> serverInstanceToSegmentsMap = + routingTable.getServerInstanceToSegmentsMap(); + if (!serverInstanceToSegmentsMap.isEmpty()) { + realtimeRoutingTable = serverInstanceToSegmentsMap; + } else { + realtimeBrokerRequest = null; + } + numPrunedSegmentsTotal += routingTable.getNumPrunedSegments(); + } else { + realtimeBrokerRequest = null; + } + } + int numUnavailableSegments = unavailableSegments.size(); + requestContext.setNumUnavailableSegments(numUnavailableSegments); + + List exceptions = new ArrayList<>(); + if (numUnavailableSegments > 0) { + String errorMessage; + if (numUnavailableSegments > MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION) { + errorMessage = String.format("%d segments unavailable, sampling %d: %s", numUnavailableSegments, + MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION, + unavailableSegments.subList(0, MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION)); + } else { + errorMessage = String.format("%d segments unavailable: %s", numUnavailableSegments, unavailableSegments); + } + exceptions.add(QueryException.getException(QueryException.BROKER_SEGMENT_UNAVAILABLE_ERROR, errorMessage)); + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_UNAVAILABLE_SEGMENTS, 1); + } + + if (offlineBrokerRequest == null && realtimeBrokerRequest == null) { + if (!exceptions.isEmpty()) { + LOGGER.info("No server found for request {}: {}", requestId, query); + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.NO_SERVER_FOUND_EXCEPTIONS, 1); + return new BrokerResponseNative(exceptions); + } else { + // When all segments have been pruned, we can just return an empty response. + return getEmptyBrokerOnlyResponse(pinotQuery, requestContext, tableName, requesterIdentity); + } + } + long routingEndTimeNs = System.nanoTime(); + _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.QUERY_ROUTING, + routingEndTimeNs - routingStartTimeNs); + + // Set timeout in the requests + long timeSpentMs = TimeUnit.NANOSECONDS.toMillis(routingEndTimeNs - compilationStartTimeNs); + // Remaining time in milliseconds for the server query execution + // NOTE: For hybrid use case, in most cases offline table and real-time table should have the same query timeout + // configured, but if necessary, we also allow different timeout for them. + // If the timeout is not the same for offline table and real-time table, use the max of offline table + // remaining time and realtime table remaining time. Server side will have different remaining time set for + // each table type, and broker should wait for both types to return. + long remainingTimeMs = 0; + try { + if (offlineBrokerRequest != null) { + remainingTimeMs = + setQueryTimeout(offlineTableName, offlineBrokerRequest.getPinotQuery().getQueryOptions(), timeSpentMs); + } + if (realtimeBrokerRequest != null) { + remainingTimeMs = Math.max(remainingTimeMs, + setQueryTimeout(realtimeTableName, realtimeBrokerRequest.getPinotQuery().getQueryOptions(), timeSpentMs)); + } + } catch (TimeoutException e) { + String errorMessage = e.getMessage(); + LOGGER.info("{} {}: {}", errorMessage, requestId, query); + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.REQUEST_TIMEOUT_BEFORE_SCATTERED_EXCEPTIONS, 1); + exceptions.add(QueryException.getException(QueryException.BROKER_TIMEOUT_ERROR, errorMessage)); + return new BrokerResponseNative(exceptions); + } + + // Set the maximum serialized response size per server, and ask server to directly return final response when only + // one server is queried + int numServers = 0; + if (offlineRoutingTable != null) { + numServers += offlineRoutingTable.size(); + } + if (realtimeRoutingTable != null) { + numServers += realtimeRoutingTable.size(); + } + if (offlineBrokerRequest != null) { + Map queryOptions = offlineBrokerRequest.getPinotQuery().getQueryOptions(); + setMaxServerResponseSizeBytes(numServers, queryOptions, offlineTableConfig); + // Set the query option to directly return final result for single server query unless it is explicitly disabled + if (numServers == 1) { + // Set the same flag in the original server request to be used in the reduce phase for hybrid table + if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null + && offlineBrokerRequest != serverBrokerRequest) { + serverBrokerRequest.getPinotQuery().getQueryOptions() + .put(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true"); + } + } + } + if (realtimeBrokerRequest != null) { + Map queryOptions = realtimeBrokerRequest.getPinotQuery().getQueryOptions(); + setMaxServerResponseSizeBytes(numServers, queryOptions, realtimeTableConfig); + // Set the query option to directly return final result for single server query unless it is explicitly disabled + if (numServers == 1) { + // Set the same flag in the original server request to be used in the reduce phase for hybrid table + if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null + && realtimeBrokerRequest != serverBrokerRequest) { + serverBrokerRequest.getPinotQuery().getQueryOptions() + .put(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true"); + } + } + } + + // Execute the query + // TODO: Replace ServerStats with ServerRoutingStatsEntry. + ServerStats serverStats = new ServerStats(); + // TODO: Handle broker specific operations for explain plan queries such as: + // - Alias handling + // - Compile time function invocation + // - Literal only queries + // - Any rewrites + if (pinotQuery.isExplain()) { + // Update routing tables to only send request to offline servers for OFFLINE and HYBRID tables. + // TODO: Assess if the Explain Plan Query should also be routed to REALTIME servers for HYBRID tables + if (offlineRoutingTable != null) { + // For OFFLINE and HYBRID tables, don't send EXPLAIN query to realtime servers. + realtimeBrokerRequest = null; + realtimeRoutingTable = null; + } + } + BrokerResponseNative brokerResponse; + if (_queriesById != null) { + // Start to track the running query for cancellation just before sending it out to servers to avoid any + // potential failures that could happen before sending it out, like failures to calculate the routing table etc. + // TODO: Even tracking the query as late as here, a potential race condition between calling cancel API and + // query being sent out to servers can still happen. If cancel request arrives earlier than query being + // sent out to servers, the servers miss the cancel request and continue to run the queries. The users + // can always list the running queries and cancel query again until it ends. Just that such race + // condition makes cancel API less reliable. This should be rare as it assumes sending queries out to + // servers takes time, but will address later if needed. + _queriesById.put(requestId, new QueryServers(query, offlineRoutingTable, realtimeRoutingTable)); + LOGGER.debug("Keep track of running query: {}", requestId); + try { + brokerResponse = processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, offlineBrokerRequest, + offlineRoutingTable, realtimeBrokerRequest, realtimeRoutingTable, remainingTimeMs, serverStats, + requestContext); + } finally { + _queriesById.remove(requestId); + LOGGER.debug("Remove track of running query: {}", requestId); + } + } else { + brokerResponse = processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, offlineBrokerRequest, + offlineRoutingTable, realtimeBrokerRequest, realtimeRoutingTable, remainingTimeMs, serverStats, + requestContext); + } + + for (ProcessingException exception : exceptions) { + brokerResponse.addException(exception); + } + brokerResponse.setNumSegmentsPrunedByBroker(numPrunedSegmentsTotal); + long executionEndTimeNs = System.nanoTime(); + _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.QUERY_EXECUTION, + executionEndTimeNs - routingEndTimeNs); + + // Track number of queries with number of groups limit reached + if (brokerResponse.isNumGroupsLimitReached()) { + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_NUM_GROUPS_LIMIT_REACHED, + 1); + } + + // Set total query processing time + long totalTimeMs = System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis(); + brokerResponse.setTimeUsedMs(totalTimeMs); + augmentStatistics(requestContext, brokerResponse); + if (QueryOptionsUtils.shouldDropResults(pinotQuery.getQueryOptions())) { + brokerResponse.setResultTable(null); + } + _brokerMetrics.addTimedTableValue(rawTableName, BrokerTimer.QUERY_TOTAL_TIME_MS, totalTimeMs, + TimeUnit.MILLISECONDS); + + // Log query and stats + _queryLogger.log( + new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse, requesterIdentity, serverStats)); + + return brokerResponse; + } finally { + Tracing.ThreadAccountantOps.clear(); + } + } + + private BrokerResponseNative getEmptyBrokerOnlyResponse(PinotQuery pinotQuery, RequestContext requestContext, + String tableName, @Nullable RequesterIdentity requesterIdentity) { + if (pinotQuery.isExplain()) { + // EXPLAIN PLAN results to show that query is evaluated exclusively by Broker. + return BrokerResponseNative.BROKER_ONLY_EXPLAIN_PLAN_OUTPUT; + } + + // Send empty response since we don't need to evaluate either offline or realtime request. + BrokerResponseNative brokerResponse = BrokerResponseNative.empty(); + brokerResponse.setTimeUsedMs(System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis()); + _queryLogger.log( + new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse, requesterIdentity, null)); + return brokerResponse; + } + + private void handleTimestampIndexOverride(PinotQuery pinotQuery, @Nullable TableConfig tableConfig) { + if (tableConfig == null || tableConfig.getFieldConfigList() == null) { + return; + } + + Set timestampIndexColumns = _tableCache.getTimestampIndexColumns(tableConfig.getTableName()); + if (CollectionUtils.isEmpty(timestampIndexColumns)) { + return; + } + for (Expression expression : pinotQuery.getSelectList()) { + setTimestampIndexExpressionOverrideHints(expression, timestampIndexColumns, pinotQuery); + } + setTimestampIndexExpressionOverrideHints(pinotQuery.getFilterExpression(), timestampIndexColumns, pinotQuery); + setTimestampIndexExpressionOverrideHints(pinotQuery.getHavingExpression(), timestampIndexColumns, pinotQuery); + List groupByList = pinotQuery.getGroupByList(); + if (CollectionUtils.isNotEmpty(groupByList)) { + groupByList.forEach( + expression -> setTimestampIndexExpressionOverrideHints(expression, timestampIndexColumns, pinotQuery)); + } + List orderByList = pinotQuery.getOrderByList(); + if (CollectionUtils.isNotEmpty(orderByList)) { + orderByList.forEach( + expression -> setTimestampIndexExpressionOverrideHints(expression, timestampIndexColumns, pinotQuery)); + } + } + + private void setTimestampIndexExpressionOverrideHints(@Nullable Expression expression, + Set timestampIndexColumns, PinotQuery pinotQuery) { + if (expression == null || expression.getFunctionCall() == null) { + return; + } + Function function = expression.getFunctionCall(); + switch (function.getOperator()) { + case "datetrunc": + String granularString = function.getOperands().get(0).getLiteral().getStringValue().toUpperCase(); + Expression timeExpression = function.getOperands().get(1); + if (((function.getOperandsSize() == 2) || (function.getOperandsSize() == 3 && "MILLISECONDS".equalsIgnoreCase( + function.getOperands().get(2).getLiteral().getStringValue()))) && TimestampIndexUtils.isValidGranularity( + granularString) && timeExpression.getIdentifier() != null) { + String timeColumn = timeExpression.getIdentifier().getName(); + String timeColumnWithGranularity = TimestampIndexUtils.getColumnWithGranularity(timeColumn, granularString); + if (timestampIndexColumns.contains(timeColumnWithGranularity)) { + pinotQuery.putToExpressionOverrideHints(expression, + RequestUtils.getIdentifierExpression(timeColumnWithGranularity)); + } + } + break; + default: + break; + } + function.getOperands() + .forEach(operand -> setTimestampIndexExpressionOverrideHints(operand, timestampIndexColumns, pinotQuery)); + } + + /** Given a {@link PinotQuery}, check if the WHERE clause will always evaluate to false. */ + private boolean isFilterAlwaysFalse(PinotQuery pinotQuery) { + return FALSE.equals(pinotQuery.getFilterExpression()); + } + + /** Given a {@link PinotQuery}, check if the WHERE clause will always evaluate to true. */ + private boolean isFilterAlwaysTrue(PinotQuery pinotQuery) { + return TRUE.equals(pinotQuery.getFilterExpression()); + } + + private String getServerTenant(String tableNameWithType) { + TableConfig tableConfig = _tableCache.getTableConfig(tableNameWithType); + if (tableConfig == null) { + LOGGER.debug("Table config is not available for table {}", tableNameWithType); + return "unknownTenant"; + } + return tableConfig.getTenantConfig().getServer(); + } + + /** + * Handles the subquery in the given query. + *

Currently only supports subquery within the filter. + */ + private void handleSubquery(PinotQuery pinotQuery, long requestId, JsonNode jsonRequest, + @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders, + AccessControl accessControl) + throws Exception { + Expression filterExpression = pinotQuery.getFilterExpression(); + if (filterExpression != null) { + handleSubquery(filterExpression, requestId, jsonRequest, requesterIdentity, requestContext, httpHeaders, + accessControl); + } + } + + /** + * Handles the subquery in the given expression. + *

When subquery is detected, first executes the subquery and gets the response, then rewrites the expression with + * the subquery response. + *

Currently only supports ID_SET subquery within the IN_SUBQUERY transform function, which will be rewritten to an + * IN_ID_SET transform function. + */ + private void handleSubquery(Expression expression, long requestId, JsonNode jsonRequest, + @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders, + AccessControl accessControl) + throws Exception { + Function function = expression.getFunctionCall(); + if (function == null) { + return; + } + List operands = function.getOperands(); + if (function.getOperator().equals(IN_SUBQUERY)) { + Preconditions.checkState(operands.size() == 2, "IN_SUBQUERY requires 2 arguments: expression, subquery"); + Literal subqueryLiteral = operands.get(1).getLiteral(); + Preconditions.checkState(subqueryLiteral != null, "Second argument of IN_SUBQUERY must be a literal (subquery)"); + String subquery = subqueryLiteral.getStringValue(); + BrokerResponse response = + handleRequest(requestId, subquery, null, jsonRequest, requesterIdentity, requestContext, httpHeaders, + accessControl); + if (response.getExceptionsSize() != 0) { + throw new RuntimeException("Caught exception while executing subquery: " + subquery); + } + String serializedIdSet = (String) response.getResultTable().getRows().get(0)[0]; + function.setOperator(IN_ID_SET); + operands.set(1, RequestUtils.getLiteralExpression(serializedIdSet)); + } else { + for (Expression operand : operands) { + handleSubquery(operand, requestId, jsonRequest, requesterIdentity, requestContext, httpHeaders, accessControl); + } + } + } + + /** + * Resolves the actual table name for: + * - Case-insensitive cluster + * + * @param tableName the table name in the query + * @param tableCache the table case-sensitive cache + * @return table name if the table name is found in Pinot registry. + */ + @VisibleForTesting + static String getActualTableName(String tableName, TableCache tableCache) { + String actualTableName = tableCache.getActualTableName(tableName); + if (actualTableName != null) { + return actualTableName; + } + return tableName; + } + + /** + * Retrieve segment partitioned columns for a table. + * For a hybrid table, a segment partitioned column has to be the intersection of both offline and realtime tables. + * + * @param tableCache + * @param tableName + * @return segment partitioned columns belong to both offline and realtime tables. + */ + private static Set getSegmentPartitionedColumns(TableCache tableCache, String tableName) { + final TableConfig offlineTableConfig = + tableCache.getTableConfig(TableNameBuilder.OFFLINE.tableNameWithType(tableName)); + final TableConfig realtimeTableConfig = + tableCache.getTableConfig(TableNameBuilder.REALTIME.tableNameWithType(tableName)); + if (offlineTableConfig == null) { + return getSegmentPartitionedColumns(realtimeTableConfig); + } + if (realtimeTableConfig == null) { + return getSegmentPartitionedColumns(offlineTableConfig); + } + Set segmentPartitionedColumns = getSegmentPartitionedColumns(offlineTableConfig); + segmentPartitionedColumns.retainAll(getSegmentPartitionedColumns(realtimeTableConfig)); + return segmentPartitionedColumns; + } + + private static Set getSegmentPartitionedColumns(@Nullable TableConfig tableConfig) { + Set segmentPartitionedColumns = new HashSet<>(); + if (tableConfig == null) { + return segmentPartitionedColumns; + } + List fieldConfigs = tableConfig.getFieldConfigList(); + if (fieldConfigs != null) { + for (FieldConfig fieldConfig : fieldConfigs) { + if (fieldConfig.getProperties() != null && Boolean.parseBoolean( + fieldConfig.getProperties().get(FieldConfig.IS_SEGMENT_PARTITIONED_COLUMN_KEY))) { + segmentPartitionedColumns.add(fieldConfig.getName()); + } + } + } + return segmentPartitionedColumns; + } + + /** + * Retrieve multivalued columns for a table. + * From the table Schema , we get the multi valued columns of dimension fields. + * + * @param tableSchema + * @param columnName + * @return multivalued columns of the table . + */ + private static boolean isMultiValueColumn(Schema tableSchema, String columnName) { + + DimensionFieldSpec dimensionFieldSpec = tableSchema.getDimensionSpec(columnName); + return dimensionFieldSpec != null && !dimensionFieldSpec.isSingleValueField(); + } + + /** + * Sets the table name in the given broker request. + * NOTE: Set table name in broker request because it is used for access control, query routing etc. + */ + private void setTableName(BrokerRequest brokerRequest, String tableName) { + brokerRequest.getQuerySource().setTableName(tableName); + brokerRequest.getPinotQuery().getDataSource().setTableName(tableName); + } + + /** + * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given query. + */ + private static void handleHLLLog2mOverride(PinotQuery pinotQuery, int hllLog2mOverride) { + List selectList = pinotQuery.getSelectList(); + for (Expression expression : selectList) { + handleHLLLog2mOverride(expression, hllLog2mOverride); + } + List orderByList = pinotQuery.getOrderByList(); + if (orderByList != null) { + for (Expression expression : orderByList) { + // NOTE: Order-by is always a Function with the ordering of the Expression + handleHLLLog2mOverride(expression.getFunctionCall().getOperands().get(0), hllLog2mOverride); + } + } + Expression havingExpression = pinotQuery.getHavingExpression(); + if (havingExpression != null) { + handleHLLLog2mOverride(havingExpression, hllLog2mOverride); + } + } + + /** + * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given expression. + */ + private static void handleHLLLog2mOverride(Expression expression, int hllLog2mOverride) { + Function function = expression.getFunctionCall(); + if (function == null) { + return; + } + switch (function.getOperator()) { + case "distinctcounthll": + case "distinctcounthllmv": + case "distinctcountrawhll": + case "distinctcountrawhllmv": + if (function.getOperandsSize() == 1) { + function.addToOperands(RequestUtils.getLiteralExpression(hllLog2mOverride)); + } + return; + default: + break; + } + for (Expression operand : function.getOperands()) { + handleHLLLog2mOverride(operand, hllLog2mOverride); + } + } + + /** + * Overrides the LIMIT of the given query if it exceeds the query limit. + */ + @VisibleForTesting + static void handleQueryLimitOverride(PinotQuery pinotQuery, int queryLimit) { + if (queryLimit > 0 && pinotQuery.getLimit() > queryLimit) { + pinotQuery.setLimit(queryLimit); + } + } + + /** + * Rewrites 'DistinctCount' to 'SegmentPartitionDistinctCount' for the given query. + */ + @VisibleForTesting + static void handleSegmentPartitionedDistinctCountOverride(PinotQuery pinotQuery, + Set segmentPartitionedColumns) { + if (segmentPartitionedColumns.isEmpty()) { + return; + } + for (Expression expression : pinotQuery.getSelectList()) { + handleSegmentPartitionedDistinctCountOverride(expression, segmentPartitionedColumns); + } + List orderByExpressions = pinotQuery.getOrderByList(); + if (orderByExpressions != null) { + for (Expression expression : orderByExpressions) { + // NOTE: Order-by is always a Function with the ordering of the Expression + handleSegmentPartitionedDistinctCountOverride(expression.getFunctionCall().getOperands().get(0), + segmentPartitionedColumns); + } + } + Expression havingExpression = pinotQuery.getHavingExpression(); + if (havingExpression != null) { + handleSegmentPartitionedDistinctCountOverride(havingExpression, segmentPartitionedColumns); + } + } + + /** + * Rewrites 'DistinctCount' to 'SegmentPartitionDistinctCount' for the given expression. + */ + private static void handleSegmentPartitionedDistinctCountOverride(Expression expression, + Set segmentPartitionedColumns) { + Function function = expression.getFunctionCall(); + if (function == null) { + return; + } + if (function.getOperator().equals("distinctcount")) { + List operands = function.getOperands(); + if (operands.size() == 1 && operands.get(0).isSetIdentifier() && segmentPartitionedColumns.contains( + operands.get(0).getIdentifier().getName())) { + function.setOperator("segmentpartitioneddistinctcount"); + } + } else { + for (Expression operand : function.getOperands()) { + handleSegmentPartitionedDistinctCountOverride(operand, segmentPartitionedColumns); + } + } + } + + /** + * Rewrites 'DistinctCount' to 'DistinctCountBitmap' for the given query. + */ + private static void handleDistinctCountBitmapOverride(PinotQuery pinotQuery) { + for (Expression expression : pinotQuery.getSelectList()) { + handleDistinctCountBitmapOverride(expression); + } + List orderByExpressions = pinotQuery.getOrderByList(); + if (orderByExpressions != null) { + for (Expression expression : orderByExpressions) { + // NOTE: Order-by is always a Function with the ordering of the Expression + handleDistinctCountBitmapOverride(expression.getFunctionCall().getOperands().get(0)); + } + } + Expression havingExpression = pinotQuery.getHavingExpression(); + if (havingExpression != null) { + handleDistinctCountBitmapOverride(havingExpression); + } + } + + /** + * Rewrites selected 'Distinct' prefixed function to 'Distinct----MV' function for the field of multivalued type. + */ + @VisibleForTesting + static void handleDistinctMultiValuedOverride(PinotQuery pinotQuery, Schema tableSchema) { + for (Expression expression : pinotQuery.getSelectList()) { + handleDistinctMultiValuedOverride(expression, tableSchema); + } + List orderByExpressions = pinotQuery.getOrderByList(); + if (orderByExpressions != null) { + for (Expression expression : orderByExpressions) { + // NOTE: Order-by is always a Function with the ordering of the Expression + handleDistinctMultiValuedOverride(expression.getFunctionCall().getOperands().get(0), tableSchema); + } + } + Expression havingExpression = pinotQuery.getHavingExpression(); + if (havingExpression != null) { + handleDistinctMultiValuedOverride(havingExpression, tableSchema); + } + } + + /** + * Rewrites selected 'Distinct' prefixed function to 'Distinct----MV' function for the field of multivalued type. + */ + private static void handleDistinctMultiValuedOverride(Expression expression, Schema tableSchema) { + Function function = expression.getFunctionCall(); + if (function == null) { + return; + } + + String overrideOperator = DISTINCT_MV_COL_FUNCTION_OVERRIDE_MAP.get(function.getOperator()); + if (overrideOperator != null) { + List operands = function.getOperands(); + if (operands.size() >= 1 && operands.get(0).isSetIdentifier() && isMultiValueColumn(tableSchema, + operands.get(0).getIdentifier().getName())) { + // we are only checking the first operand that if its a MV column as all the overriding agg. fn.'s have + // first operator is column name + function.setOperator(overrideOperator); + } + } else { + for (Expression operand : function.getOperands()) { + handleDistinctMultiValuedOverride(operand, tableSchema); + } + } + } + + /** + * Rewrites 'DistinctCount' to 'DistinctCountBitmap' for the given expression. + */ + private static void handleDistinctCountBitmapOverride(Expression expression) { + Function function = expression.getFunctionCall(); + if (function == null) { + return; + } + if (function.getOperator().equals("distinctcount")) { + function.setOperator("distinctcountbitmap"); + } else { + for (Expression operand : function.getOperands()) { + handleDistinctCountBitmapOverride(operand); + } + } + } + + private HandlerContext getHandlerContext(@Nullable TableConfig offlineTableConfig, + @Nullable TableConfig realtimeTableConfig) { + Boolean disableGroovyOverride = null; + Boolean useApproximateFunctionOverride = null; + if (offlineTableConfig != null && offlineTableConfig.getQueryConfig() != null) { + QueryConfig offlineTableQueryConfig = offlineTableConfig.getQueryConfig(); + Boolean disableGroovyOfflineTableOverride = offlineTableQueryConfig.getDisableGroovy(); + if (disableGroovyOfflineTableOverride != null) { + disableGroovyOverride = disableGroovyOfflineTableOverride; + } + Boolean useApproximateFunctionOfflineTableOverride = offlineTableQueryConfig.getUseApproximateFunction(); + if (useApproximateFunctionOfflineTableOverride != null) { + useApproximateFunctionOverride = useApproximateFunctionOfflineTableOverride; + } + } + if (realtimeTableConfig != null && realtimeTableConfig.getQueryConfig() != null) { + QueryConfig realtimeTableQueryConfig = realtimeTableConfig.getQueryConfig(); + Boolean disableGroovyRealtimeTableOverride = realtimeTableQueryConfig.getDisableGroovy(); + if (disableGroovyRealtimeTableOverride != null) { + if (disableGroovyOverride == null) { + disableGroovyOverride = disableGroovyRealtimeTableOverride; + } else { + // Disable Groovy if either offline or realtime table config disables Groovy + disableGroovyOverride |= disableGroovyRealtimeTableOverride; + } + } + Boolean useApproximateFunctionRealtimeTableOverride = realtimeTableQueryConfig.getUseApproximateFunction(); + if (useApproximateFunctionRealtimeTableOverride != null) { + if (useApproximateFunctionOverride == null) { + useApproximateFunctionOverride = useApproximateFunctionRealtimeTableOverride; + } else { + // Use approximate function if both offline and realtime table config uses approximate function + useApproximateFunctionOverride &= useApproximateFunctionRealtimeTableOverride; + } + } + } + + boolean disableGroovy = disableGroovyOverride != null ? disableGroovyOverride : _disableGroovy; + boolean useApproximateFunction = + useApproximateFunctionOverride != null ? useApproximateFunctionOverride : _useApproximateFunction; + return new HandlerContext(disableGroovy, useApproximateFunction); + } + + private static class HandlerContext { + final boolean _disableGroovy; + final boolean _useApproximateFunction; + + HandlerContext(boolean disableGroovy, boolean useApproximateFunction) { + _disableGroovy = disableGroovy; + _useApproximateFunction = useApproximateFunction; + } + } + + /** + * Verifies that no groovy is present in the PinotQuery when disabled. + */ + @VisibleForTesting + static void rejectGroovyQuery(PinotQuery pinotQuery) { + List selectList = pinotQuery.getSelectList(); + for (Expression expression : selectList) { + rejectGroovyQuery(expression); + } + List orderByList = pinotQuery.getOrderByList(); + if (orderByList != null) { + for (Expression expression : orderByList) { + // NOTE: Order-by is always a Function with the ordering of the Expression + rejectGroovyQuery(expression.getFunctionCall().getOperands().get(0)); + } + } + Expression havingExpression = pinotQuery.getHavingExpression(); + if (havingExpression != null) { + rejectGroovyQuery(havingExpression); + } + Expression filterExpression = pinotQuery.getFilterExpression(); + if (filterExpression != null) { + rejectGroovyQuery(filterExpression); + } + List groupByList = pinotQuery.getGroupByList(); + if (groupByList != null) { + for (Expression expression : groupByList) { + rejectGroovyQuery(expression); + } + } + } + + private static void rejectGroovyQuery(Expression expression) { + Function function = expression.getFunctionCall(); + if (function == null) { + return; + } + if (function.getOperator().equals("groovy")) { + throw new BadQueryRequestException("Groovy transform functions are disabled for queries"); + } + for (Expression operandExpression : function.getOperands()) { + rejectGroovyQuery(operandExpression); + } + } + + /** + * Rewrites potential expensive functions to their approximation counterparts. + * - DISTINCT_COUNT -> DISTINCT_COUNT_SMART_HLL + * - PERCENTILE -> PERCENTILE_SMART_TDIGEST + */ + @VisibleForTesting + static void handleApproximateFunctionOverride(PinotQuery pinotQuery) { + for (Expression expression : pinotQuery.getSelectList()) { + handleApproximateFunctionOverride(expression); + } + List orderByExpressions = pinotQuery.getOrderByList(); + if (orderByExpressions != null) { + for (Expression expression : orderByExpressions) { + // NOTE: Order-by is always a Function with the ordering of the Expression + handleApproximateFunctionOverride(expression.getFunctionCall().getOperands().get(0)); + } + } + Expression havingExpression = pinotQuery.getHavingExpression(); + if (havingExpression != null) { + handleApproximateFunctionOverride(havingExpression); + } + } + + private static void handleApproximateFunctionOverride(Expression expression) { + Function function = expression.getFunctionCall(); + if (function == null) { + return; + } + String functionName = function.getOperator(); + if (functionName.equals("distinctcount") || functionName.equals("distinctcountmv")) { + function.setOperator("distinctcountsmarthll"); + } else if (functionName.startsWith("percentile")) { + String remainingFunctionName = functionName.substring(10); + if (remainingFunctionName.isEmpty() || remainingFunctionName.equals("mv")) { + function.setOperator("percentilesmarttdigest"); + } else if (remainingFunctionName.matches("\\d+")) { + try { + int percentile = Integer.parseInt(remainingFunctionName); + function.setOperator("percentilesmarttdigest"); + function.addToOperands(RequestUtils.getLiteralExpression(percentile)); + } catch (Exception e) { + throw new BadQueryRequestException("Illegal function name: " + functionName); + } + } else if (remainingFunctionName.matches("\\d+mv")) { + try { + int percentile = Integer.parseInt(remainingFunctionName.substring(0, remainingFunctionName.length() - 2)); + function.setOperator("percentilesmarttdigest"); + function.addToOperands(RequestUtils.getLiteralExpression(percentile)); + } catch (Exception e) { + throw new BadQueryRequestException("Illegal function name: " + functionName); + } + } + } else { + for (Expression operand : function.getOperands()) { + handleApproximateFunctionOverride(operand); + } + } + } + + private static void handleExpressionOverride(PinotQuery pinotQuery, + @Nullable Map expressionOverrideMap) { + if (expressionOverrideMap == null) { + return; + } + pinotQuery.getSelectList().replaceAll(o -> handleExpressionOverride(o, expressionOverrideMap)); + Expression filterExpression = pinotQuery.getFilterExpression(); + if (filterExpression != null) { + pinotQuery.setFilterExpression(handleExpressionOverride(filterExpression, expressionOverrideMap)); + } + List groupByList = pinotQuery.getGroupByList(); + if (groupByList != null) { + groupByList.replaceAll(o -> handleExpressionOverride(o, expressionOverrideMap)); + } + List orderByList = pinotQuery.getOrderByList(); + if (orderByList != null) { + for (Expression expression : orderByList) { + // NOTE: Order-by is always a Function with the ordering of the Expression + expression.getFunctionCall().getOperands().replaceAll(o -> handleExpressionOverride(o, expressionOverrideMap)); + } + } + Expression havingExpression = pinotQuery.getHavingExpression(); + if (havingExpression != null) { + pinotQuery.setHavingExpression(handleExpressionOverride(havingExpression, expressionOverrideMap)); + } + } + + private static Expression handleExpressionOverride(Expression expression, + Map expressionOverrideMap) { + Expression override = expressionOverrideMap.get(expression); + if (override != null) { + return new Expression(override); + } + Function function = expression.getFunctionCall(); + if (function != null) { + function.getOperands().replaceAll(o -> handleExpressionOverride(o, expressionOverrideMap)); + } + return expression; + } + + /** + * Returns {@code true} if the given query only contains literals, {@code false} otherwise. + */ + @VisibleForTesting + static boolean isLiteralOnlyQuery(PinotQuery pinotQuery) { + for (Expression expression : pinotQuery.getSelectList()) { + if (!CalciteSqlParser.isLiteralOnlyExpression(expression)) { + return false; + } + } + return true; + } + + /** + * Processes the literal only query. + */ + private BrokerResponseNative processLiteralOnlyQuery(long requestId, PinotQuery pinotQuery, + RequestContext requestContext) { + BrokerResponseNative brokerResponse = new BrokerResponseNative(); + List columnNames = new ArrayList<>(); + List columnTypes = new ArrayList<>(); + List row = new ArrayList<>(); + for (Expression expression : pinotQuery.getSelectList()) { + computeResultsForExpression(expression, columnNames, columnTypes, row); + } + DataSchema dataSchema = + new DataSchema(columnNames.toArray(new String[0]), columnTypes.toArray(new DataSchema.ColumnDataType[0])); + List rows = new ArrayList<>(); + rows.add(row.toArray()); + ResultTable resultTable = new ResultTable(dataSchema, rows); + brokerResponse.setResultTable(resultTable); + brokerResponse.setTimeUsedMs(System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis()); + augmentStatistics(requestContext, brokerResponse); + if (QueryOptionsUtils.shouldDropResults(pinotQuery.getQueryOptions())) { + brokerResponse.setResultTable(null); + } + return brokerResponse; + } + + // TODO(xiangfu): Move Literal function computation here from Calcite Parser. + private void computeResultsForExpression(Expression e, List columnNames, + List columnTypes, List row) { + if (e.getType() == ExpressionType.LITERAL) { + computeResultsForLiteral(e.getLiteral(), columnNames, columnTypes, row); + } + if (e.getType() == ExpressionType.FUNCTION) { + if (e.getFunctionCall().getOperator().equals("as")) { + String columnName = e.getFunctionCall().getOperands().get(1).getIdentifier().getName(); + computeResultsForExpression(e.getFunctionCall().getOperands().get(0), columnNames, columnTypes, row); + columnNames.set(columnNames.size() - 1, columnName); + } else { + throw new IllegalStateException( + "No able to compute results for function - " + e.getFunctionCall().getOperator()); + } + } + } + + private void computeResultsForLiteral(Literal literal, List columnNames, + List columnTypes, List row) { + columnNames.add(RequestUtils.prettyPrint(literal)); + switch (literal.getSetField()) { + case NULL_VALUE: + columnTypes.add(DataSchema.ColumnDataType.UNKNOWN); + row.add(null); + break; + case BOOL_VALUE: + columnTypes.add(DataSchema.ColumnDataType.BOOLEAN); + row.add(literal.getBoolValue()); + break; + case INT_VALUE: + columnTypes.add(DataSchema.ColumnDataType.INT); + row.add(literal.getIntValue()); + break; + case LONG_VALUE: + columnTypes.add(DataSchema.ColumnDataType.LONG); + row.add(literal.getLongValue()); + break; + case FLOAT_VALUE: + columnTypes.add(DataSchema.ColumnDataType.FLOAT); + row.add(Float.intBitsToFloat(literal.getFloatValue())); + break; + case DOUBLE_VALUE: + columnTypes.add(DataSchema.ColumnDataType.DOUBLE); + row.add(literal.getDoubleValue()); + break; + case BIG_DECIMAL_VALUE: + columnTypes.add(DataSchema.ColumnDataType.BIG_DECIMAL); + row.add(BigDecimalUtils.deserialize(literal.getBigDecimalValue())); + break; + case STRING_VALUE: + columnTypes.add(DataSchema.ColumnDataType.STRING); + row.add(literal.getStringValue()); + break; + case BINARY_VALUE: + columnTypes.add(DataSchema.ColumnDataType.BYTES); + row.add(BytesUtils.toHexString(literal.getBinaryValue())); + break; + // TODO: Revisit the array handling. Currently we are setting List into the row. + case INT_ARRAY_VALUE: + columnTypes.add(DataSchema.ColumnDataType.INT_ARRAY); + row.add(literal.getIntArrayValue()); + break; + case LONG_ARRAY_VALUE: + columnTypes.add(DataSchema.ColumnDataType.LONG_ARRAY); + row.add(literal.getLongArrayValue()); + break; + case FLOAT_ARRAY_VALUE: + columnTypes.add(DataSchema.ColumnDataType.FLOAT_ARRAY); + row.add(literal.getFloatArrayValue().stream().map(Float::intBitsToFloat).collect(Collectors.toList())); + break; + case DOUBLE_ARRAY_VALUE: + columnTypes.add(DataSchema.ColumnDataType.DOUBLE_ARRAY); + row.add(literal.getDoubleArrayValue()); + break; + case STRING_ARRAY_VALUE: + columnTypes.add(DataSchema.ColumnDataType.STRING_ARRAY); + row.add(literal.getStringArrayValue()); + break; + default: + throw new IllegalStateException("Unsupported literal: " + literal); + } + } + + /** + * Fixes the column names to the actual column names in the given query. + */ + @VisibleForTesting + static void updateColumnNames(String rawTableName, PinotQuery pinotQuery, boolean isCaseInsensitive, + Map columnNameMap) { + if (pinotQuery != null) { + boolean hasStar = false; + for (Expression expression : pinotQuery.getSelectList()) { + fixColumnName(rawTableName, expression, columnNameMap, isCaseInsensitive); + //check if the select expression is '*' + if (!hasStar && expression.equals(STAR)) { + hasStar = true; + } + } + //if query has a '*' selection along with other columns + if (hasStar) { + expandStarExpressionsToActualColumns(pinotQuery, columnNameMap); + } + Expression filterExpression = pinotQuery.getFilterExpression(); + if (filterExpression != null) { + // We don't support alias in filter expression, so we don't need to pass aliasMap + fixColumnName(rawTableName, filterExpression, columnNameMap, isCaseInsensitive); + } + List groupByList = pinotQuery.getGroupByList(); + if (groupByList != null) { + for (Expression expression : groupByList) { + fixColumnName(rawTableName, expression, columnNameMap, isCaseInsensitive); + } + } + List orderByList = pinotQuery.getOrderByList(); + if (orderByList != null) { + for (Expression expression : orderByList) { + // NOTE: Order-by is always a Function with the ordering of the Expression + fixColumnName(rawTableName, expression.getFunctionCall().getOperands().get(0), columnNameMap, + isCaseInsensitive); + } + } + Expression havingExpression = pinotQuery.getHavingExpression(); + if (havingExpression != null) { + fixColumnName(rawTableName, havingExpression, columnNameMap, isCaseInsensitive); + } + } + } + + private static void expandStarExpressionsToActualColumns(PinotQuery pinotQuery, Map columnNameMap) { + List originalSelections = pinotQuery.getSelectList(); + //expand '*' + List expandedSelections = new ArrayList<>(); + for (String tableCol : columnNameMap.values()) { + Expression newSelection = RequestUtils.getIdentifierExpression(tableCol); + //we exclude default virtual columns + if (tableCol.charAt(0) != '$') { + expandedSelections.add(newSelection); + } + } + //sort naturally + expandedSelections.sort(null); + List newSelections = new ArrayList<>(); + for (Expression originalSelection : originalSelections) { + if (originalSelection.equals(STAR)) { + newSelections.addAll(expandedSelections); + } else { + newSelections.add(originalSelection); + } + } + pinotQuery.setSelectList(newSelections); + } + + /** + * Fixes the column names to the actual column names in the given expression. + */ + private static void fixColumnName(String rawTableName, Expression expression, Map columnNameMap, + boolean ignoreCase) { + ExpressionType expressionType = expression.getType(); + if (expressionType == ExpressionType.IDENTIFIER) { + Identifier identifier = expression.getIdentifier(); + identifier.setName(getActualColumnName(rawTableName, identifier.getName(), columnNameMap, ignoreCase)); + } else if (expressionType == ExpressionType.FUNCTION) { + final Function functionCall = expression.getFunctionCall(); + switch (functionCall.getOperator()) { + case "as": + fixColumnName(rawTableName, functionCall.getOperands().get(0), columnNameMap, ignoreCase); + break; + case "lookup": + // LOOKUP function looks up another table's schema, skip the check for now. + break; + default: + for (Expression operand : functionCall.getOperands()) { + fixColumnName(rawTableName, operand, columnNameMap, ignoreCase); + } + break; + } + } + } + + /** + * Returns the actual column name for the given column name for: + * - Case-insensitive cluster + * - Column name in the format of [{@code rawTableName}].[column_name] + * - Column name in the format of [logical_table_name].[column_name] while {@code rawTableName} is a translated name + */ + @VisibleForTesting + static String getActualColumnName(String rawTableName, String columnName, @Nullable Map columnNameMap, + boolean ignoreCase) { + if ("*".equals(columnName)) { + return columnName; + } + String columnNameToCheck = trimTableName(rawTableName, columnName, ignoreCase); + if (ignoreCase) { + columnNameToCheck = columnNameToCheck.toLowerCase(); + } + if (columnNameMap != null) { + String actualColumnName = columnNameMap.get(columnNameToCheck); + if (actualColumnName != null) { + return actualColumnName; + } + } + if (columnName.charAt(0) == '$') { + return columnName; + } + throw new BadQueryRequestException("Unknown columnName '" + columnName + "' found in the query"); + } + + private static String trimTableName(String rawTableName, String columnName, boolean ignoreCase) { + int columnNameLength = columnName.length(); + int rawTableNameLength = rawTableName.length(); + if (columnNameLength > rawTableNameLength && columnName.charAt(rawTableNameLength) == '.' + && columnName.regionMatches(ignoreCase, 0, rawTableName, 0, rawTableNameLength)) { + return columnName.substring(rawTableNameLength + 1); + } + // Check if raw table name is translated name ([database_name].[logical_table_name]]) + String[] split = StringUtils.split(rawTableName, '.'); + if (split.length == 2) { + String logicalTableName = split[1]; + int logicalTableNameLength = logicalTableName.length(); + if (columnNameLength > logicalTableNameLength && columnName.charAt(logicalTableNameLength) == '.' + && columnName.regionMatches(ignoreCase, 0, logicalTableName, 0, logicalTableNameLength)) { + return columnName.substring(logicalTableNameLength + 1); + } + } + return columnName; + } + + /** + * Helper function to decide whether to force the log + * + * TODO: come up with other criteria for forcing a log and come up with better numbers + */ + private boolean forceLog(BrokerResponse brokerResponse, long totalTimeMs) { + if (brokerResponse.isNumGroupsLimitReached()) { + return true; + } + + if (brokerResponse.getExceptionsSize() > 0) { + return true; + } + + // If response time is more than 1 sec, force the log + return totalTimeMs > 1000L; + } + + /** + * Sets the query timeout (remaining time in milliseconds) into the query options, and returns the remaining time in + * milliseconds. + *

For the overall query timeout, use query-level timeout (in the query options) if exists, or use table-level + * timeout (in the table config) if exists, or use instance-level timeout (in the broker config). + */ + private long setQueryTimeout(String tableNameWithType, Map queryOptions, long timeSpentMs) + throws TimeoutException { + long queryTimeoutMs; + Long queryLevelTimeoutMs = QueryOptionsUtils.getTimeoutMs(queryOptions); + if (queryLevelTimeoutMs != null) { + // Use query-level timeout if exists + queryTimeoutMs = queryLevelTimeoutMs; + } else { + Long tableLevelTimeoutMs = _routingManager.getQueryTimeoutMs(tableNameWithType); + if (tableLevelTimeoutMs != null) { + // Use table-level timeout if exists + queryTimeoutMs = tableLevelTimeoutMs; + } else { + // Use instance-level timeout + queryTimeoutMs = _brokerTimeoutMs; + } + } + + long remainingTimeMs = queryTimeoutMs - timeSpentMs; + if (remainingTimeMs <= 0) { + String errorMessage = + String.format("Query timed out (time spent: %dms, timeout: %dms) for table: %s before scattering the request", + timeSpentMs, queryTimeoutMs, tableNameWithType); + throw new TimeoutException(errorMessage); + } + queryOptions.put(QueryOptionKey.TIMEOUT_MS, Long.toString(remainingTimeMs)); + return remainingTimeMs; + } + + /** + * Sets a query option indicating the maximum response size that can be sent from a server to the broker. This size + * is measured for the serialized response. + * + * The overriding order of priority is: + * 1. QueryOption -> maxServerResponseSizeBytes + * 2. QueryOption -> maxQueryResponseSizeBytes + * 3. TableConfig -> maxServerResponseSizeBytes + * 4. TableConfig -> maxQueryResponseSizeBytes + * 5. BrokerConfig -> maxServerResponseSizeBytes + * 6. BrokerConfig -> maxServerResponseSizeBytes + */ + private void setMaxServerResponseSizeBytes(int numServers, Map queryOptions, + @Nullable TableConfig tableConfig) { + // QueryOption + if (QueryOptionsUtils.getMaxServerResponseSizeBytes(queryOptions) != null) { + return; + } + Long maxQueryResponseSizeQueryOption = QueryOptionsUtils.getMaxQueryResponseSizeBytes(queryOptions); + if (maxQueryResponseSizeQueryOption != null) { + queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, + Long.toString(maxQueryResponseSizeQueryOption / numServers)); + return; + } + + // TableConfig + if (tableConfig != null && tableConfig.getQueryConfig() != null) { + QueryConfig queryConfig = tableConfig.getQueryConfig(); + if (queryConfig.getMaxServerResponseSizeBytes() != null) { + queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, + Long.toString(queryConfig.getMaxServerResponseSizeBytes())); + return; + } + if (queryConfig.getMaxQueryResponseSizeBytes() != null) { + queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, + Long.toString(queryConfig.getMaxQueryResponseSizeBytes() / numServers)); + return; + } + } + + // BrokerConfig + String maxServerResponseSizeBrokerConfig = _config.getProperty(Broker.CONFIG_OF_MAX_SERVER_RESPONSE_SIZE_BYTES); + if (maxServerResponseSizeBrokerConfig != null) { + queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, + Long.toString(DataSizeUtils.toBytes(maxServerResponseSizeBrokerConfig))); + return; + } + + String maxQueryResponseSizeBrokerConfig = _config.getProperty(Broker.CONFIG_OF_MAX_QUERY_RESPONSE_SIZE_BYTES); + if (maxQueryResponseSizeBrokerConfig != null) { + queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES, + Long.toString(DataSizeUtils.toBytes(maxQueryResponseSizeBrokerConfig) / numServers)); + } + } + + /** + * Broker side validation on the query. + *

Throw exception if query does not pass validation. + *

Current validations are: + *

    + *
  • Value for 'LIMIT' <= configured value
    • + *
  • Query options must be set to SQL mode
    • + *
  • Check if numReplicaGroupsToQuery option provided is valid
    • + *
+ */ + @VisibleForTesting + static void validateRequest(PinotQuery pinotQuery, int queryResponseLimit) { + // Verify LIMIT + int limit = pinotQuery.getLimit(); + if (limit > queryResponseLimit) { + throw new IllegalStateException( + "Value for 'LIMIT' (" + limit + ") exceeds maximum allowed value of " + queryResponseLimit); + } + + Map queryOptions = pinotQuery.getQueryOptions(); + try { + // throw errors if options is less than 1 or invalid + Integer numReplicaGroupsToQuery = QueryOptionsUtils.getNumReplicaGroupsToQuery(queryOptions); + if (numReplicaGroupsToQuery != null) { + Preconditions.checkState(numReplicaGroupsToQuery > 0, "numReplicaGroups must be " + "positive number, got: %d", + numReplicaGroupsToQuery); + } + } catch (NumberFormatException e) { + String numReplicaGroupsToQuery = queryOptions.get(QueryOptionKey.NUM_REPLICA_GROUPS_TO_QUERY); + throw new IllegalStateException( + String.format("numReplicaGroups must be a positive number, got: %s", numReplicaGroupsToQuery)); + } + + if (pinotQuery.getDataSource().getSubquery() != null) { + validateRequest(pinotQuery.getDataSource().getSubquery(), queryResponseLimit); + } + } + + /** + * Helper method to attach the time boundary to the given PinotQuery. + */ + private static void attachTimeBoundary(PinotQuery pinotQuery, TimeBoundaryInfo timeBoundaryInfo, + boolean isOfflineRequest) { + String functionName = isOfflineRequest ? FilterKind.LESS_THAN_OR_EQUAL.name() : FilterKind.GREATER_THAN.name(); + String timeColumn = timeBoundaryInfo.getTimeColumn(); + String timeValue = timeBoundaryInfo.getTimeValue(); + Expression timeFilterExpression = + RequestUtils.getFunctionExpression(functionName, RequestUtils.getIdentifierExpression(timeColumn), + RequestUtils.getLiteralExpression(timeValue)); + + Expression filterExpression = pinotQuery.getFilterExpression(); + if (filterExpression != null) { + pinotQuery.setFilterExpression( + RequestUtils.getFunctionExpression(FilterKind.AND.name(), filterExpression, timeFilterExpression)); + } else { + pinotQuery.setFilterExpression(timeFilterExpression); + } + } + + /** + * Processes the optimized broker requests for both OFFLINE and REALTIME table. + * TODO: Directly take PinotQuery + */ + protected abstract BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest, + BrokerRequest serverBrokerRequest, @Nullable BrokerRequest offlineBrokerRequest, + @Nullable Map, List>> offlineRoutingTable, + @Nullable BrokerRequest realtimeBrokerRequest, + @Nullable Map, List>> realtimeRoutingTable, long timeoutMs, + ServerStats serverStats, RequestContext requestContext) + throws Exception; + + private String getGlobalQueryId(long requestId) { + return _brokerId + "_" + requestId; + } + + /** + * Helper class to pass the per server statistics. + */ + public static class ServerStats { + private String _serverStats; + + public String getServerStats() { + return _serverStats; + } + + public void setServerStats(String serverStats) { + _serverStats = serverStats; + } + } + + /** + * Helper class to track the query plaintext and the requested servers. + */ + private static class QueryServers { + final String _query; + final Set _servers = new HashSet<>(); + + QueryServers(String query, @Nullable Map, List>> offlineRoutingTable, + @Nullable Map, List>> realtimeRoutingTable) { + _query = query; + if (offlineRoutingTable != null) { + _servers.addAll(offlineRoutingTable.keySet()); + } + if (realtimeRoutingTable != null) { + _servers.addAll(realtimeRoutingTable.keySet()); + } + } + } +} diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandler.java index 4aa7b26eaab4..9744a1cd5ec0 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandler.java @@ -19,6 +19,8 @@ package org.apache.pinot.broker.requesthandler; import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.common.annotations.VisibleForTesting; import java.util.Map; import java.util.concurrent.Executor; import javax.annotation.Nullable; @@ -28,6 +30,10 @@ import org.apache.pinot.broker.api.RequesterIdentity; import org.apache.pinot.common.response.BrokerResponse; import org.apache.pinot.spi.trace.RequestContext; +import org.apache.pinot.spi.trace.RequestScope; +import org.apache.pinot.spi.trace.Tracing; +import org.apache.pinot.spi.utils.CommonConstants.Broker.Request; +import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.sql.parsers.SqlNodeAndOptions; @@ -42,10 +48,15 @@ BrokerResponse handleRequest(JsonNode request, @Nullable SqlNodeAndOptions sqlNo @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders) throws Exception; - default BrokerResponse handleRequest(JsonNode request, @Nullable RequesterIdentity requesterIdentity, - RequestContext requestContext, @Nullable HttpHeaders httpHeaders) + @VisibleForTesting + default BrokerResponse handleRequest(String sql) throws Exception { - return handleRequest(request, null, requesterIdentity, requestContext, httpHeaders); + ObjectNode request = JsonUtils.newObjectNode(); + request.put(Request.SQL, sql); + try (RequestScope requestContext = Tracing.getTracer().createRequestScope()) { + requestContext.setRequestArrivalTimeMillis(System.currentTimeMillis()); + return handleRequest(request, null, null, requestContext, null); + } } Map getRunningQueries(); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandlerDelegate.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandlerDelegate.java index 0da84e6eadc4..cdbf64e3bc72 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandlerDelegate.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandlerDelegate.java @@ -26,16 +26,13 @@ import org.apache.http.conn.HttpClientConnectionManager; import org.apache.pinot.broker.api.RequesterIdentity; import org.apache.pinot.common.exception.QueryException; -import org.apache.pinot.common.metrics.BrokerMeter; -import org.apache.pinot.common.metrics.BrokerMetrics; import org.apache.pinot.common.response.BrokerResponse; import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.spi.trace.RequestContext; -import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.Broker.Request; import org.apache.pinot.sql.parsers.SqlNodeAndOptions; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; /** @@ -45,26 +42,18 @@ * {@see: @CommonConstant */ public class BrokerRequestHandlerDelegate implements BrokerRequestHandler { - private static final Logger LOGGER = LoggerFactory.getLogger(BrokerRequestHandlerDelegate.class); + private final BaseSingleStageBrokerRequestHandler _singleStageBrokerRequestHandler; + private final MultiStageBrokerRequestHandler _multiStageBrokerRequestHandler; - private final BrokerRequestHandler _singleStageBrokerRequestHandler; - private final BrokerRequestHandler _multiStageBrokerRequestHandler; - private final BrokerMetrics _brokerMetrics; - private final String _brokerId; - - public BrokerRequestHandlerDelegate(String brokerId, BrokerRequestHandler singleStageBrokerRequestHandler, - @Nullable BrokerRequestHandler multiStageBrokerRequestHandler, BrokerMetrics brokerMetrics) { - _brokerId = brokerId; + public BrokerRequestHandlerDelegate(BaseSingleStageBrokerRequestHandler singleStageBrokerRequestHandler, + @Nullable MultiStageBrokerRequestHandler multiStageBrokerRequestHandler) { _singleStageBrokerRequestHandler = singleStageBrokerRequestHandler; _multiStageBrokerRequestHandler = multiStageBrokerRequestHandler; - _brokerMetrics = brokerMetrics; } @Override public void start() { - if (_singleStageBrokerRequestHandler != null) { - _singleStageBrokerRequestHandler.start(); - } + _singleStageBrokerRequestHandler.start(); if (_multiStageBrokerRequestHandler != null) { _multiStageBrokerRequestHandler.start(); } @@ -72,9 +61,7 @@ public void start() { @Override public void shutDown() { - if (_singleStageBrokerRequestHandler != null) { - _singleStageBrokerRequestHandler.shutDown(); - } + _singleStageBrokerRequestHandler.shutDown(); if (_multiStageBrokerRequestHandler != null) { _multiStageBrokerRequestHandler.shutDown(); } @@ -84,25 +71,20 @@ public void shutDown() { public BrokerResponse handleRequest(JsonNode request, @Nullable SqlNodeAndOptions sqlNodeAndOptions, @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders) throws Exception { - requestContext.setBrokerId(_brokerId); + // Parse the query if needed if (sqlNodeAndOptions == null) { try { - sqlNodeAndOptions = RequestUtils.parseQuery(request.get(CommonConstants.Broker.Request.SQL).asText(), request); + sqlNodeAndOptions = RequestUtils.parseQuery(request.get(Request.SQL).asText(), request); } catch (Exception e) { - LOGGER.info("Caught exception while compiling SQL: {}, {}", request, e.getMessage()); - _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_COMPILATION_EXCEPTIONS, 1); + // Do not log or emit metric here because it is pure user error requestContext.setErrorCode(QueryException.SQL_PARSING_ERROR_CODE); return new BrokerResponseNative(QueryException.getException(QueryException.SQL_PARSING_ERROR, e)); } } - if (request.has(CommonConstants.Broker.Request.QUERY_OPTIONS)) { - sqlNodeAndOptions.setExtraOptions( - RequestUtils.getOptionsFromJson(request, CommonConstants.Broker.Request.QUERY_OPTIONS)); - } - - if (_multiStageBrokerRequestHandler != null && Boolean.parseBoolean( - sqlNodeAndOptions.getOptions().get(CommonConstants.Broker.Request.QueryOptionKey.USE_MULTISTAGE_ENGINE))) { - return _multiStageBrokerRequestHandler.handleRequest(request, requesterIdentity, requestContext, httpHeaders); + if (_multiStageBrokerRequestHandler != null && QueryOptionsUtils.isUseMultistageEngine( + sqlNodeAndOptions.getOptions())) { + return _multiStageBrokerRequestHandler.handleRequest(request, sqlNodeAndOptions, requesterIdentity, + requestContext, httpHeaders); } else { return _singleStageBrokerRequestHandler.handleRequest(request, sqlNodeAndOptions, requesterIdentity, requestContext, httpHeaders); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java index c2cd3d60276a..96e06ffec5a9 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java @@ -23,6 +23,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; import javax.annotation.Nullable; import javax.annotation.concurrent.ThreadSafe; import org.apache.commons.lang3.tuple.Pair; @@ -30,9 +31,7 @@ import org.apache.pinot.broker.queryquota.QueryQuotaManager; import org.apache.pinot.broker.routing.BrokerRoutingManager; import org.apache.pinot.common.config.GrpcConfig; -import org.apache.pinot.common.config.TlsConfig; import org.apache.pinot.common.config.provider.TableCache; -import org.apache.pinot.common.metrics.BrokerMetrics; import org.apache.pinot.common.proto.Server; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.response.broker.BrokerResponseNative; @@ -43,37 +42,23 @@ import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.env.PinotConfiguration; -import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListener; import org.apache.pinot.spi.trace.RequestContext; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; /** * The GrpcBrokerRequestHandler class communicates query request via GRPC. */ @ThreadSafe -public class GrpcBrokerRequestHandler extends BaseBrokerRequestHandler { - private static final Logger LOGGER = LoggerFactory.getLogger(GrpcBrokerRequestHandler.class); - - private final GrpcConfig _grpcConfig; +public class GrpcBrokerRequestHandler extends BaseSingleStageBrokerRequestHandler { private final StreamingReduceService _streamingReduceService; private final PinotStreamingQueryClient _streamingQueryClient; // TODO: Support TLS public GrpcBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, - AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, - BrokerMetrics brokerMetrics, TlsConfig tlsConfig, BrokerQueryEventListener brokerQueryEventListener) { - super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, brokerMetrics, - brokerQueryEventListener); - LOGGER.info("Using Grpc BrokerRequestHandler."); - _grpcConfig = GrpcConfig.buildGrpcQueryConfig(config); - - // create streaming query client - _streamingQueryClient = new PinotStreamingQueryClient(_grpcConfig); - - // create streaming reduce service + AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); _streamingReduceService = new StreamingReduceService(config); + _streamingQueryClient = new PinotStreamingQueryClient(GrpcConfig.buildGrpcQueryConfig(config)); } @Override @@ -81,7 +66,7 @@ public void start() { } @Override - public synchronized void shutDown() { + public void shutDown() { _streamingQueryClient.shutdown(); _streamingReduceService.shutDown(); } @@ -95,6 +80,7 @@ protected BrokerResponseNative processBrokerRequest(long requestId, BrokerReques ServerStats serverStats, RequestContext requestContext) throws Exception { // TODO: Support failure detection + // TODO: Add servers queried/responded stats assert offlineBrokerRequest != null || realtimeBrokerRequest != null; Map> responseMap = new HashMap<>(); if (offlineBrokerRequest != null) { @@ -107,10 +93,10 @@ protected BrokerResponseNative processBrokerRequest(long requestId, BrokerReques sendRequest(requestId, TableType.REALTIME, realtimeBrokerRequest, realtimeRoutingTable, responseMap, requestContext.isSampledRequest()); } - final long startReduceTimeNanos = System.nanoTime(); + long reduceStartTimeNs = System.nanoTime(); BrokerResponseNative brokerResponse = _streamingReduceService.reduceOnStreamResponse(originalBrokerRequest, responseMap, timeoutMs, _brokerMetrics); - requestContext.setReduceTimeNanos(System.nanoTime() - startReduceTimeNanos); + brokerResponse.setBrokerReduceTimeMs(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - reduceStartTimeNs)); return brokerResponse; } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index 01bfe456b5b1..8e608473086b 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -19,19 +19,19 @@ package org.apache.pinot.broker.requesthandler; import com.fasterxml.jackson.databind.JsonNode; -import com.google.common.collect.Maps; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.concurrent.TimeUnit; +import java.util.concurrent.Executor; import java.util.concurrent.TimeoutException; +import java.util.stream.Collectors; import javax.annotation.Nullable; import javax.ws.rs.WebApplicationException; import javax.ws.rs.core.HttpHeaders; import javax.ws.rs.core.Response; -import org.apache.commons.lang3.tuple.Pair; +import org.apache.commons.lang.StringUtils; +import org.apache.http.conn.HttpClientConnectionManager; import org.apache.pinot.broker.api.AccessControl; import org.apache.pinot.broker.api.RequesterIdentity; import org.apache.pinot.broker.broker.AccessControlFactory; @@ -42,14 +42,10 @@ import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.exception.QueryException; import org.apache.pinot.common.metrics.BrokerMeter; -import org.apache.pinot.common.metrics.BrokerMetrics; import org.apache.pinot.common.metrics.BrokerQueryPhase; -import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.response.BrokerResponse; import org.apache.pinot.common.response.broker.BrokerResponseNative; import org.apache.pinot.common.response.broker.BrokerResponseNativeV2; -import org.apache.pinot.common.response.broker.BrokerResponseStats; -import org.apache.pinot.common.response.broker.QueryProcessingException; import org.apache.pinot.common.response.broker.ResultTable; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.DatabaseUtils; @@ -58,19 +54,20 @@ import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.core.auth.Actions; import org.apache.pinot.core.auth.TargetType; -import org.apache.pinot.core.query.reduce.ExecutionStatsAggregator; -import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.query.QueryEnvironment; import org.apache.pinot.query.catalog.PinotCatalog; import org.apache.pinot.query.mailbox.MailboxService; import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; +import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.routing.WorkerManager; +import org.apache.pinot.query.runtime.MultiStageStatsTreeBuilder; +import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; import org.apache.pinot.query.service.dispatch.QueryDispatcher; import org.apache.pinot.query.type.TypeFactory; import org.apache.pinot.query.type.TypeSystem; +import org.apache.pinot.spi.auth.TableAuthorizationResult; import org.apache.pinot.spi.env.PinotConfiguration; -import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListener; import org.apache.pinot.spi.exception.DatabaseConflictException; import org.apache.pinot.spi.trace.RequestContext; import org.apache.pinot.spi.utils.CommonConstants; @@ -84,51 +81,56 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageBrokerRequestHandler.class); private final WorkerManager _workerManager; - private final MailboxService _mailboxService; private final QueryDispatcher _queryDispatcher; public MultiStageBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, - AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, - BrokerMetrics brokerMetrics, BrokerQueryEventListener brokerQueryEventListener) { - super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, brokerMetrics, - brokerQueryEventListener); - LOGGER.info("Using Multi-stage BrokerRequestHandler."); + AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); String hostname = config.getProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME); int port = Integer.parseInt(config.getProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT)); _workerManager = new WorkerManager(hostname, port, _routingManager); - _mailboxService = new MailboxService(hostname, port, config); - _queryDispatcher = new QueryDispatcher(_mailboxService); + _queryDispatcher = new QueryDispatcher(new MailboxService(hostname, port, config)); + LOGGER.info("Initialized MultiStageBrokerRequestHandler on host: {}, port: {} with broker id: {}, timeout: {}ms, " + + "query log max length: {}, query log max rate: {}", hostname, port, _brokerId, _brokerTimeoutMs, + _queryLogger.getMaxQueryLengthToLog(), _queryLogger.getLogRateLimit()); + } - // TODO: move this to a startUp() function. - _mailboxService.start(); + @Override + public void start() { + _queryDispatcher.start(); + } + + @Override + public void shutDown() { + _queryDispatcher.shutdown(); } @Override protected BrokerResponse handleRequest(long requestId, String query, @Nullable SqlNodeAndOptions sqlNodeAndOptions, JsonNode request, @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, - HttpHeaders httpHeaders) { + HttpHeaders httpHeaders, AccessControl accessControl) { LOGGER.debug("SQL query for request {}: {}", requestId, query); - long compilationStartTimeNs; + // Parse the query if needed + if (sqlNodeAndOptions == null) { + try { + sqlNodeAndOptions = RequestUtils.parseQuery(query, request); + } catch (Exception e) { + // Do not log or emit metric here because it is pure user error + requestContext.setErrorCode(QueryException.SQL_PARSING_ERROR_CODE); + return new BrokerResponseNative(QueryException.getException(QueryException.SQL_PARSING_ERROR, e)); + } + } + + // Compile the request + Map queryOptions = sqlNodeAndOptions.getOptions(); + long compilationStartTimeNs = System.nanoTime(); long queryTimeoutMs; QueryEnvironment.QueryPlannerResult queryPlanResult; try { - // Parse the request - sqlNodeAndOptions = sqlNodeAndOptions != null ? sqlNodeAndOptions : RequestUtils.parseQuery(query, request); - } catch (RuntimeException e) { - String consolidatedMessage = ExceptionUtils.consolidateExceptionMessages(e); - LOGGER.info("Caught exception parsing request {}: {}, {}", requestId, query, consolidatedMessage); - _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_COMPILATION_EXCEPTIONS, 1); - requestContext.setErrorCode(QueryException.SQL_PARSING_ERROR_CODE); - return new BrokerResponseNative( - QueryException.getException(QueryException.SQL_PARSING_ERROR, consolidatedMessage)); - } - try { - Long timeoutMsFromQueryOption = QueryOptionsUtils.getTimeoutMs(sqlNodeAndOptions.getOptions()); - queryTimeoutMs = timeoutMsFromQueryOption == null ? _brokerTimeoutMs : timeoutMsFromQueryOption; - String database = DatabaseUtils.extractDatabaseFromQueryRequest(sqlNodeAndOptions.getOptions(), httpHeaders); - // Compile the request - compilationStartTimeNs = System.nanoTime(); + Long timeoutMsFromQueryOption = QueryOptionsUtils.getTimeoutMs(queryOptions); + queryTimeoutMs = timeoutMsFromQueryOption != null ? timeoutMsFromQueryOption : _brokerTimeoutMs; + String database = DatabaseUtils.extractDatabaseFromQueryRequest(queryOptions, httpHeaders); QueryEnvironment queryEnvironment = new QueryEnvironment(new TypeFactory(new TypeSystem()), CalciteSchemaBuilder.asRootSchema(new PinotCatalog(database, _tableCache), database), _workerManager, _tableCache); @@ -137,10 +139,16 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S queryPlanResult = queryEnvironment.explainQuery(query, sqlNodeAndOptions, requestId); String plan = queryPlanResult.getExplainPlan(); Set tableNames = queryPlanResult.getTableNames(); - if (!hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders)) { - throw new WebApplicationException("Permission denied", Response.Status.FORBIDDEN); + TableAuthorizationResult tableAuthorizationResult = + hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders); + if (!tableAuthorizationResult.hasAccess()) { + String failureMessage = tableAuthorizationResult.getFailureMessage(); + if (StringUtils.isNotBlank(failureMessage)) { + failureMessage = "Reason: " + failureMessage; + } + throw new WebApplicationException("Permission denied. " + failureMessage, + Response.Status.FORBIDDEN); } - return constructMultistageExplainPlan(query, plan); case SELECT: default: @@ -184,8 +192,15 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S updatePhaseTimingForTables(tableNames, BrokerQueryPhase.REQUEST_COMPILATION, compilationTimeNs); // Validate table access. - if (!hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders)) { - throw new WebApplicationException("Permission denied", Response.Status.FORBIDDEN); + TableAuthorizationResult tableAuthorizationResult = + hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders); + if (!tableAuthorizationResult.hasAccess()) { + String failureMessage = tableAuthorizationResult.getFailureMessage(); + if (StringUtils.isNotBlank(failureMessage)) { + failureMessage = "Reason: " + failureMessage; + } + throw new WebApplicationException("Permission denied." + failureMessage, + Response.Status.FORBIDDEN); } // Validate QPS quota @@ -194,25 +209,11 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S return new BrokerResponseNative(QueryException.getException(QueryException.QUOTA_EXCEEDED_ERROR, errorMessage)); } - Map queryOptions = sqlNodeAndOptions.getOptions(); - boolean traceEnabled = Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.TRACE)); - Map stageIdStatsMap; - if (!traceEnabled) { - stageIdStatsMap = Collections.singletonMap(0, new ExecutionStatsAggregator(false)); - } else { - List stagePlans = dispatchableSubPlan.getQueryStageList(); - int numStages = stagePlans.size(); - stageIdStatsMap = Maps.newHashMapWithExpectedSize(numStages); - for (int stageId = 0; stageId < numStages; stageId++) { - stageIdStatsMap.put(stageId, new ExecutionStatsAggregator(true)); - } - } - long executionStartTimeNs = System.nanoTime(); - ResultTable queryResults; + QueryDispatcher.QueryResult queryResults; try { - queryResults = _queryDispatcher.submitAndReduce(requestContext, dispatchableSubPlan, queryTimeoutMs, queryOptions, - stageIdStatsMap); + queryResults = + _queryDispatcher.submitAndReduce(requestContext, dispatchableSubPlan, queryTimeoutMs, queryOptions); } catch (TimeoutException e) { for (String table : tableNames) { _brokerMetrics.addMeteredTableValue(table, BrokerMeter.BROKER_RESPONSES_WITH_TIMEOUTS, 1); @@ -231,7 +232,9 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S updatePhaseTimingForTables(tableNames, BrokerQueryPhase.QUERY_EXECUTION, executionEndTimeNs - executionStartTimeNs); BrokerResponseNativeV2 brokerResponse = new BrokerResponseNativeV2(); - brokerResponse.setResultTable(queryResults); + brokerResponse.setResultTable(queryResults.getResultTable()); + // TODO: Add servers queried/responded stats + brokerResponse.setBrokerReduceTimeMs(queryResults.getBrokerReduceTimeMs()); // Attach unavailable segments int numUnavailableSegments = 0; @@ -239,59 +242,68 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S String tableName = entry.getKey(); Set unavailableSegments = entry.getValue(); numUnavailableSegments += unavailableSegments.size(); - brokerResponse.addToExceptions(new QueryProcessingException(QueryException.SERVER_SEGMENT_MISSING_ERROR_CODE, + brokerResponse.addException(QueryException.getException(QueryException.SERVER_SEGMENT_MISSING_ERROR, String.format("Find unavailable segments: %s for table: %s", unavailableSegments, tableName))); } + requestContext.setNumUnavailableSegments(numUnavailableSegments); - for (Map.Entry entry : stageIdStatsMap.entrySet()) { - if (entry.getKey() == 0) { - // Root stats are aggregated and added separately to broker response for backward compatibility - entry.getValue().setStats(brokerResponse); - continue; - } - - BrokerResponseStats brokerResponseStats = new BrokerResponseStats(); - if (!tableNames.isEmpty()) { - //TODO: Only using first table to assign broker metrics - // find a way to split metrics in case of multiple table - String rawTableName = TableNameBuilder.extractRawTableName(tableNames.iterator().next()); - entry.getValue().setStageLevelStats(rawTableName, brokerResponseStats, _brokerMetrics); - } else { - entry.getValue().setStageLevelStats(null, brokerResponseStats, null); - } - brokerResponse.addStageStat(entry.getKey(), brokerResponseStats); - } - - // Set partial result flag - brokerResponse.setPartialResult(isPartialResult(brokerResponse)); + fillOldBrokerResponseStats(brokerResponse, queryResults.getQueryStats(), dispatchableSubPlan); // Set total query processing time // TODO: Currently we don't emit metric for QUERY_TOTAL_TIME_MS - long totalTimeMs = TimeUnit.NANOSECONDS.toMillis( - sqlNodeAndOptions.getParseTimeNs() + (executionEndTimeNs - compilationStartTimeNs)); + long totalTimeMs = System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis(); brokerResponse.setTimeUsedMs(totalTimeMs); - requestContext.setQueryProcessingTime(totalTimeMs); - requestContext.setTraceInfo(brokerResponse.getTraceInfo()); augmentStatistics(requestContext, brokerResponse); + if (QueryOptionsUtils.shouldDropResults(queryOptions)) { + brokerResponse.setResultTable(null); + } // Log query and stats _queryLogger.log( - new QueryLogger.QueryLogParams(requestId, query, requestContext, tableNames.toString(), numUnavailableSegments, - null, brokerResponse, totalTimeMs, requesterIdentity)); + new QueryLogger.QueryLogParams(requestContext, tableNames.toString(), brokerResponse, requesterIdentity, null)); return brokerResponse; } + private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse, + List queryStats, DispatchableSubPlan dispatchableSubPlan) { + List stagePlans = dispatchableSubPlan.getQueryStageList(); + List planNodes = new ArrayList<>(stagePlans.size()); + for (DispatchablePlanFragment stagePlan : stagePlans) { + planNodes.add(stagePlan.getPlanFragment().getFragmentRoot()); + } + MultiStageStatsTreeBuilder treeBuilder = new MultiStageStatsTreeBuilder(planNodes, queryStats); + brokerResponse.setStageStats(treeBuilder.jsonStatsByStage(0)); + for (MultiStageQueryStats.StageStats.Closed stageStats : queryStats) { + if (stageStats != null) { // for example pipeline breaker may not have stats + stageStats.forEach((type, stats) -> type.mergeInto(brokerResponse, stats)); + } + } + } + /** * Validates whether the requester has access to all the tables. */ - private boolean hasTableAccess(RequesterIdentity requesterIdentity, Set tableNames, + private TableAuthorizationResult hasTableAccess(RequesterIdentity requesterIdentity, Set tableNames, RequestContext requestContext, HttpHeaders httpHeaders) { final long startTimeNs = System.nanoTime(); AccessControl accessControl = _accessControlFactory.create(); - boolean hasAccess = accessControl.hasAccess(requesterIdentity, tableNames) && tableNames.stream() - .allMatch(table -> accessControl.hasAccess(httpHeaders, TargetType.TABLE, table, Actions.Table.QUERY)); - if (!hasAccess) { + + TableAuthorizationResult tableAuthorizationResult = accessControl.authorize(requesterIdentity, tableNames); + + Set failedTables = tableNames.stream() + .filter(table -> !accessControl.hasAccess(httpHeaders, TargetType.TABLE, table, Actions.Table.QUERY)) + .collect(Collectors.toSet()); + + failedTables.addAll(tableAuthorizationResult.getFailedTables()); + + if (!failedTables.isEmpty()) { + tableAuthorizationResult = new TableAuthorizationResult(failedTables); + } else { + tableAuthorizationResult = TableAuthorizationResult.success(); + } + + if (!tableAuthorizationResult.hasAccess()) { _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1); LOGGER.warn("Access denied for requestId {}", requestContext.getRequestId()); requestContext.setErrorCode(QueryException.ACCESS_DENIED_ERROR_CODE); @@ -299,7 +311,7 @@ private boolean hasTableAccess(RequesterIdentity requesterIdentity, Set updatePhaseTimingForTables(tableNames, BrokerQueryPhase.AUTHORIZATION, System.nanoTime() - startTimeNs); - return hasAccess; + return tableAuthorizationResult; } /** @@ -325,7 +337,7 @@ private void updatePhaseTimingForTables(Set tableNames, BrokerQueryPhase } } - private BrokerResponseNative constructMultistageExplainPlan(String sql, String plan) { + private BrokerResponse constructMultistageExplainPlan(String sql, String plan) { BrokerResponseNative brokerResponse = BrokerResponseNative.empty(); List rows = new ArrayList<>(); rows.add(new Object[]{sql, plan}); @@ -336,24 +348,15 @@ private BrokerResponseNative constructMultistageExplainPlan(String sql, String p } @Override - protected BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest, - BrokerRequest serverBrokerRequest, @Nullable BrokerRequest offlineBrokerRequest, - @Nullable Map, List>> offlineRoutingTable, - @Nullable BrokerRequest realtimeBrokerRequest, - @Nullable Map, List>> realtimeRoutingTable, long timeoutMs, - ServerStats serverStats, RequestContext requestContext) - throws Exception { + public Map getRunningQueries() { + // TODO: Support running query tracking for multi-stage engine throw new UnsupportedOperationException(); } @Override - public void start() { - // no-op - } - - @Override - public void shutDown() { - _queryDispatcher.shutdown(); - _mailboxService.shutdown(); + public boolean cancelQuery(long queryId, int timeoutMs, Executor executor, HttpClientConnectionManager connMgr, + Map serverResponses) { + // TODO: Support query cancellation for multi-stage engine + throw new UnsupportedOperationException(); } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java index 68ae70f9eb57..a51634c983f8 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java @@ -18,8 +18,8 @@ */ package org.apache.pinot.broker.requesthandler; +import com.google.common.collect.Maps; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -37,12 +37,10 @@ import org.apache.pinot.common.datatable.DataTable; import org.apache.pinot.common.exception.QueryException; import org.apache.pinot.common.metrics.BrokerMeter; -import org.apache.pinot.common.metrics.BrokerMetrics; import org.apache.pinot.common.metrics.BrokerQueryPhase; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.response.broker.BrokerResponseNative; import org.apache.pinot.common.response.broker.QueryProcessingException; -import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.core.query.reduce.BrokerReduceService; import org.apache.pinot.core.transport.AsyncQueryResponse; import org.apache.pinot.core.transport.QueryResponse; @@ -52,7 +50,6 @@ import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; import org.apache.pinot.spi.env.PinotConfiguration; -import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListener; import org.apache.pinot.spi.trace.RequestContext; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.builder.TableNameBuilder; @@ -65,7 +62,8 @@ * connection per server to route the queries. */ @ThreadSafe -public class SingleConnectionBrokerRequestHandler extends BaseBrokerRequestHandler implements FailureDetector.Listener { +public class SingleConnectionBrokerRequestHandler extends BaseSingleStageBrokerRequestHandler + implements FailureDetector.Listener { private static final Logger LOGGER = LoggerFactory.getLogger(SingleConnectionBrokerRequestHandler.class); private final BrokerReduceService _brokerReduceService; @@ -74,16 +72,12 @@ public class SingleConnectionBrokerRequestHandler extends BaseBrokerRequestHandl public SingleConnectionBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, - QueryQuotaManager queryQuotaManager, TableCache tableCache, BrokerMetrics brokerMetrics, NettyConfig nettyConfig, - TlsConfig tlsConfig, ServerRoutingStatsManager serverRoutingStatsManager, - BrokerQueryEventListener brokerQueryEventListener) { - super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, brokerMetrics, - brokerQueryEventListener); - LOGGER.info("Using Netty BrokerRequestHandler."); - + QueryQuotaManager queryQuotaManager, TableCache tableCache, NettyConfig nettyConfig, TlsConfig tlsConfig, + ServerRoutingStatsManager serverRoutingStatsManager) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); _brokerReduceService = new BrokerReduceService(_config); - _queryRouter = new QueryRouter(_brokerId, brokerMetrics, nettyConfig, tlsConfig, serverRoutingStatsManager); - _failureDetector = FailureDetectorFactory.getFailureDetector(config, brokerMetrics); + _queryRouter = new QueryRouter(_brokerId, _brokerMetrics, nettyConfig, tlsConfig, serverRoutingStatsManager); + _failureDetector = FailureDetectorFactory.getFailureDetector(config, _brokerMetrics); } @Override @@ -93,7 +87,7 @@ public void start() { } @Override - public synchronized void shutDown() { + public void shutDown() { _failureDetector.stop(); _queryRouter.shutDown(); _brokerReduceService.shutDown(); @@ -130,7 +124,7 @@ protected BrokerResponseNative processBrokerRequest(long requestId, BrokerReques int numServersQueried = finalResponses.size(); long totalResponseSize = 0; - Map dataTableMap = new HashMap<>(HashUtil.getHashMapCapacity(numServersQueried)); + Map dataTableMap = Maps.newHashMapWithExpectedSize(numServersQueried); List serversNotResponded = new ArrayList<>(); for (Map.Entry entry : finalResponses.entrySet()) { ServerResponse serverResponse = entry.getValue(); @@ -145,27 +139,26 @@ protected BrokerResponseNative processBrokerRequest(long requestId, BrokerReques int numServersResponded = dataTableMap.size(); long reduceStartTimeNs = System.nanoTime(); - long reduceTimeOutMs = timeoutMs - TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - scatterGatherStartTimeNs); + long reduceTimeoutMs = timeoutMs - TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - scatterGatherStartTimeNs); BrokerResponseNative brokerResponse = _brokerReduceService.reduceOnDataTable(originalBrokerRequest, serverBrokerRequest, dataTableMap, - reduceTimeOutMs, _brokerMetrics); - final long reduceTimeNanos = System.nanoTime() - reduceStartTimeNs; - requestContext.setTraceInfo(brokerResponse.getTraceInfo()); - requestContext.setReduceTimeNanos(reduceTimeNanos); + reduceTimeoutMs, _brokerMetrics); + long reduceTimeNanos = System.nanoTime() - reduceStartTimeNs; _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.REDUCE, reduceTimeNanos); brokerResponse.setNumServersQueried(numServersQueried); brokerResponse.setNumServersResponded(numServersResponded); + brokerResponse.setBrokerReduceTimeMs(TimeUnit.NANOSECONDS.toMillis(reduceTimeNanos)); Exception brokerRequestSendException = asyncQueryResponse.getException(); if (brokerRequestSendException != null) { String errorMsg = QueryException.getTruncatedStackTrace(brokerRequestSendException); - brokerResponse.addToExceptions( + brokerResponse.addException( new QueryProcessingException(QueryException.BROKER_REQUEST_SEND_ERROR_CODE, errorMsg)); } int numServersNotResponded = serversNotResponded.size(); if (numServersNotResponded != 0) { - brokerResponse.addToExceptions(new QueryProcessingException(QueryException.SERVER_NOT_RESPONDING_ERROR_CODE, + brokerResponse.addException(new QueryProcessingException(QueryException.SERVER_NOT_RESPONDING_ERROR_CODE, String.format("%d servers %s not responded", numServersNotResponded, serversNotResponded))); _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_PARTIAL_SERVERS_RESPONDED, 1); } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java index c2e6b20cce54..59aa65406da8 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java @@ -253,7 +253,7 @@ private List getFilterTimeIntervals(Expression filterExpression) { return null; case RANGE: if (isTimeColumn(operands.get(0))) { - return parseInterval(operands.get(1).getLiteral().getFieldValue().toString()); + return parseInterval(operands.get(1).getLiteral().getStringValue()); } return null; default: @@ -375,7 +375,22 @@ private boolean isTimeColumn(Expression expression) { private long toMillisSinceEpoch(Expression expression) { Literal literal = expression.getLiteral(); Preconditions.checkArgument(literal != null, "Literal is required for time column filter, got: %s", expression); - return _timeFormatSpec.fromFormatToMillis(literal.getFieldValue().toString()); + String value; + Literal._Fields type = literal.getSetField(); + switch (type) { + case INT_VALUE: + value = Integer.toString(literal.getIntValue()); + break; + case LONG_VALUE: + value = Long.toString(literal.getLongValue()); + break; + case STRING_VALUE: + value = literal.getStringValue(); + break; + default: + throw new IllegalStateException("Unsupported literal type: " + type + " as time column filter"); + } + return _timeFormatSpec.fromFormatToMillis(value); } /** diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/api/AccessControlBackwardCompatibleTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/api/AccessControlBackwardCompatibleTest.java new file mode 100644 index 000000000000..ad781c9f885c --- /dev/null +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/api/AccessControlBackwardCompatibleTest.java @@ -0,0 +1,97 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.broker.api; + +import java.util.Set; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.spi.auth.AuthorizationResult; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; + + +public class AccessControlBackwardCompatibleTest { + + @Test + public void testBackwardCompatibleHasAccessBrokerRequest() { + AccessControl accessControl = new AllFalseAccessControlImpl(); + HttpRequesterIdentity identity = new HttpRequesterIdentity(); + BrokerRequest request = new BrokerRequest(); + assertFalse(accessControl.authorize(identity, request).hasAccess()); + } + + @Test + public void testBackwardCompatibleHasAccessMutliTable() { + AccessControl accessControl = new AllFalseAccessControlImpl(); + HttpRequesterIdentity identity = new HttpRequesterIdentity(); + Set tables = Set.of("table1", "table2"); + AuthorizationResult result = accessControl.authorize(identity, tables); + assertFalse(result.hasAccess()); + assertEquals(result.getFailureMessage(), "Authorization Failed for tables: [table1, table2]"); + } + + @Test(expectedExceptions = UnsupportedOperationException.class) + public void testExceptionForNoImplAccessControlMultiTable() { + AccessControl accessControl = new NoImplAccessControl(); + HttpRequesterIdentity identity = new HttpRequesterIdentity(); + Set tables = Set.of("table1", "table2"); + accessControl.hasAccess(identity, tables); + } + + @Test(expectedExceptions = UnsupportedOperationException.class) + public void testExceptionForNoImplAccessControlBrokerRequest() { + AccessControl accessControl = new NoImplAccessControl(); + HttpRequesterIdentity identity = new HttpRequesterIdentity(); + BrokerRequest request = new BrokerRequest(); + accessControl.hasAccess(identity, request); + } + + @Test(expectedExceptions = UnsupportedOperationException.class) + public void testExceptionForNoImplAccessControlAuthorizeMultiTable() { + AccessControl accessControl = new NoImplAccessControl(); + HttpRequesterIdentity identity = new HttpRequesterIdentity(); + Set tables = Set.of("table1", "table2"); + accessControl.authorize(identity, tables); + } + + @Test(expectedExceptions = UnsupportedOperationException.class) + public void testExceptionForNoImplAccessControlAuthorizeBrokerRequest() { + AccessControl accessControl = new NoImplAccessControl(); + HttpRequesterIdentity identity = new HttpRequesterIdentity(); + BrokerRequest request = new BrokerRequest(); + accessControl.authorize(identity, request); + } + + class AllFalseAccessControlImpl implements AccessControl { + + @Override + public boolean hasAccess(RequesterIdentity requesterIdentity, BrokerRequest brokerRequest) { + return false; + } + + @Override + public boolean hasAccess(RequesterIdentity requesterIdentity, Set tables) { + return false; + } + } + + class NoImplAccessControl implements AccessControl { + } +} diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/BasicAuthAccessControlTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/BasicAuthAccessControlTest.java index 42f121ecefd1..5eb444431aaa 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/BasicAuthAccessControlTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/BasicAuthAccessControlTest.java @@ -29,6 +29,7 @@ import org.apache.pinot.broker.api.HttpRequesterIdentity; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.QuerySource; +import org.apache.pinot.spi.auth.AuthorizationResult; import org.apache.pinot.spi.env.PinotConfiguration; import org.testng.Assert; import org.testng.annotations.BeforeClass; @@ -66,7 +67,7 @@ public void setup() { @Test(expectedExceptions = IllegalArgumentException.class) public void testNullEntity() { - _accessControl.hasAccess(null, (BrokerRequest) null); + _accessControl.authorize(null, (BrokerRequest) null); } @Test @@ -77,7 +78,7 @@ public void testNullToken() { identity.setHttpHeaders(headers); try { - _accessControl.hasAccess(identity, (BrokerRequest) null); + _accessControl.authorize(identity, (BrokerRequest) null); } catch (WebApplicationException e) { Assert.assertEquals(e.getResponse().getStatus(), 401, "must return 401"); } @@ -97,8 +98,8 @@ public void testAllow() { BrokerRequest request = new BrokerRequest(); request.setQuerySource(source); - Assert.assertTrue(_accessControl.hasAccess(identity, request)); - Assert.assertTrue(_accessControl.hasAccess(identity, _tableNames)); + Assert.assertTrue(_accessControl.authorize(identity, request).hasAccess()); + Assert.assertTrue(_accessControl.authorize(identity, _tableNames).hasAccess()); } @Test @@ -114,16 +115,27 @@ public void testDeny() { BrokerRequest request = new BrokerRequest(); request.setQuerySource(source); - - Assert.assertFalse(_accessControl.hasAccess(identity, request)); + AuthorizationResult authorizationResult = _accessControl.authorize(identity, request); + Assert.assertFalse(authorizationResult.hasAccess()); + Assert.assertEquals(authorizationResult.getFailureMessage(), + "Authorization Failed for tables: [veryImportantStuff]"); Set tableNames = new HashSet<>(); tableNames.add("veryImportantStuff"); - Assert.assertFalse(_accessControl.hasAccess(identity, tableNames)); + authorizationResult = _accessControl.authorize(identity, tableNames); + Assert.assertFalse(authorizationResult.hasAccess()); + Assert.assertEquals(authorizationResult.getFailureMessage(), + "Authorization Failed for tables: [veryImportantStuff]"); tableNames.add("lessImportantStuff"); - Assert.assertFalse(_accessControl.hasAccess(identity, tableNames)); + authorizationResult = _accessControl.authorize(identity, tableNames); + Assert.assertFalse(authorizationResult.hasAccess()); + Assert.assertEquals(authorizationResult.getFailureMessage(), + "Authorization Failed for tables: [veryImportantStuff]"); tableNames.add("lesserImportantStuff"); - Assert.assertFalse(_accessControl.hasAccess(identity, tableNames)); + authorizationResult = _accessControl.authorize(identity, tableNames); + Assert.assertFalse(authorizationResult.hasAccess()); + Assert.assertEquals(authorizationResult.getFailureMessage(), + "Authorization Failed for tables: [veryImportantStuff]"); } @Test @@ -139,15 +151,18 @@ public void testAllowAll() { BrokerRequest request = new BrokerRequest(); request.setQuerySource(source); - - Assert.assertTrue(_accessControl.hasAccess(identity, request)); + AuthorizationResult authorizationResult = _accessControl.authorize(identity, request); + Assert.assertTrue(authorizationResult.hasAccess()); + Assert.assertEquals(authorizationResult.getFailureMessage(), ""); Set tableNames = new HashSet<>(); tableNames.add("lessImportantStuff"); tableNames.add("veryImportantStuff"); tableNames.add("lesserImportantStuff"); - Assert.assertTrue(_accessControl.hasAccess(identity, tableNames)); + authorizationResult = _accessControl.authorize(identity, tableNames); + Assert.assertTrue(authorizationResult.hasAccess()); + Assert.assertEquals(authorizationResult.getFailureMessage(), ""); } @Test @@ -159,11 +174,12 @@ public void testAllowNonTable() { identity.setHttpHeaders(headers); BrokerRequest request = new BrokerRequest(); - - Assert.assertTrue(_accessControl.hasAccess(identity, request)); + AuthorizationResult authorizationResult = _accessControl.authorize(identity, request); + Assert.assertTrue(authorizationResult.hasAccess()); Set tableNames = new HashSet<>(); - Assert.assertTrue(_accessControl.hasAccess(identity, tableNames)); + authorizationResult = _accessControl.authorize(identity, tableNames); + Assert.assertTrue(authorizationResult.hasAccess()); } @Test @@ -180,7 +196,7 @@ public void testNormalizeToken() { BrokerRequest request = new BrokerRequest(); request.setQuerySource(source); - Assert.assertTrue(_accessControl.hasAccess(identity, request)); - Assert.assertTrue(_accessControl.hasAccess(identity, _tableNames)); + Assert.assertTrue(_accessControl.authorize(identity, request).hasAccess()); + Assert.assertTrue(_accessControl.authorize(identity, _tableNames).hasAccess()); } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/querylog/QueryLoggerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/querylog/QueryLoggerTest.java index bf173f615b63..1e2d909a45cc 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/querylog/QueryLoggerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/querylog/QueryLoggerTest.java @@ -25,10 +25,8 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; -import java.util.stream.Collectors; -import java.util.stream.IntStream; import org.apache.pinot.broker.api.RequesterIdentity; -import org.apache.pinot.broker.requesthandler.BaseBrokerRequestHandler; +import org.apache.pinot.broker.requesthandler.BaseSingleStageBrokerRequestHandler.ServerStats; import org.apache.pinot.common.response.ProcessingException; import org.apache.pinot.common.response.broker.BrokerResponseNative; import org.apache.pinot.spi.trace.DefaultRequestContext; @@ -78,7 +76,8 @@ public void setUp() { } @AfterMethod - public void tearDown() throws Exception { + public void tearDown() + throws Exception { _closeMocks.close(); } @@ -94,23 +93,25 @@ public void shouldFormatLogLineProperly() { // Then: Assert.assertEquals(_infoLog.size(), 1); + //@formatter:off Assert.assertEquals(_infoLog.get(0), "requestId=123," + "table=table," + "timeMs=456," + "docs=1/2," + "entries=3/4," + "segments(queried/processed/matched/consumingQueried/consumingProcessed/consumingMatched/unavailable)" - + ":5/6/7/8/9/10/24," + + ":5/6/7/8/9/10/21," + "consumingFreshnessTimeMs=11," + "servers=12/13," + "groupLimitReached=false," - + "brokerReduceTimeMs=22," + + "brokerReduceTimeMs=20," + "exceptions=0," + "serverStats=serverStats," - + "offlineThreadCpuTimeNs(total/thread/sysActivity/resSer):14/15/16/17," - + "realtimeThreadCpuTimeNs(total/thread/sysActivity/resSer):18/19/20/21," + + "offlineThreadCpuTimeNs(total/thread/sysActivity/resSer):45/14/15/16," + + "realtimeThreadCpuTimeNs(total/thread/sysActivity/resSer):54/17/18/19," + "clientIp=ip," + "query=SELECT * FROM foo"); + //@formatter:on } @Test @@ -125,8 +126,7 @@ public void shouldOmitClientId() { // Then: Assert.assertEquals(_infoLog.size(), 1); - Assert.assertFalse( - _infoLog.get(0).contains("clientId"), + Assert.assertFalse(_infoLog.get(0).contains("clientId"), "did not expect to see clientId Logs. Got: " + _infoLog.get(0)); } @@ -191,8 +191,7 @@ public void shouldHandleRaceConditionsWithDroppedQueries() throws InterruptedException { // Given: final CountDownLatch logLatch = new CountDownLatch(1); - Mockito.when(_logRateLimiter.tryAcquire()) - .thenReturn(false) + Mockito.when(_logRateLimiter.tryAcquire()).thenReturn(false) .thenReturn(true) // this one will block when it hits tryAcquire() .thenReturn(false) // this one just increments the dropped logs .thenAnswer(invocation -> { @@ -205,12 +204,11 @@ public void shouldHandleRaceConditionsWithDroppedQueries() // ensure that the tryAcquire only succeeds after three other // logs have went through (see logAndDecrement) final CountDownLatch dropLogLatch = new CountDownLatch(3); - Mockito.when(_droppedRateLimiter.tryAcquire()) - .thenAnswer(invocation -> { - logLatch.countDown(); - dropLogLatch.await(); - return true; - }).thenReturn(true); + Mockito.when(_droppedRateLimiter.tryAcquire()).thenAnswer(invocation -> { + logLatch.countDown(); + dropLogLatch.await(); + return true; + }).thenReturn(true); QueryLogger.QueryLogParams params = generateParams(false, 0, 456); QueryLogger queryLogger = new QueryLogger(_logRateLimiter, 100, true, _logger, _droppedRateLimiter); @@ -240,8 +238,18 @@ public void shouldHandleRaceConditionsWithDroppedQueries() Assert.assertEquals((long) _numDropped.get(0), 2L); } - private QueryLogger.QueryLogParams generateParams(boolean isGroupLimitHit, int numExceptions, long timeUsed) { + private QueryLogger.QueryLogParams generateParams(boolean numGroupsLimitReached, int numExceptions, long timeUsedMs) { + RequestContext requestContext = new DefaultRequestContext(); + requestContext.setRequestId(123); + requestContext.setQuery("SELECT * FROM foo"); + requestContext.setNumUnavailableSegments(21); + BrokerResponseNative response = new BrokerResponseNative(); + response.setNumGroupsLimitReached(numGroupsLimitReached); + for (int i = 0; i < numExceptions; i++) { + response.addException(new ProcessingException()); + } + response.setTimeUsedMs(timeUsedMs); response.setNumDocsScanned(1); response.setTotalDocs(2); response.setNumEntriesScannedInFilter(3); @@ -255,40 +263,24 @@ private QueryLogger.QueryLogParams generateParams(boolean isGroupLimitHit, int n response.setMinConsumingFreshnessTimeMs(11); response.setNumServersResponded(12); response.setNumServersQueried(13); - response.setNumGroupsLimitReached(isGroupLimitHit); - response.setExceptions( - IntStream.range(0, numExceptions) - .mapToObj(i -> new ProcessingException()).collect(Collectors.toList())); - response.setOfflineTotalCpuTimeNs(14); - response.setOfflineThreadCpuTimeNs(15); - response.setOfflineSystemActivitiesCpuTimeNs(16); - response.setOfflineResponseSerializationCpuTimeNs(17); - response.setRealtimeTotalCpuTimeNs(18); - response.setRealtimeThreadCpuTimeNs(19); - response.setRealtimeSystemActivitiesCpuTimeNs(20); - response.setRealtimeResponseSerializationCpuTimeNs(21); - - RequestContext request = new DefaultRequestContext(); - request.setReduceTimeMillis(22); - - BaseBrokerRequestHandler.ServerStats serverStats = new BaseBrokerRequestHandler.ServerStats(); - serverStats.setServerStats("serverStats"); + response.setOfflineThreadCpuTimeNs(14); + response.setOfflineSystemActivitiesCpuTimeNs(15); + response.setOfflineResponseSerializationCpuTimeNs(16); + response.setRealtimeThreadCpuTimeNs(17); + response.setRealtimeSystemActivitiesCpuTimeNs(18); + response.setRealtimeResponseSerializationCpuTimeNs(19); + response.setBrokerReduceTimeMs(20); + RequesterIdentity identity = new RequesterIdentity() { - @Override public String getClientIp() { + @Override + public String getClientIp() { return "ip"; } }; - return new QueryLogger.QueryLogParams( - 123, - "SELECT * FROM foo", - request, - "table", - 24, - serverStats, - response, - timeUsed, - identity - ); + ServerStats serverStats = new ServerStats(); + serverStats.setServerStats("serverStats"); + + return new QueryLogger.QueryLogParams(requestContext, "table", response, identity, serverStats); } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/MaxHitRateTrackerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/MaxHitRateTrackerTest.java index b4d9cc5f936e..f19c55504dbf 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/MaxHitRateTrackerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/MaxHitRateTrackerTest.java @@ -37,24 +37,29 @@ public void testMaxHitRateTracker() { long latestTimeStamp = currentTimestamp + (timeInSec - 1) * 1000; Assert.assertNotNull(hitCounter); Assert.assertEquals(5, hitCounter.getMaxCountPerBucket(latestTimeStamp)); + Assert.assertEquals(5 * 60, hitCounter.getHitCount(latestTimeStamp)); // 2 seconds have passed, the hit counter should return 5 as well since the count in the last bucket could increase. latestTimeStamp = latestTimeStamp + 2000L; Assert.assertEquals(5, hitCounter.getMaxCountPerBucket(latestTimeStamp)); + Assert.assertEquals(5 * (60 - 2), hitCounter.getHitCount(latestTimeStamp)); // This time it should return 0 as the internal lastAccessTimestamp has already been updated and there is no more // hits between the gap. latestTimeStamp = latestTimeStamp + 2000L; Assert.assertEquals(0, hitCounter.getMaxCountPerBucket(latestTimeStamp)); + Assert.assertEquals(5 * (60 - 4), hitCounter.getHitCount(latestTimeStamp)); // Increment the hit in this second and we should see the result becomes 1. hitCounter.hit(latestTimeStamp); latestTimeStamp = latestTimeStamp + 2000L; Assert.assertEquals(1, hitCounter.getMaxCountPerBucket(latestTimeStamp)); + Assert.assertEquals(5 * (60 - 6) + 1, hitCounter.getHitCount(latestTimeStamp)); // More than a time range period has passed and the hit counter should return 0 as there is no hits. hitCounter.hit(latestTimeStamp); latestTimeStamp = latestTimeStamp + timeInSec * 2 * 1000L + 2000L; Assert.assertEquals(0, hitCounter.getMaxCountPerBucket(latestTimeStamp)); + Assert.assertEquals(0, hitCounter.getHitCount(latestTimeStamp)); } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java similarity index 70% rename from pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandlerTest.java rename to pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java index 7bb85737f0d4..f1a6dfe33f11 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandlerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java @@ -18,9 +18,6 @@ */ package org.apache.pinot.broker.requesthandler; -import com.fasterxml.jackson.databind.JsonNode; -import com.google.common.collect.ImmutableMap; -import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -28,8 +25,6 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import javax.annotation.Nullable; -import javax.ws.rs.core.HttpHeaders; -import javax.ws.rs.core.MultivaluedHashMap; import org.apache.commons.lang3.tuple.Pair; import org.apache.helix.model.InstanceConfig; import org.apache.pinot.broker.broker.AllowAllAccessControlFactory; @@ -46,11 +41,10 @@ import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TenantConfig; import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; import org.apache.pinot.spi.exception.BadQueryRequestException; -import org.apache.pinot.spi.metrics.PinotMetricUtils; import org.apache.pinot.spi.trace.RequestContext; -import org.apache.pinot.spi.trace.Tracing; -import org.apache.pinot.spi.utils.JsonUtils; +import org.apache.pinot.spi.utils.CommonConstants.Broker; import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.apache.pinot.util.TestUtils; import org.mockito.Mockito; @@ -63,15 +57,15 @@ import static org.mockito.Mockito.when; -public class BaseBrokerRequestHandlerTest { +public class BaseSingleStageBrokerRequestHandlerTest { @Test public void testUpdateColumnNames() { String query = "SELECT database.my_table.column_name_1st, column_name_2nd from database.my_table"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); Map columnNameMap = - ImmutableMap.of("column_name_1st", "column_name_1st", "column_name_2nd", "column_name_2nd"); - BaseBrokerRequestHandler.updateColumnNames("database.my_table", pinotQuery, false, columnNameMap); + Map.of("column_name_1st", "column_name_1st", "column_name_2nd", "column_name_2nd"); + BaseSingleStageBrokerRequestHandler.updateColumnNames("database.my_table", pinotQuery, false, columnNameMap); Assert.assertEquals(pinotQuery.getSelectList().size(), 2); for (Expression expression : pinotQuery.getSelectList()) { String columnName = expression.getIdentifier().getName(); @@ -90,17 +84,18 @@ public void testGetActualColumnNameCaseSensitive() { Map columnNameMap = new HashMap<>(); columnNameMap.put("student_name", "student_name"); String actualColumnName = - BaseBrokerRequestHandler.getActualColumnName("mytable", "mytable.student_name", columnNameMap, false); + BaseSingleStageBrokerRequestHandler.getActualColumnName("mytable", "mytable.student_name", columnNameMap, + false); Assert.assertEquals(actualColumnName, "student_name"); Assert.assertEquals( - BaseBrokerRequestHandler.getActualColumnName("db1.mytable", "db1.mytable.student_name", columnNameMap, false), - "student_name"); + BaseSingleStageBrokerRequestHandler.getActualColumnName("db1.mytable", "db1.mytable.student_name", + columnNameMap, false), "student_name"); Assert.assertEquals( - BaseBrokerRequestHandler.getActualColumnName("db1.mytable", "mytable.student_name", columnNameMap, false), - "student_name"); + BaseSingleStageBrokerRequestHandler.getActualColumnName("db1.mytable", "mytable.student_name", columnNameMap, + false), "student_name"); boolean exceptionThrown = false; try { - BaseBrokerRequestHandler.getActualColumnName("mytable", "mytable2.student_name", columnNameMap, false); + BaseSingleStageBrokerRequestHandler.getActualColumnName("mytable", "mytable2.student_name", columnNameMap, false); Assert.fail("should throw exception if column is not known"); } catch (BadQueryRequestException ex) { exceptionThrown = true; @@ -108,7 +103,7 @@ public void testGetActualColumnNameCaseSensitive() { Assert.assertTrue(exceptionThrown, "should throw exception if column is not known"); exceptionThrown = false; try { - BaseBrokerRequestHandler.getActualColumnName("mytable", "MYTABLE.student_name", columnNameMap, false); + BaseSingleStageBrokerRequestHandler.getActualColumnName("mytable", "MYTABLE.student_name", columnNameMap, false); Assert.fail("should throw exception if case sensitive and table name different"); } catch (BadQueryRequestException ex) { exceptionThrown = true; @@ -116,12 +111,13 @@ public void testGetActualColumnNameCaseSensitive() { Assert.assertTrue(exceptionThrown, "should throw exception if column is not known"); columnNameMap.put("mytable_student_name", "mytable_student_name"); String wrongColumnName2 = - BaseBrokerRequestHandler.getActualColumnName("mytable", "mytable_student_name", columnNameMap, false); + BaseSingleStageBrokerRequestHandler.getActualColumnName("mytable", "mytable_student_name", columnNameMap, + false); Assert.assertEquals(wrongColumnName2, "mytable_student_name"); columnNameMap.put("mytable", "mytable"); String wrongColumnName3 = - BaseBrokerRequestHandler.getActualColumnName("mytable", "mytable", columnNameMap, false); + BaseSingleStageBrokerRequestHandler.getActualColumnName("mytable", "mytable", columnNameMap, false); Assert.assertEquals(wrongColumnName3, "mytable"); } @@ -130,17 +126,17 @@ public void testGetActualColumnNameCaseInSensitive() { Map columnNameMap = new HashMap<>(); columnNameMap.put("student_name", "student_name"); String actualColumnName = - BaseBrokerRequestHandler.getActualColumnName("mytable", "MYTABLE.student_name", columnNameMap, true); + BaseSingleStageBrokerRequestHandler.getActualColumnName("mytable", "MYTABLE.student_name", columnNameMap, true); Assert.assertEquals(actualColumnName, "student_name"); Assert.assertEquals( - BaseBrokerRequestHandler.getActualColumnName("db1.MYTABLE", "DB1.mytable.student_name", columnNameMap, true), - "student_name"); + BaseSingleStageBrokerRequestHandler.getActualColumnName("db1.MYTABLE", "DB1.mytable.student_name", + columnNameMap, true), "student_name"); Assert.assertEquals( - BaseBrokerRequestHandler.getActualColumnName("db1.mytable", "MYTABLE.student_name", columnNameMap, true), - "student_name"); + BaseSingleStageBrokerRequestHandler.getActualColumnName("db1.mytable", "MYTABLE.student_name", columnNameMap, + true), "student_name"); boolean exceptionThrown = false; try { - BaseBrokerRequestHandler.getActualColumnName("student", "MYTABLE2.student_name", columnNameMap, true); + BaseSingleStageBrokerRequestHandler.getActualColumnName("student", "MYTABLE2.student_name", columnNameMap, true); Assert.fail("should throw exception if column is not known"); } catch (BadQueryRequestException ex) { exceptionThrown = true; @@ -148,47 +144,43 @@ public void testGetActualColumnNameCaseInSensitive() { Assert.assertTrue(exceptionThrown, "should throw exception if column is not known"); columnNameMap.put("mytable_student_name", "mytable_student_name"); String wrongColumnName2 = - BaseBrokerRequestHandler.getActualColumnName("mytable", "MYTABLE_student_name", columnNameMap, true); + BaseSingleStageBrokerRequestHandler.getActualColumnName("mytable", "MYTABLE_student_name", columnNameMap, true); Assert.assertEquals(wrongColumnName2, "mytable_student_name"); columnNameMap.put("mytable", "mytable"); String wrongColumnName3 = - BaseBrokerRequestHandler.getActualColumnName("MYTABLE", "mytable", columnNameMap, true); + BaseSingleStageBrokerRequestHandler.getActualColumnName("MYTABLE", "mytable", columnNameMap, true); Assert.assertEquals(wrongColumnName3, "mytable"); } @Test - public void testCancelQuery() - throws Exception { + public void testCancelQuery() { String tableName = "myTable_OFFLINE"; // Mock pretty much everything until the query can be submitted. TableCache tableCache = mock(TableCache.class); TableConfig tableCfg = mock(TableConfig.class); when(tableCache.getActualTableName(anyString())).thenReturn(tableName); - HttpHeaders headers = mock(HttpHeaders.class); - when(headers.getRequestHeaders()).thenReturn(new MultivaluedHashMap<>()); - when(headers.getHeaderString(anyString())).thenReturn(null); TenantConfig tenant = new TenantConfig("tier_BROKER", "tier_SERVER", null); when(tableCfg.getTenantConfig()).thenReturn(tenant); - when(tableCache.getTableConfig(anyString())).thenReturn(tableCfg); + when(tableCache.getTableConfig(tableName)).thenReturn(tableCfg); BrokerRoutingManager routingManager = mock(BrokerRoutingManager.class); - when(routingManager.routingExists(anyString())).thenReturn(true); + when(routingManager.routingExists(tableName)).thenReturn(true); + when(routingManager.getQueryTimeoutMs(tableName)).thenReturn(10000L); RoutingTable rt = mock(RoutingTable.class); when(rt.getServerInstanceToSegmentsMap()).thenReturn( - Collections.singletonMap(new ServerInstance(new InstanceConfig("server01_9000")), - Pair.of(Collections.singletonList("segment01"), Collections.emptyList()))); + Map.of(new ServerInstance(new InstanceConfig("server01_9000")), Pair.of(List.of("segment01"), List.of()))); when(routingManager.getRoutingTable(any(), Mockito.anyLong())).thenReturn(rt); QueryQuotaManager queryQuotaManager = mock(QueryQuotaManager.class); when(queryQuotaManager.acquire(anyString())).thenReturn(true); CountDownLatch latch = new CountDownLatch(1); - final long[] testRequestId = {-1}; + long[] testRequestId = {-1}; + BrokerMetrics.register(mock(BrokerMetrics.class)); PinotConfiguration config = - new PinotConfiguration(Collections.singletonMap("pinot.broker.enable.query.cancellation", "true")); - BaseBrokerRequestHandler requestHandler = - new BaseBrokerRequestHandler(config, "testBrokerId", routingManager, new AllowAllAccessControlFactory(), - queryQuotaManager, tableCache, - new BrokerMetrics("", PinotMetricUtils.getPinotMetricsRegistry(), true, - Collections.emptySet()), null) { + new PinotConfiguration(Map.of(Broker.CONFIG_OF_BROKER_ENABLE_QUERY_CANCELLATION, "true")); + BrokerQueryEventListenerFactory.init(config); + BaseSingleStageBrokerRequestHandler requestHandler = + new BaseSingleStageBrokerRequestHandler(config, "testBrokerId", routingManager, + new AllowAllAccessControlFactory(), queryQuotaManager, tableCache) { @Override public void start() { } @@ -212,10 +204,7 @@ protected BrokerResponseNative processBrokerRequest(long requestId, BrokerReques }; CompletableFuture.runAsync(() -> { try { - JsonNode request = JsonUtils.stringToJsonNode( - String.format("{\"sql\":\"select * from %s limit 10\",\"queryOptions\":\"timeoutMs=10000\"}", tableName)); - RequestContext requestStats = Tracing.getTracer().createRequestScope(); - requestHandler.handleRequest(request, null, requestStats, headers); + requestHandler.handleRequest(String.format("select * from %s limit 10", tableName)); } catch (Exception e) { throw new RuntimeException(e); } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BrokerRequestOptionsTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BrokerRequestOptionsTest.java index 14d1d39bdba4..e6ebccd16641 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BrokerRequestOptionsTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BrokerRequestOptionsTest.java @@ -26,68 +26,40 @@ import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.apache.pinot.sql.parsers.SqlNodeAndOptions; -import org.testng.Assert; import org.testng.annotations.Test; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + /** * Tests the various options set in the broker request */ public class BrokerRequestOptionsTest { - // TODO: remove this legacy option size checker after 0.11 release cut. - private static final int LEGACY_PQL_QUERY_OPTION_SIZE = 2; @Test public void testSetOptions() { - long requestId = 1; String query = "select * from testTable"; // None of the options ObjectNode jsonRequest = JsonUtils.newObjectNode(); - SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query);; + SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 0 + LEGACY_PQL_QUERY_OPTION_SIZE); + assertTrue(sqlNodeAndOptions.getOptions().isEmpty()); // TRACE // Has trace false jsonRequest.put(Request.TRACE, false); sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 0 + LEGACY_PQL_QUERY_OPTION_SIZE); + assertTrue(sqlNodeAndOptions.getOptions().isEmpty()); // Has trace true jsonRequest.put(Request.TRACE, true); sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 1 + LEGACY_PQL_QUERY_OPTION_SIZE); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get(Request.TRACE), "true"); - - // DEBUG_OPTIONS (debug options will also be included as query options) - // Has debugOptions - jsonRequest = JsonUtils.newObjectNode(); - jsonRequest.put(Request.DEBUG_OPTIONS, "debugOption1=foo"); - sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); - RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 1 + LEGACY_PQL_QUERY_OPTION_SIZE); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("debugOption1"), "foo"); - - // Has multiple debugOptions - jsonRequest.put(Request.DEBUG_OPTIONS, "debugOption1=foo;debugOption2=bar"); - sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); - RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 2 + LEGACY_PQL_QUERY_OPTION_SIZE); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("debugOption1"), "foo"); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("debugOption2"), "bar"); - - // Invalid debug options - jsonRequest.put(Request.DEBUG_OPTIONS, "debugOption1"); - sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); - try { - RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.fail(); - } catch (Exception e) { - // Expected - } + assertEquals(sqlNodeAndOptions.getOptions().size(), 1); + assertEquals(sqlNodeAndOptions.getOptions().get(Request.TRACE), "true"); // QUERY_OPTIONS jsonRequest = JsonUtils.newObjectNode(); @@ -97,41 +69,28 @@ public void testSetOptions() { queryOptions.put("queryOption1", "foo"); sqlNodeAndOptions.getOptions().putAll(queryOptions); RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 1 + LEGACY_PQL_QUERY_OPTION_SIZE); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "foo"); + assertEquals(sqlNodeAndOptions.getOptions().size(), 1); + assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "foo"); // Has queryOptions in query sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions("SET queryOption1='foo'; select * from testTable"); RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 1 + LEGACY_PQL_QUERY_OPTION_SIZE); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "foo"); + assertEquals(sqlNodeAndOptions.getOptions().size(), 1); + assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "foo"); // Has query options in json payload jsonRequest.put(Request.QUERY_OPTIONS, "queryOption1=foo"); sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 1 + LEGACY_PQL_QUERY_OPTION_SIZE); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "foo"); + assertEquals(sqlNodeAndOptions.getOptions().size(), 1); + assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "foo"); // Has query options in both json payload and sqlNodeAndOptions, sqlNodeAndOptions takes priority jsonRequest.put(Request.QUERY_OPTIONS, "queryOption1=bar;queryOption2=moo"); sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions("SET queryOption1='foo'; select * from testTable;"); RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 2 + LEGACY_PQL_QUERY_OPTION_SIZE); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "foo"); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("queryOption2"), "moo"); - - // Has all 3 - jsonRequest = JsonUtils.newObjectNode(); - jsonRequest.put(Request.TRACE, true); - jsonRequest.put(Request.DEBUG_OPTIONS, "debugOption1=foo"); - jsonRequest.put(Request.QUERY_OPTIONS, "queryOption1=bar;queryOption2=moo"); - sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); - RequestUtils.setOptions(sqlNodeAndOptions, jsonRequest); - Assert.assertEquals(sqlNodeAndOptions.getOptions().size(), 4 + LEGACY_PQL_QUERY_OPTION_SIZE); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "bar"); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("queryOption2"), "moo"); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get(Request.TRACE), "true"); - Assert.assertEquals(sqlNodeAndOptions.getOptions().get("debugOption1"), "foo"); + assertEquals(sqlNodeAndOptions.getOptions().size(), 2); + assertEquals(sqlNodeAndOptions.getOptions().get("queryOption1"), "foo"); + assertEquals(sqlNodeAndOptions.getOptions().get("queryOption2"), "moo"); } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java index 749afeb233b6..0b68d2b843ba 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java @@ -18,8 +18,6 @@ */ package org.apache.pinot.broker.requesthandler; -import com.fasterxml.jackson.databind.JsonNode; -import java.util.Collections; import java.util.List; import java.util.Random; import java.util.concurrent.TimeUnit; @@ -31,17 +29,17 @@ import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; import org.apache.pinot.spi.env.PinotConfiguration; -import org.apache.pinot.spi.eventlistener.query.PinotBrokerQueryEventListenerFactory; -import org.apache.pinot.spi.metrics.PinotMetricUtils; -import org.apache.pinot.spi.trace.RequestContext; -import org.apache.pinot.spi.trace.Tracing; +import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; import org.apache.pinot.spi.utils.BytesUtils; -import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.sql.parsers.CalciteSqlParser; -import org.testng.Assert; +import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static org.apache.pinot.broker.requesthandler.BaseSingleStageBrokerRequestHandler.isLiteralOnlyQuery; import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; public class LiteralOnlyBrokerRequestTest { @@ -49,132 +47,120 @@ public class LiteralOnlyBrokerRequestTest { private static final Random RANDOM = new Random(System.currentTimeMillis()); private static final long ONE_HOUR_IN_MS = TimeUnit.HOURS.toMillis(1); + @BeforeClass + public void setUp() { + BrokerMetrics.register(mock(BrokerMetrics.class)); + BrokerQueryEventListenerFactory.init(new PinotConfiguration()); + } + @Test public void testStringLiteralBrokerRequestFromSQL() { - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 'a'"))); - Assert.assertTrue( - BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 'a', 'b'"))); - Assert.assertTrue( - BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 'a' FROM myTable"))); - Assert.assertTrue(BaseBrokerRequestHandler - .isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 'a', 'b' FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 'a'"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 'a', 'b'"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 'a' FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 'a', 'b' FROM myTable"))); } @Test public void testSelectStarBrokerRequestFromSQL() { - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT '*'"))); - Assert.assertTrue( - BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT '*' FROM myTable"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT *"))); - Assert.assertFalse( - BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT * FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT '*'"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT '*' FROM myTable"))); + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT *"))); + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT * FROM myTable"))); } @Test public void testNumberLiteralBrokerRequestFromSQL() { - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 1"))); - Assert.assertTrue( - BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 1, '2', 3"))); - Assert.assertTrue( - BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 1 FROM myTable"))); - Assert.assertTrue(BaseBrokerRequestHandler - .isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 1, '2', 3 FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 1"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 1, '2', 3"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 1 FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT 1, '2', 3 FROM myTable"))); } @Test public void testLiteralOnlyTransformBrokerRequestFromSQL() { - Assert - .assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT now()"))); - Assert.assertTrue( - BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT ago('PT1H')"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT now()"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT ago('PT1H')"))); + assertTrue(isLiteralOnlyQuery( CalciteSqlParser.compileToPinotQuery("SELECT now(), fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z')"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery( + assertTrue(isLiteralOnlyQuery( CalciteSqlParser.compileToPinotQuery("SELECT ago('PT1H'), fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z')"))); - Assert.assertTrue( - BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT now() FROM myTable"))); - Assert.assertTrue(BaseBrokerRequestHandler - .isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT ago('PT1H') FROM myTable"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT now(), fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') FROM myTable"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT ago('PT1H'), fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') FROM myTable"))); - Assert.assertFalse(BaseBrokerRequestHandler - .isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT count(*) from foo where bar > ago('PT1H')"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT encodeUrl('key1=value 1&key2=value@!$2&key3=value%3')," + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT now() FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT ago('PT1H') FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT now(), fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT ago('PT1H'), fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') FROM myTable"))); + assertFalse( + isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT count(*) from foo where bar > ago('PT1H')"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT encodeUrl('key1=value 1&key2=value@!$2&key3=value%3')," + " decodeUrl('key1%3Dvalue+1%26key2%3Dvalue%40%21%242%26key3%3Dvalue%253') FROM myTable"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT count(*) from foo " - + "where bar = encodeUrl('key1=value 1&key2=value@!$2&key3=value%3')"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT count(*) from foo " - + "where bar = decodeUrl('key1%3Dvalue+1%26key2%3Dvalue%40%21%242%26key3%3Dvalue%253')"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( - "SELECT toUtf8('hello!')," + " fromUtf8(toUtf8('hello!')) FROM myTable"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( - "SELECT reverse(fromUtf8(foo))," + " toUtf8('hello!') FROM myTable"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT count(*) from foo " + "where bar = encodeUrl('key1=value 1&key2=value@!$2&key3=value%3')"))); + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("SELECT count(*) from foo " + + "where bar = decodeUrl('key1%3Dvalue+1%26key2%3Dvalue%40%21%242%26key3%3Dvalue%253')"))); + assertTrue(isLiteralOnlyQuery( + CalciteSqlParser.compileToPinotQuery("SELECT toUtf8('hello!')," + " fromUtf8(toUtf8('hello!')) FROM myTable"))); + assertFalse(isLiteralOnlyQuery( + CalciteSqlParser.compileToPinotQuery("SELECT reverse(fromUtf8(foo))," + " toUtf8('hello!') FROM myTable"))); + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "SELECT toBase64(toUtf8('hello!'))," + " fromBase64('aGVsbG8h') FROM myTable"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "SELECT reverse(toBase64(foo))," + " toBase64(fromBase64('aGVsbG8h')) FROM myTable"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery( + assertFalse(isLiteralOnlyQuery( CalciteSqlParser.compileToPinotQuery("SELECT fromBase64(toBase64(to_utf8(foo))) FROM myTable"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery( + assertFalse(isLiteralOnlyQuery( CalciteSqlParser.compileToPinotQuery("SELECT count(*) from foo " + "where bar = toBase64(toASCII('hello!'))"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery( + assertFalse(isLiteralOnlyQuery( CalciteSqlParser.compileToPinotQuery("SELECT count(*) from foo " + "where bar = fromBase64('aGVsbG8h')"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "SELECT count(*) from foo " + "where bar = fromUtf8(fromBase64('aGVsbG8h'))"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\") = " - + "\"fooXarYXazY\""))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10) = " - + "\"fooXarYXazY\""))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10 , 1) = " + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\") = " + "\"fooXarYXazY\""))); + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10) = " + "\"fooXarYXazY\""))); + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10 , 1) = " + "\"fooXarYXazY\""))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10 , 1, " - + "\"i\") = " + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10 , 1, " + "\"i\") = " + "\"fooXarYXazY\""))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser - .compileToPinotQuery("SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10 , 1, " - + "\"m\") = " + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + "SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10 , 1, " + "\"m\") = " + "\"fooXarYXazY\""))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery( + assertTrue(isLiteralOnlyQuery( CalciteSqlParser.compileToPinotQuery("select isSubnetOf('1.2.3.128/0', '192.168.5.1') from mytable"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "select isSubnetOf('1.2.3.128/0', rtrim('192.168.5.1 ')) from mytable"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "select isSubnetOf('123:db8:85a3::8a2e:370:7334/72', '124:db8:85a3::8a2e:370:7334') from mytable"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery( - CalciteSqlParser.compileToPinotQuery("select isSubnetOf('1.2.3.128/0', foo) from mytable"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertFalse( + isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery("select isSubnetOf('1.2.3.128/0', foo) from mytable"))); + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "select count(*) from mytable where isSubnetOf('7890:db8:113::8a2e:370:7334/127', ltrim(' " + "7890:db8:113::8a2e:370:7336'))"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertFalse(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "select count(*) from mytable where isSubnetOf('7890:db8:113::8a2e:370:7334/127', " + "'7890:db8:113::8a2e:370:7336')"))); - Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery( + assertFalse(isLiteralOnlyQuery( CalciteSqlParser.compileToPinotQuery("select count(*) from mytable where isSubnetOf(foo, bar)"))); } @Test public void testLiteralOnlyWithAsBrokerRequestFromSQL() { - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "SELECT now() AS currentTs, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') AS firstDayOf2020"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "SELECT ago('PT1H') AS currentTs, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') AS firstDayOf2020"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "SELECT encodeUrl('key1=value 1&key2=value@!$2&key3=value%3') AS encoded, " + "decodeUrl('key1%3Dvalue+1%26key2%3Dvalue%40%21%242%26key3%3Dvalue%253') AS decoded"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "SELECT toUtf8('hello!') AS encoded, " + "fromUtf8(toUtf8('hello!')) AS decoded"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "SELECT toBase64(toUtf8('hello!')) AS encoded, " + "fromBase64('aGVsbG8h') AS decoded"))); - Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( + assertTrue(isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery( "select isSubnetOf('1.2.3.128/0', '192.168.5.1') AS booleanCol from mytable"))); } @@ -183,28 +169,23 @@ public void testBrokerRequestHandler() throws Exception { SingleConnectionBrokerRequestHandler requestHandler = new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY, - null, null, new BrokerMetrics("", PinotMetricUtils.getPinotMetricsRegistry(), true, Collections.emptySet()), - null, null, mock(ServerRoutingStatsManager.class), - PinotBrokerQueryEventListenerFactory.getBrokerQueryEventListener()); + null, null, null, null, mock(ServerRoutingStatsManager.class)); long randNum = RANDOM.nextLong(); byte[] randBytes = new byte[12]; RANDOM.nextBytes(randBytes); String ranStr = BytesUtils.toHexString(randBytes); - JsonNode request = JsonUtils.stringToJsonNode(String.format("{\"sql\":\"SELECT %d, '%s'\"}", randNum, ranStr)); - RequestContext requestStats = Tracing.getTracer().createRequestScope(); - BrokerResponse brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(0), String.format("%d", randNum)); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(0), - DataSchema.ColumnDataType.LONG); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(1), ranStr); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(1), + BrokerResponse brokerResponse = requestHandler.handleRequest(String.format("SELECT %d, '%s'", randNum, ranStr)); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(0), String.format("%d", randNum)); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.LONG); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(1), String.format("'%s'", ranStr)); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(1), DataSchema.ColumnDataType.STRING); - Assert.assertEquals(brokerResponse.getResultTable().getRows().size(), 1); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0).length, 2); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0)[0], randNum); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0)[1], ranStr); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); + assertEquals(brokerResponse.getResultTable().getRows().size(), 1); + assertEquals(brokerResponse.getResultTable().getRows().get(0).length, 2); + assertEquals(brokerResponse.getResultTable().getRows().get(0)[0], randNum); + assertEquals(brokerResponse.getResultTable().getRows().get(0)[1], ranStr); + assertEquals(brokerResponse.getTotalDocs(), 0); } @Test @@ -212,206 +193,152 @@ public void testBrokerRequestHandlerWithAsFunction() throws Exception { SingleConnectionBrokerRequestHandler requestHandler = new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY, - null, null, new BrokerMetrics("", PinotMetricUtils.getPinotMetricsRegistry(), true, Collections.emptySet()), - null, null, mock(ServerRoutingStatsManager.class), - PinotBrokerQueryEventListenerFactory.getBrokerQueryEventListener()); + null, null, null, null, mock(ServerRoutingStatsManager.class)); long currentTsMin = System.currentTimeMillis(); - JsonNode request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT now() as currentTs, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') as firstDayOf2020\"}"); - RequestContext requestStats = Tracing.getTracer().createRequestScope(); - BrokerResponse brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + BrokerResponse brokerResponse = requestHandler.handleRequest( + "SELECT now() AS currentTs, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') AS firstDayOf2020"); long currentTsMax = System.currentTimeMillis(); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(0), "currentTs"); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(0), - DataSchema.ColumnDataType.LONG); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(1), "firstDayOf2020"); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(1), - DataSchema.ColumnDataType.LONG); - Assert.assertEquals(brokerResponse.getResultTable().getRows().size(), 1); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0).length, 2); - Assert.assertTrue(Long.parseLong(brokerResponse.getResultTable().getRows().get(0)[0].toString()) > currentTsMin); - Assert.assertTrue(Long.parseLong(brokerResponse.getResultTable().getRows().get(0)[0].toString()) < currentTsMax); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0)[1], 1577836800000L); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(0), "currentTs"); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.LONG); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(1), "firstDayOf2020"); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(1), DataSchema.ColumnDataType.LONG); + assertEquals(brokerResponse.getResultTable().getRows().size(), 1); + assertEquals(brokerResponse.getResultTable().getRows().get(0).length, 2); + assertTrue(Long.parseLong(brokerResponse.getResultTable().getRows().get(0)[0].toString()) > currentTsMin); + assertTrue(Long.parseLong(brokerResponse.getResultTable().getRows().get(0)[0].toString()) < currentTsMax); + assertEquals(brokerResponse.getResultTable().getRows().get(0)[1], 1577836800000L); + assertEquals(brokerResponse.getTotalDocs(), 0); long oneHourAgoTsMin = System.currentTimeMillis() - ONE_HOUR_IN_MS; - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT ago('PT1H') as oneHourAgoTs, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') as " - + "firstDayOf2020\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + brokerResponse = requestHandler.handleRequest( + "SELECT ago('PT1H') AS oneHourAgoTs, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') AS firstDayOf2020"); long oneHourAgoTsMax = System.currentTimeMillis() - ONE_HOUR_IN_MS; - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(0), "oneHourAgoTs"); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(0), - DataSchema.ColumnDataType.LONG); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(1), "firstDayOf2020"); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(1), - DataSchema.ColumnDataType.LONG); - Assert.assertEquals(brokerResponse.getResultTable().getRows().size(), 1); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0).length, 2); - Assert - .assertTrue(Long.parseLong(brokerResponse.getResultTable().getRows().get(0)[0].toString()) >= oneHourAgoTsMin); - Assert - .assertTrue(Long.parseLong(brokerResponse.getResultTable().getRows().get(0)[0].toString()) <= oneHourAgoTsMax); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0)[1], 1577836800000L); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); - - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT encodeUrl('key1=value 1&key2=value@!$2&key3=value%3') AS encoded, " - + "decodeUrl('key1%3Dvalue+1%26key2%3Dvalue%40%21%242%26key3%3Dvalue%253') AS decoded\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - System.out.println(brokerResponse.getResultTable()); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(0), "encoded"); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(0), + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(0), "oneHourAgoTs"); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.LONG); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(1), "firstDayOf2020"); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(1), DataSchema.ColumnDataType.LONG); + assertEquals(brokerResponse.getResultTable().getRows().size(), 1); + assertEquals(brokerResponse.getResultTable().getRows().get(0).length, 2); + assertTrue(Long.parseLong(brokerResponse.getResultTable().getRows().get(0)[0].toString()) >= oneHourAgoTsMin); + assertTrue(Long.parseLong(brokerResponse.getResultTable().getRows().get(0)[0].toString()) <= oneHourAgoTsMax); + assertEquals(brokerResponse.getResultTable().getRows().get(0)[1], 1577836800000L); + assertEquals(brokerResponse.getTotalDocs(), 0); + + brokerResponse = requestHandler.handleRequest( + "SELECT encodeUrl('key1=value 1&key2=value@!$2&key3=value%3') AS encoded, " + + "decodeUrl('key1%3Dvalue+1%26key2%3Dvalue%40%21%242%26key3%3Dvalue%253') AS decoded"); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(0), "encoded"); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.STRING); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(1), "decoded"); - Assert.assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(1), + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnName(1), "decoded"); + assertEquals(brokerResponse.getResultTable().getDataSchema().getColumnDataType(1), DataSchema.ColumnDataType.STRING); - Assert.assertEquals(brokerResponse.getResultTable().getRows().size(), 1); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0).length, 2); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0)[0].toString(), + assertEquals(brokerResponse.getResultTable().getRows().size(), 1); + assertEquals(brokerResponse.getResultTable().getRows().get(0).length, 2); + assertEquals(brokerResponse.getResultTable().getRows().get(0)[0].toString(), "key1%3Dvalue+1%26key2%3Dvalue%40%21%242%26key3%3Dvalue%253"); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0)[1].toString(), + assertEquals(brokerResponse.getResultTable().getRows().get(0)[1].toString(), "key1=value 1&key2=value@!$2&key3=value%3"); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); + assertEquals(brokerResponse.getTotalDocs(), 0); - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT toBase64(toUtf8('hello!')) AS encoded, " + "fromUtf8(fromBase64('aGVsbG8h')) AS decoded\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + brokerResponse = requestHandler.handleRequest( + "SELECT toBase64(toUtf8('hello!')) AS encoded, fromUtf8(fromBase64('aGVsbG8h')) AS decoded"); ResultTable resultTable = brokerResponse.getResultTable(); DataSchema dataSchema = resultTable.getDataSchema(); List rows = resultTable.getRows(); - Assert.assertEquals(dataSchema.getColumnName(0), "encoded"); - Assert.assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.STRING); - Assert.assertEquals(dataSchema.getColumnName(1), "decoded"); - Assert.assertEquals(dataSchema.getColumnDataType(1), DataSchema.ColumnDataType.STRING); - Assert.assertEquals(rows.size(), 1); - Assert.assertEquals(rows.get(0).length, 2); - Assert.assertEquals(rows.get(0)[0].toString(), "aGVsbG8h"); - Assert.assertEquals(rows.get(0)[1].toString(), "hello!"); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); - - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT fromUtf8(fromBase64(toBase64(toUtf8('nested')))) AS output\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + assertEquals(dataSchema.getColumnName(0), "encoded"); + assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.STRING); + assertEquals(dataSchema.getColumnName(1), "decoded"); + assertEquals(dataSchema.getColumnDataType(1), DataSchema.ColumnDataType.STRING); + assertEquals(rows.size(), 1); + assertEquals(rows.get(0).length, 2); + assertEquals(rows.get(0)[0].toString(), "aGVsbG8h"); + assertEquals(rows.get(0)[1].toString(), "hello!"); + assertEquals(brokerResponse.getTotalDocs(), 0); + + brokerResponse = requestHandler.handleRequest("SELECT fromUtf8(fromBase64(toBase64(toUtf8('nested')))) AS output"); resultTable = brokerResponse.getResultTable(); dataSchema = resultTable.getDataSchema(); rows = resultTable.getRows(); - Assert.assertEquals(dataSchema.getColumnName(0), "output"); - Assert.assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.STRING); - Assert.assertEquals(rows.size(), 1); - Assert.assertEquals(rows.get(0).length, 1); - Assert.assertEquals(rows.get(0)[0].toString(), "nested"); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); - - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT toBase64(toUtf8('this is a long string that will encode to more than 76 characters using " - + "base64'))" - + " AS encoded\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + assertEquals(dataSchema.getColumnName(0), "output"); + assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.STRING); + assertEquals(rows.size(), 1); + assertEquals(rows.get(0).length, 1); + assertEquals(rows.get(0)[0].toString(), "nested"); + assertEquals(brokerResponse.getTotalDocs(), 0); + + brokerResponse = requestHandler.handleRequest( + "SELECT toBase64(toUtf8('this is a long string that will encode to more than 76 characters using base64')) " + + "AS encoded"); resultTable = brokerResponse.getResultTable(); dataSchema = resultTable.getDataSchema(); rows = resultTable.getRows(); - Assert.assertEquals(dataSchema.getColumnName(0), "encoded"); - Assert.assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.STRING); - Assert.assertEquals(rows.size(), 1); - Assert.assertEquals(rows.get(0).length, 1); - Assert.assertEquals(rows.get(0)[0].toString(), + assertEquals(dataSchema.getColumnName(0), "encoded"); + assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.STRING); + assertEquals(rows.size(), 1); + assertEquals(rows.get(0).length, 1); + assertEquals(rows.get(0)[0].toString(), "dGhpcyBpcyBhIGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBlbmNvZGUgdG8gbW9yZSB0aGFuIDc2IGNoYXJhY3RlcnMgdXNpbmcgYmFzZTY0"); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); + assertEquals(brokerResponse.getTotalDocs(), 0); - request = JsonUtils.stringToJsonNode("{\"sql\":\"SELECT fromUtf8(fromBase64" - + "('dGhpcyBpcyBhIGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBlbmNvZGUgdG8gbW9yZSB0aGFuIDc2IGNoYXJhY3RlcnMgdXNpbmcgYmFzZTY0" - + "')) AS decoded\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + brokerResponse = requestHandler.handleRequest("SELECT fromUtf8(fromBase64(" + + "'dGhpcyBpcyBhIGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBlbmNvZGUgdG8gbW9yZSB0aGFuIDc2IGNoYXJhY3RlcnMgdXNpbmcgYmFzZTY0'" + + ")) AS decoded"); resultTable = brokerResponse.getResultTable(); dataSchema = resultTable.getDataSchema(); rows = resultTable.getRows(); - Assert.assertEquals(dataSchema.getColumnName(0), "decoded"); - Assert.assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.STRING); - Assert.assertEquals(rows.size(), 1); - Assert.assertEquals(rows.get(0).length, 1); - Assert.assertEquals(rows.get(0)[0].toString(), + assertEquals(dataSchema.getColumnName(0), "decoded"); + assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.STRING); + assertEquals(rows.size(), 1); + assertEquals(rows.get(0).length, 1); + assertEquals(rows.get(0)[0].toString(), "this is a long string that will encode to more than 76 characters using base64"); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); - - request = JsonUtils.stringToJsonNode("{\"sql\":\"SELECT fromBase64" + "(0) AS decoded\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - Assert.assertTrue( - brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException")); - - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT isSubnetOf('2001:db8:85a3::8a2e:370:7334/62', '2001:0db8:85a3:0003:ffff:ffff:ffff:ffff')" - + " as booleanCol\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + assertEquals(brokerResponse.getTotalDocs(), 0); + + brokerResponse = requestHandler.handleRequest("SELECT fromBase64(0) AS decoded"); + assertTrue(brokerResponse.getExceptions().get(0).getMessage().contains("IllegalArgumentException")); + + brokerResponse = requestHandler.handleRequest( + "SELECT isSubnetOf('2001:db8:85a3::8a2e:370:7334/62', '2001:0db8:85a3:0003:ffff:ffff:ffff:ffff') " + + "AS booleanCol"); resultTable = brokerResponse.getResultTable(); dataSchema = resultTable.getDataSchema(); rows = resultTable.getRows(); - Assert.assertEquals(dataSchema.getColumnName(0), "booleanCol"); - Assert.assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.BOOLEAN); - Assert.assertEquals(rows.size(), 1); - Assert.assertEquals(rows.get(0).length, 1); - Assert.assertTrue((boolean) rows.get(0)[0]); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); + assertEquals(dataSchema.getColumnName(0), "booleanCol"); + assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.BOOLEAN); + assertEquals(rows.size(), 1); + assertEquals(rows.get(0).length, 1); + assertTrue((boolean) rows.get(0)[0]); + assertEquals(brokerResponse.getTotalDocs(), 0); // first argument must be in prefix format - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT isSubnetOf('2001:db8:85a3::8a2e:370:7334', '2001:0db8:85a3:0003:ffff:ffff:ffff:ffff') as" - + " booleanCol\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - Assert.assertTrue( - brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException")); + brokerResponse = requestHandler.handleRequest( + "SELECT isSubnetOf('2001:db8:85a3::8a2e:370:7334', '2001:0db8:85a3:0003:ffff:ffff:ffff:ffff') AS booleanCol"); + assertTrue(brokerResponse.getExceptions().get(0).getMessage().contains("IllegalArgumentException")); // first argument must be in prefix format - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT isSubnetOf('105.25.245.115', '105.25.245.115') as" + " booleanCol\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - Assert.assertTrue( - brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException")); + brokerResponse = + requestHandler.handleRequest("SELECT isSubnetOf('105.25.245.115', '105.25.245.115') AS booleanCol"); + assertTrue(brokerResponse.getExceptions().get(0).getMessage().contains("IllegalArgumentException")); // second argument should not be a prefix - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT isSubnetOf('1.2.3.128/26', '3.175.47.239/26') as" + " booleanCol\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - Assert.assertTrue( - brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException")); + brokerResponse = requestHandler.handleRequest("SELECT isSubnetOf('1.2.3.128/26', '3.175.47.239/26') AS booleanCol"); + assertTrue(brokerResponse.getExceptions().get(0).getMessage().contains("IllegalArgumentException")); // second argument should not be a prefix - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT isSubnetOf('5f3f:bfdb:1bbe:a824:6bf9:0fbb:d358:1889/64', " - + "'4275:386f:b2b5:0664:04aa:d7bd:0589:6909/64') as" - + " booleanCol\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - Assert.assertTrue( - brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException")); + brokerResponse = requestHandler.handleRequest("SELECT isSubnetOf('5f3f:bfdb:1bbe:a824:6bf9:0fbb:d358:1889/64', " + + "'4275:386f:b2b5:0664:04aa:d7bd:0589:6909/64') AS booleanCol"); + assertTrue(brokerResponse.getExceptions().get(0).getMessage().contains("IllegalArgumentException")); // invalid prefix length - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT isSubnetOf('2001:4801:7825:103:be76:4eff::/129', '2001:4801:7825:103:be76:4eff::') as" - + " booleanCol\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - Assert.assertTrue( - brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException")); + brokerResponse = requestHandler.handleRequest( + "SELECT isSubnetOf('2001:4801:7825:103:be76:4eff::/129', '2001:4801:7825:103:be76:4eff::') AS booleanCol"); + assertTrue(brokerResponse.getExceptions().get(0).getMessage().contains("IllegalArgumentException")); // invalid prefix length - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"SELECT isSubnetOf('170.189.0.175/33', '170.189.0.175') as" + " booleanCol\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); - Assert.assertTrue( - brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException")); + brokerResponse = + requestHandler.handleRequest("SELECT isSubnetOf('170.189.0.175/33', '170.189.0.175') AS booleanCol"); + assertTrue(brokerResponse.getExceptions().get(0).getMessage().contains("IllegalArgumentException")); } /** Tests for EXPLAIN PLAN for literal only queries. */ @@ -420,54 +347,40 @@ public void testExplainPlanLiteralOnly() throws Exception { SingleConnectionBrokerRequestHandler requestHandler = new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY, - null, null, new BrokerMetrics("", PinotMetricUtils.getPinotMetricsRegistry(), true, Collections.emptySet()), - null, null, mock(ServerRoutingStatsManager.class), - PinotBrokerQueryEventListenerFactory.getBrokerQueryEventListener()); + null, null, null, null, mock(ServerRoutingStatsManager.class)); // Test 1: select constant - JsonNode request = JsonUtils.stringToJsonNode("{\"sql\":\"EXPLAIN PLAN FOR SELECT 1.5, 'test'\"}"); - RequestContext requestStats = Tracing.getTracer().createRequestScope(); - BrokerResponse brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + BrokerResponse brokerResponse = requestHandler.handleRequest("EXPLAIN PLAN FOR SELECT 1.5, 'test'"); checkExplainResultSchema(brokerResponse.getResultTable().getDataSchema(), - new String[]{"Operator", "Operator_Id", "Parent_Id"}, - new DataSchema.ColumnDataType[]{ - DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.INT, - DataSchema.ColumnDataType.INT + new String[]{"Operator", "Operator_Id", "Parent_Id"}, new DataSchema.ColumnDataType[]{ + DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.INT }); - Assert.assertEquals(brokerResponse.getResultTable().getRows().size(), 1); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0), - new Object[]{"BROKER_EVALUATE", 0, -1}); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); + assertEquals(brokerResponse.getResultTable().getRows().size(), 1); + assertEquals(brokerResponse.getResultTable().getRows().get(0), new Object[]{"BROKER_EVALUATE", 0, -1}); + assertEquals(brokerResponse.getTotalDocs(), 0); // Test 2: invoke compile time function -> literal only - long currentTsMin = System.currentTimeMillis(); - request = JsonUtils.stringToJsonNode( - "{\"sql\":\"EXPLAIN PLAN FOR SELECT 6+8 as addition, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') as " - + "firstDayOf2020\"}"); - requestStats = Tracing.getTracer().createRequestScope(); - brokerResponse = requestHandler.handleRequest(request, null, requestStats, null); + requestHandler.handleRequest( + "EXPLAIN PLAN FOR SELECT 6+8 AS addition, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') AS firstDayOf2020"); checkExplainResultSchema(brokerResponse.getResultTable().getDataSchema(), - new String[]{"Operator", "Operator_Id", "Parent_Id"}, - new DataSchema.ColumnDataType[]{ - DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.INT, - DataSchema.ColumnDataType.INT + new String[]{"Operator", "Operator_Id", "Parent_Id"}, new DataSchema.ColumnDataType[]{ + DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.INT }); - Assert.assertEquals(brokerResponse.getResultTable().getRows().size(), 1); - Assert.assertEquals(brokerResponse.getResultTable().getRows().get(0), - new Object[]{"BROKER_EVALUATE", 0, -1}); + assertEquals(brokerResponse.getResultTable().getRows().size(), 1); + assertEquals(brokerResponse.getResultTable().getRows().get(0), new Object[]{"BROKER_EVALUATE", 0, -1}); - Assert.assertEquals(brokerResponse.getTotalDocs(), 0); + assertEquals(brokerResponse.getTotalDocs(), 0); } private void checkExplainResultSchema(DataSchema schema, String[] columnNames, DataSchema.ColumnDataType[] columnTypes) { for (int i = 0; i < columnNames.length; i++) { - Assert.assertEquals(schema.getColumnName(i), columnNames[i]); - Assert.assertEquals(schema.getColumnDataType(i), columnTypes[i]); + assertEquals(schema.getColumnName(i), columnNames[i]); + assertEquals(schema.getColumnDataType(i), columnTypes[i]); } } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java deleted file mode 100644 index 45d41ffc6f95..000000000000 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java +++ /dev/null @@ -1,104 +0,0 @@ -/** - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package org.apache.pinot.broker.requesthandler; - -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import java.util.ArrayList; -import java.util.List; -import org.apache.pinot.broker.broker.AccessControlFactory; -import org.apache.pinot.broker.broker.AllowAllAccessControlFactory; -import org.apache.pinot.broker.queryquota.QueryQuotaManager; -import org.apache.pinot.broker.routing.BrokerRoutingManager; -import org.apache.pinot.common.config.provider.TableCache; -import org.apache.pinot.common.metrics.BrokerMetrics; -import org.apache.pinot.spi.env.PinotConfiguration; -import org.apache.pinot.spi.eventlistener.query.PinotBrokerQueryEventListenerFactory; -import org.apache.pinot.spi.trace.DefaultRequestContext; -import org.apache.pinot.spi.trace.RequestContext; -import org.apache.pinot.spi.utils.CommonConstants; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import org.testng.Assert; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; - - -public class MultiStageBrokerRequestHandlerTest { - - private PinotConfiguration _config; - @Mock - private BrokerRoutingManager _routingManager; - - private AccessControlFactory _accessControlFactory; - @Mock - private QueryQuotaManager _queryQuotaManager; - @Mock - private TableCache _tableCache; - - @Mock - private BrokerMetrics _brokerMetrics; - - private MultiStageBrokerRequestHandler _requestHandler; - - @BeforeClass - public void setUp() { - MockitoAnnotations.openMocks(this); - _config = new PinotConfiguration(); - _config.setProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME, "localhost"); - _config.setProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, "12345"); - _accessControlFactory = new AllowAllAccessControlFactory(); - _requestHandler = - new MultiStageBrokerRequestHandler(_config, "Broker_localhost", _routingManager, - _accessControlFactory, _queryQuotaManager, _tableCache, _brokerMetrics, - PinotBrokerQueryEventListenerFactory.getBrokerQueryEventListener()); - } - - @Test - public void testSetRequestId() - throws Exception { - String sampleSqlQuery = "SELECT * FROM testTable"; - String sampleJsonRequest = String.format("{\"sql\":\"%s\"}", sampleSqlQuery); - ObjectMapper objectMapper = new ObjectMapper(); - JsonNode jsonRequest = objectMapper.readTree(sampleJsonRequest); - RequestContext requestContext = new DefaultRequestContext(); - - List requestIds = new ArrayList<>(); - // Request id should be unique each time, and there should be a difference of 1 between consecutive requestIds. - for (int iteration = 0; iteration < 10; iteration++) { - _requestHandler.handleRequest(jsonRequest, null, null, requestContext, null); - Assert.assertTrue(requestContext.getRequestId() >= 0, "Request ID should be non-negative"); - requestIds.add(requestContext.getRequestId()); - if (iteration != 0) { - Assert.assertEquals(1, requestIds.get(iteration) - requestIds.get(iteration - 1), - "Request Id should have difference of 1"); - } - } - Assert.assertEquals(10, requestIds.stream().distinct().count(), "Request Id should be unique"); - Assert.assertEquals(1, requestIds.stream().map(x -> (x >> 32)).distinct().count(), - "Request Id should have a broker-id specific mask for the 32 MSB"); - Assert.assertTrue(requestIds.stream().noneMatch(x -> x < 0), "Request Id should not be negative"); - } - - @AfterClass - public void tearDown() { - _requestHandler.shutDown(); - } -} diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryOverrideTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryOverrideTest.java index df9e07b44694..22137c137bb6 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryOverrideTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryOverrideTest.java @@ -50,7 +50,7 @@ public void testLimitOverride() { private void testLimitOverride(String query, int expectedLimit) { PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleQueryLimitOverride(pinotQuery, QUERY_LIMIT); + BaseSingleStageBrokerRequestHandler.handleQueryLimitOverride(pinotQuery, QUERY_LIMIT); assertEquals(pinotQuery.getLimit(), expectedLimit); } @@ -58,9 +58,10 @@ private void testLimitOverride(String query, int expectedLimit) { public void testDistinctCountOverride() { String query = "SELECT DISTINCT_COUNT(col1) FROM myTable"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleSegmentPartitionedDistinctCountOverride(pinotQuery, ImmutableSet.of("col2", "col3")); + BaseSingleStageBrokerRequestHandler.handleSegmentPartitionedDistinctCountOverride(pinotQuery, + ImmutableSet.of("col2", "col3")); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcount"); - BaseBrokerRequestHandler.handleSegmentPartitionedDistinctCountOverride(pinotQuery, + BaseSingleStageBrokerRequestHandler.handleSegmentPartitionedDistinctCountOverride(pinotQuery, ImmutableSet.of("col1", "col2", "col3")); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "segmentpartitioneddistinctcount"); } @@ -77,9 +78,9 @@ public void testDistinctMultiValuedOverride() + "{\"name\":\"col3\",\"dataType\":\"LONG\",\"singleValueField\":\"false\"}]," + "\"dateTimeFieldSpecs\":[{\"name\":\"dt1\",\"dataType\":\"INT\",\"format\":\"x:HOURS:EPOCH\"," + "\"granularity\":\"1:HOURS\"}]}"); - BaseBrokerRequestHandler.handleDistinctMultiValuedOverride(pinotQuery1, tableSchema); + BaseSingleStageBrokerRequestHandler.handleDistinctMultiValuedOverride(pinotQuery1, tableSchema); assertEquals(pinotQuery1.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcount"); - BaseBrokerRequestHandler.handleDistinctMultiValuedOverride(pinotQuery2, tableSchema); + BaseSingleStageBrokerRequestHandler.handleDistinctMultiValuedOverride(pinotQuery2, tableSchema); assertEquals(pinotQuery2.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcountmv"); } @@ -89,7 +90,7 @@ public void testApproximateFunctionOverride() { String query = "SELECT DISTINCT_COUNT(col1) FROM myTable GROUP BY col2 HAVING DISTINCT_COUNT(col1) > 10 " + "ORDER BY DISTINCT_COUNT(col1) DESC"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); + BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcountsmarthll"); assertEquals( pinotQuery.getOrderByList().get(0).getFunctionCall().getOperands().get(0).getFunctionCall().getOperator(), @@ -100,22 +101,22 @@ public void testApproximateFunctionOverride() { query = "SELECT DISTINCT_COUNT_MV(col1) FROM myTable"; pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); + BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcountsmarthll"); query = "SELECT DISTINCT col1 FROM myTable"; pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); + BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinct"); query = "SELECT DISTINCT_COUNT_HLL(col1) FROM myTable"; pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); + BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcounthll"); query = "SELECT DISTINCT_COUNT_BITMAP(col1) FROM myTable"; pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); + BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcountbitmap"); } @@ -123,7 +124,7 @@ public void testApproximateFunctionOverride() { "SELECT PERCENTILE_MV(col1, 95) FROM myTable", "SELECT PERCENTILE95(col1) FROM myTable", "SELECT PERCENTILE95MV(col1) FROM myTable")) { PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); + BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "percentilesmarttdigest"); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().get(1), RequestUtils.getLiteralExpression(95)); @@ -131,12 +132,12 @@ public void testApproximateFunctionOverride() { { String query = "SELECT PERCENTILE_TDIGEST(col1, 95) FROM myTable"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); + BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "percentiletdigest"); query = "SELECT PERCENTILE_EST(col1, 95) FROM myTable"; pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); + BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride(pinotQuery); assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "percentileest"); } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java index 4772c0ddacb3..f46833b0783b 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java @@ -115,16 +115,17 @@ public void testRejectGroovyQuery() { @Test public void testReplicaGroupToQueryInvalidQuery() { - PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery( - "SET numReplicaGroupsToQuery='illegal'; SELECT COUNT(*) FROM MY_TABLE"); - Assert.assertThrows(IllegalStateException.class, () -> BaseBrokerRequestHandler.validateRequest(pinotQuery, 10)); + PinotQuery pinotQuery = + CalciteSqlParser.compileToPinotQuery("SET numReplicaGroupsToQuery='illegal'; SELECT COUNT(*) FROM MY_TABLE"); + Assert.assertThrows(IllegalStateException.class, + () -> BaseSingleStageBrokerRequestHandler.validateRequest(pinotQuery, 10)); } private void testRejectGroovyQuery(String query, boolean queryContainsGroovy) { PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); try { - BaseBrokerRequestHandler.rejectGroovyQuery(pinotQuery); + BaseSingleStageBrokerRequestHandler.rejectGroovyQuery(pinotQuery); if (queryContainsGroovy) { Assert.fail("Query should have failed since groovy was found in query: " + pinotQuery); } @@ -136,7 +137,7 @@ private void testRejectGroovyQuery(String query, boolean queryContainsGroovy) { private void testUnsupportedQuery(String query, String errorMessage) { try { PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.validateRequest(pinotQuery, 1000); + BaseSingleStageBrokerRequestHandler.validateRequest(pinotQuery, 1000); Assert.fail("Query should have failed"); } catch (Exception e) { Assert.assertEquals(e.getMessage(), errorMessage); @@ -147,7 +148,7 @@ private void testNonExistingColumns(String rawTableName, boolean isCaseInsensiti String query, String errorMessage) { try { PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap); + BaseSingleStageBrokerRequestHandler.updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap); Assert.fail("Query should have failed"); } catch (Exception e) { Assert.assertEquals(errorMessage, e.getMessage()); @@ -158,7 +159,7 @@ private void testExistingColumns(String rawTableName, boolean isCaseInsensitive, String query) { try { PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); - BaseBrokerRequestHandler.updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap); + BaseSingleStageBrokerRequestHandler.updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap); } catch (Exception e) { Assert.fail("Query should have succeeded"); } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/SelectStarWithOtherColsRewriteTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/SelectStarWithOtherColsRewriteTest.java index 4c490b94ac31..3b8c6ac7cd2b 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/SelectStarWithOtherColsRewriteTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/SelectStarWithOtherColsRewriteTest.java @@ -57,7 +57,7 @@ public class SelectStarWithOtherColsRewriteTest { public void testShouldExpandWhenOnlyStarIsSelected() { String sql = "SELECT * FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); Map countMap = new HashMap<>(); for (Expression selection : newSelections) { @@ -79,7 +79,7 @@ public void testShouldExpandWhenOnlyStarIsSelected() { public void testShouldNotReturnExtraDefaultColumns() { String sql = "SELECT $docId,*,$segmentName FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); int docIdCnt = 0; int segmentNameCnt = 0; @@ -108,7 +108,7 @@ public void testShouldNotReturnExtraDefaultColumns() { public void testShouldNotDedupMultipleRequestedColumns() { String sql = "SELECT playerID,*,G_old FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); int playerIdCnt = 0; int goldCount = 0; @@ -135,7 +135,7 @@ public void testShouldNotDedupMultipleRequestedColumns() { public void testSelectionOrder() { String sql = "SELECT playerID,*,G_old FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); Assert.assertEquals(newSelections.get(0).getIdentifier().getName(), "playerID"); Assert.assertEquals(newSelections.get(newSelections.size() - 1).getIdentifier().getName(), "G_old"); @@ -154,7 +154,7 @@ public void testSelectionOrder() { public void testAliasing() { String sql = "SELECT playerID as pid,* FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); Assert.assertTrue(newSelections.get(0).isSetFunctionCall()); Assert.assertEquals(newSelections.get(0).getFunctionCall().getOperator(), "as"); @@ -178,7 +178,7 @@ public void testAliasing() { public void testFuncOnColumns1() { String sql = "SELECT sqrt(homeRuns),* FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); Assert.assertTrue(newSelections.get(0).isSetFunctionCall()); Assert.assertEquals(newSelections.get(0).getFunctionCall().getOperator(), "sqrt"); @@ -201,7 +201,7 @@ public void testFuncOnColumns1() { public void testFuncOnColumns2() { String sql = "SELECT add(homeRuns,groundedIntoDoublePlays),* FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); Assert.assertTrue(newSelections.get(0).isSetFunctionCall()); Assert.assertEquals(newSelections.get(0).getFunctionCall().getOperator(), "add"); @@ -230,7 +230,7 @@ public void testFuncOnColumns2() { public void testMultipleUnqualifiedStars() { String sql = "SELECT *,* FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); Assert.assertEquals(newSelections.get(0).getIdentifier().getName(), "G_old"); Assert.assertEquals(newSelections.get(1).getIdentifier().getName(), "groundedIntoDoublePlays"); @@ -250,7 +250,7 @@ public void testAll() { "SELECT abs(homeRuns),sqrt(groundedIntoDoublePlays),*,$segmentName,$hostName,playerStint as pstint,playerID " + "FROM baseballStats"; PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql); - BaseBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); + BaseSingleStageBrokerRequestHandler.updateColumnNames("baseballStats", pinotQuery, false, COL_MAP); List newSelections = pinotQuery.getSelectList(); Assert.assertTrue(newSelections.get(0).isSetFunctionCall()); Assert.assertEquals(newSelections.get(0).getFunctionCall().getOperator(), "abs"); diff --git a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotDriver.java b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotDriver.java index 3597b69c0332..93ba1ab551d1 100644 --- a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotDriver.java +++ b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotDriver.java @@ -87,7 +87,7 @@ public Connection connect(String url, Properties info) if (!this.acceptsURL(url)) { return null; } - LOGGER.info("Initiating connection to database for url: " + url); + LOGGER.info("Initiating connection to database for url: {}", url); Map urlParams = DriverUtils.getURLParams(url); info.putAll(urlParams); diff --git a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java index ac52810af93c..a55c4d1a3299 100644 --- a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java +++ b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java @@ -29,7 +29,6 @@ import java.util.List; import java.util.Map; import java.util.Properties; -import java.util.regex.Matcher; import java.util.regex.Pattern; import javax.net.ssl.SSLContext; import org.apache.commons.configuration2.MapConfiguration; @@ -49,6 +48,8 @@ public class DriverUtils { public static final String DRIVER = "pinot"; public static final Logger LOG = LoggerFactory.getLogger(DriverUtils.class); private static final String LIMIT_STATEMENT_REGEX = "\\s(limit)\\s"; + private static final Pattern LIMIT_STATEMENT_REGEX_PATTERN = + Pattern.compile(LIMIT_STATEMENT_REGEX, Pattern.CASE_INSENSITIVE); // SSL Properties public static final String PINOT_JDBC_TLS_PREFIX = "pinot.jdbc.tls"; @@ -62,8 +63,8 @@ private DriverUtils() { } public static SSLContext getSSLContextFromJDBCProps(Properties properties) { - TlsConfig tlsConfig = TlsUtils.extractTlsConfig( - new PinotConfiguration(new MapConfiguration(properties)), PINOT_JDBC_TLS_PREFIX); + TlsConfig tlsConfig = + TlsUtils.extractTlsConfig(new PinotConfiguration(new MapConfiguration(properties)), PINOT_JDBC_TLS_PREFIX); TlsUtils.installDefaultSSLSocketFactory(tlsConfig); return TlsUtils.getSslContext(); } @@ -213,9 +214,7 @@ public static String getJavaClassName(String columnDataType) { } public static boolean queryContainsLimitStatement(String query) { - Pattern pattern = Pattern.compile(LIMIT_STATEMENT_REGEX, Pattern.CASE_INSENSITIVE); - Matcher matcher = pattern.matcher(query); - return matcher.find(); + return LIMIT_STATEMENT_REGEX_PATTERN.matcher(query).find(); } public static String enableQueryOptions(String sql, Map options) { diff --git a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotConnectionTest.java b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotConnectionTest.java index ac3a418ccd10..5fd2b3ab6240 100644 --- a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotConnectionTest.java +++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotConnectionTest.java @@ -18,14 +18,22 @@ */ package org.apache.pinot.client; +import java.sql.DatabaseMetaData; import java.sql.Statement; import org.testng.Assert; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; public class PinotConnectionTest { - private DummyPinotClientTransport _dummyPinotClientTransport = new DummyPinotClientTransport(); - private DummyPinotControllerTransport _dummyPinotControllerTransport = DummyPinotControllerTransport.create(); + private DummyPinotClientTransport _dummyPinotClientTransport; + private DummyPinotControllerTransport _dummyPinotControllerTransport; + + @BeforeMethod + public void beforeTestMethod() { + _dummyPinotClientTransport = new DummyPinotClientTransport(); + _dummyPinotControllerTransport = DummyPinotControllerTransport.create(); + } @Test public void createStatementTest() @@ -36,6 +44,31 @@ public void createStatementTest() Assert.assertNotNull(statement); } + @Test + public void getMetaDataTest() + throws Exception { + try (PinotConnection pinotConnection = + new PinotConnection("dummy", _dummyPinotClientTransport, "dummy", _dummyPinotControllerTransport)) { + DatabaseMetaData metaData = pinotConnection.getMetaData(); + Assert.assertNotNull(metaData); + Assert.assertTrue(metaData.getDatabaseMajorVersion() > 0); + Assert.assertTrue(metaData.getDatabaseMinorVersion() >= 0); + Assert.assertEquals(metaData.getDatabaseProductName(), "APACHE_PINOT"); + Assert.assertEquals(metaData.getDriverName(), "APACHE_PINOT_DRIVER"); + Assert.assertNotNull(metaData.getConnection()); + } + } + + @Test + public void isClosedTest() + throws Exception { + PinotConnection pinotConnection = + new PinotConnection("dummy", _dummyPinotClientTransport, "dummy", _dummyPinotControllerTransport); + Assert.assertFalse(pinotConnection.isClosed()); + pinotConnection.close(); + Assert.assertTrue(pinotConnection.isClosed()); + } + @Test public void setUserAgentTest() throws Exception { diff --git a/pinot-common/pom.xml b/pinot-common/pom.xml index 32cd2eb9dbce..0a3ad400ccca 100644 --- a/pinot-common/pom.xml +++ b/pinot-common/pom.xml @@ -177,6 +177,8 @@ org.apache.calcite calcite-babel + + org.glassfish.jersey.core jersey-server @@ -206,10 +208,14 @@ swagger-jersey2-jaxrs - org.testng - testng - test + org.webjars + swagger-ui + + jakarta.servlet + jakarta.servlet-api + + com.google.guava guava @@ -262,15 +268,6 @@ net.sf.jopt-simple jopt-simple - - nl.jqno.equalsverifier - equalsverifier - test - - - org.webjars - swagger-ui - com.google.code.findbugs jsr305 @@ -287,11 +284,6 @@ com.jayway.jsonpath json-path - - org.mockito - mockito-core - test - org.slf4j jcl-over-slf4j @@ -341,6 +333,27 @@ io.github.hakky54 sslcontext-kickstart-for-netty + + com.google.re2j + re2j + 1.7 + + + + org.testng + testng + test + + + org.mockito + mockito-core + test + + + nl.jqno.equalsverifier + equalsverifier + test + @@ -363,7 +376,7 @@ generate-sources generate-sources - + @@ -387,10 +400,11 @@ - + - - + + + run diff --git a/pinot-common/src/main/codegen/config.fmpp b/pinot-common/src/main/codegen/config.fmpp index 95c5a3d33fb2..aebbab93975e 100644 --- a/pinot-common/src/main/codegen/config.fmpp +++ b/pinot-common/src/main/codegen/config.fmpp @@ -17,7 +17,7 @@ # under the License. # -# Copied from Calcite 1.36.0 babel and modified for Pinot syntax. Update this file when upgrading Calcite version. +# Copied from Calcite 1.37.0 babel and modified for Pinot syntax. Update this file when upgrading Calcite version. data: { default: tdd("../default_config.fmpp") diff --git a/pinot-common/src/main/codegen/default_config.fmpp b/pinot-common/src/main/codegen/default_config.fmpp index 78191c0c11cc..31092dce926c 100644 --- a/pinot-common/src/main/codegen/default_config.fmpp +++ b/pinot-common/src/main/codegen/default_config.fmpp @@ -17,7 +17,7 @@ # under the License. # -# Copied from Calcite 1.36.0 and modified for Pinot syntax. Update this file when upgrading Calcite version. +# Copied from Calcite 1.37.0 and modified for Pinot syntax. Update this file when upgrading Calcite version. # Default data declarations for parsers. # Each of these may be overridden in a parser's config.fmpp file. diff --git a/pinot-common/src/main/codegen/templates/Parser.jj b/pinot-common/src/main/codegen/templates/Parser.jj index 1e86c8c39436..965a87b21db4 100644 --- a/pinot-common/src/main/codegen/templates/Parser.jj +++ b/pinot-common/src/main/codegen/templates/Parser.jj @@ -17,7 +17,7 @@ * under the License. */ -// Copied from Calcite 1.36.0 and modified for Pinot syntax. Update this file when upgrading Calcite version. +// Copied from Calcite 1.37.0 and modified for Pinot syntax. Update this file when upgrading Calcite version. // Modified parts are marked with "PINOT CUSTOMIZATION START/END". <@pp.dropOutputFile /> @@ -79,6 +79,7 @@ import org.apache.calcite.sql.SqlJsonQueryWrapperBehavior; import org.apache.calcite.sql.SqlJsonValueEmptyOrErrorBehavior; import org.apache.calcite.sql.SqlJsonValueReturning; import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlLambda; import org.apache.calcite.sql.SqlLiteral; import org.apache.calcite.sql.SqlMatchRecognize; import org.apache.calcite.sql.SqlMerge; @@ -1039,6 +1040,9 @@ void AddArg0(List list, ExprContext exprContext) : ) ( e = Default() + | + LOOKAHEAD((SimpleIdentifierOrList() | ) ) + e = LambdaExpression() | LOOKAHEAD(3) e = TableParam() @@ -1066,6 +1070,9 @@ void AddArg(List list, ExprContext exprContext) : ) ( e = Default() + | + LOOKAHEAD((SimpleIdentifierOrList() | ) ) + e = LambdaExpression() | e = Expression(exprContext) | @@ -2199,12 +2206,20 @@ SqlNode TableRef3(ExprContext exprContext, boolean lateral) : { ( LOOKAHEAD(2) - tableName = CompoundTableIdentifier() - ( tableRef = TableHints(tableName) | { tableRef = tableName; } ) - [ tableRef = ExtendTable(tableRef) ] - tableRef = Over(tableRef) - [ tableRef = Snapshot(tableRef) ] - [ tableRef = MatchRecognize(tableRef) ] + tableName = CompoundTableIdentifier() { s = span(); } + ( + // Table call syntax like FROM a.b() instead of FROM TABLE(a.b()) + // Three tokens needed to disambiguate EXTEND syntax from CALCITE-493. + // Example: "FROM EventLog(lastGCTime TIME)". + LOOKAHEAD(3) + tableRef = ImplicitTableFunctionCallArgs(tableName) + | + ( tableRef = TableHints(tableName) | { tableRef = tableName; } ) + [ tableRef = ExtendTable(tableRef) ] + tableRef = Over(tableRef) + [ tableRef = Snapshot(tableRef) ] + [ tableRef = MatchRecognize(tableRef) ] + ) | LOOKAHEAD(2) [ { lateral = true; } ] @@ -2394,6 +2409,38 @@ void AddCompoundIdentifierType(List list, List extendList) : } } +SqlNode ImplicitTableFunctionCallArgs(SqlIdentifier name) : +{ + final List tableFuncArgs = new ArrayList(); + final SqlNode call; + final Span s; +} +{ + // Table call syntax like FROM a.b() instead of FROM TABLE(a.b()) + // We've already parsed the name, so we don't use NamedRoutineCall. + { s = span(); } + + [ + AddArg0(tableFuncArgs, ExprContext.ACCEPT_CURSOR) + ( + { + // a comma-list can't appear where only a query is expected + checkNonQueryExpression(ExprContext.ACCEPT_CURSOR); + } + AddArg(tableFuncArgs, ExprContext.ACCEPT_CURSOR) + )* + ] + + { + final SqlParserPos pos = s.end(this); + call = createCall(name, pos, + SqlFunctionCategory.USER_DEFINED_TABLE_FUNCTION, null, + tableFuncArgs); + return SqlStdOperatorTable.COLLECTION_TABLE.createCall(pos, + call); + } +} + SqlNode TableFunctionCall() : { final Span s; @@ -3918,6 +3965,43 @@ SqlNode Expression3(ExprContext exprContext) : } } +/** + * Parses a lambda expression. + */ +SqlNode LambdaExpression() : +{ + final SqlNodeList parameters; + final SqlNode expression; + final Span s; +} +{ + parameters = SimpleIdentifierOrListOrEmpty() + { s = span(); } + expression = Expression(ExprContext.ACCEPT_NON_QUERY) + { + return new SqlLambda(s.end(this), parameters, expression); + } +} + +/** + * List of simple identifiers in parentheses or empty parentheses or one simple identifier. + *
    Examples: + *
  • {@code ()} + *
  • {@code DEPTNO} + *
  • {@code (EMPNO, DEPTNO)} + *
+ */ +SqlNodeList SimpleIdentifierOrListOrEmpty() : +{ + SqlNodeList list; +} +{ + LOOKAHEAD(2) + { return SqlNodeList.EMPTY; } +| + list = SimpleIdentifierOrList() { return list; } +} + SqlOperator periodOperator() : { } @@ -4709,6 +4793,7 @@ SqlLiteral DateTimeLiteral() : { final String p; final Span s; + boolean local = false; } { { @@ -4740,6 +4825,7 @@ SqlLiteral DateTimeLiteral() : return SqlLiteral.createUnknown("DATETIME", p, s.end(this)); } | + LOOKAHEAD(2)