| 153 | } |
| 154 | |
| 155 | Status RunMain() { |
| 156 | #ifdef ARROW_WITH_OPENTELEMETRY |
| 157 | ARROW_ASSIGN_OR_RAISE(auto otel_scope, OtelScope::Make()); |
| 158 | #endif |
| 159 | |
| 160 | ARROW_ASSIGN_OR_RAISE(auto location, Location::ForGrpcTcp(FLAGS_host, FLAGS_port)); |
| 161 | auto client_options = FlightClientOptions::Defaults(); |
| 162 | client_options.middleware.push_back( |
| 163 | arrow::flight::MakeTracingClientMiddlewareFactory()); |
| 164 | ARROW_ASSIGN_OR_RAISE(auto client, FlightClient::Connect(location, client_options)); |
| 165 | |
| 166 | FlightCallOptions call_options; |
| 167 | |
| 168 | if (!FLAGS_username.empty() || !FLAGS_password.empty()) { |
| 169 | Result<std::pair<std::string, std::string>> bearer_result = |
| 170 | client->AuthenticateBasicToken({}, FLAGS_username, FLAGS_password); |
| 171 | ARROW_RETURN_NOT_OK(bearer_result); |
| 172 | |
| 173 | call_options.headers.push_back(bearer_result.ValueOrDie()); |
| 174 | } |
| 175 | |
| 176 | FlightSqlClient sql_client(std::move(client)); |
| 177 | |
| 178 | if (FLAGS_command == "ExecuteUpdate") { |
| 179 | ARROW_ASSIGN_OR_RAISE(auto rows, sql_client.ExecuteUpdate(call_options, FLAGS_query)); |
| 180 | |
| 181 | std::cout << "Result: " << rows << std::endl; |
| 182 | |
| 183 | return Status::OK(); |
| 184 | } |
| 185 | |
| 186 | std::unique_ptr<FlightInfo> info; |
| 187 | |
| 188 | if (FLAGS_command == "Execute") { |
| 189 | ARROW_ASSIGN_OR_RAISE(info, sql_client.Execute(call_options, FLAGS_query)); |
| 190 | } else if (FLAGS_command == "GetCatalogs") { |
| 191 | ARROW_ASSIGN_OR_RAISE(info, sql_client.GetCatalogs(call_options)); |
| 192 | } else if (FLAGS_command == "PreparedStatementExecute") { |
| 193 | ARROW_ASSIGN_OR_RAISE(auto prepared_statement, |
| 194 | sql_client.Prepare(call_options, FLAGS_query)); |
| 195 | ARROW_ASSIGN_OR_RAISE(info, prepared_statement->Execute()); |
| 196 | } else if (FLAGS_command == "PreparedStatementExecuteParameterBinding") { |
| 197 | ARROW_ASSIGN_OR_RAISE(auto prepared_statement, sql_client.Prepare({}, FLAGS_query)); |
| 198 | auto parameter_schema = prepared_statement->parameter_schema(); |
| 199 | auto result_set_schema = prepared_statement->dataset_schema(); |
| 200 | |
| 201 | std::cout << result_set_schema->ToString(false) << std::endl; |
| 202 | arrow::Int64Builder int_builder; |
| 203 | ARROW_RETURN_NOT_OK(int_builder.Append(1)); |
| 204 | std::shared_ptr<arrow::Array> int_array; |
| 205 | ARROW_RETURN_NOT_OK(int_builder.Finish(&int_array)); |
| 206 | std::shared_ptr<arrow::RecordBatch> result; |
| 207 | result = arrow::RecordBatch::Make(parameter_schema, 1, {int_array}); |
| 208 | |
| 209 | ARROW_RETURN_NOT_OK(prepared_statement->SetParameters(result)); |
| 210 | ARROW_ASSIGN_OR_RAISE(info, prepared_statement->Execute()); |
| 211 | } else if (FLAGS_command == "GetDbSchemas") { |
| 212 | ARROW_ASSIGN_OR_RAISE( |
no test coverage detected