diff --git a/cpp/src/arrow/flight/sql/client.cc b/cpp/src/arrow/flight/sql/client.cc index 6fe4a4056675..a93e038a0c11 100644 --- a/cpp/src/arrow/flight/sql/client.cc +++ b/cpp/src/arrow/flight/sql/client.cc @@ -128,12 +128,14 @@ FlightSqlClient::FlightSqlClient(std::shared_ptr client) PreparedStatement::PreparedStatement(FlightSqlClient* client, std::string handle, std::shared_ptr dataset_schema, - std::shared_ptr parameter_schema) + std::shared_ptr parameter_schema, + bool is_update) : client_(client), handle_(std::move(handle)), dataset_schema_(std::move(dataset_schema)), parameter_schema_(std::move(parameter_schema)), - is_closed_(false) {} + is_closed_(false), + is_update_(is_update) {} PreparedStatement::~PreparedStatement() { if (IsClosed()) return; @@ -632,7 +634,8 @@ arrow::Result> PreparedStatement::ParseRespon auto handle = prepared_statement_result.prepared_statement_handle(); return std::make_shared(client, handle, dataset_schema, - parameter_schema); + parameter_schema, + prepared_statement_result.is_update()); } arrow::Result> PreparedStatement::Execute( @@ -700,6 +703,8 @@ Status PreparedStatement::SetParameters( bool PreparedStatement::IsClosed() const { return is_closed_; } +bool PreparedStatement::is_update() const { return is_update_; } + const std::shared_ptr& PreparedStatement::dataset_schema() const { return dataset_schema_; } diff --git a/cpp/src/arrow/flight/sql/client.h b/cpp/src/arrow/flight/sql/client.h index 9541432114f9..2f4c1b5c526d 100644 --- a/cpp/src/arrow/flight/sql/client.h +++ b/cpp/src/arrow/flight/sql/client.h @@ -440,9 +440,11 @@ class ARROW_FLIGHT_SQL_EXPORT PreparedStatement { /// \param[in] handle Handle for this prepared statement. /// \param[in] dataset_schema Schema of the resulting dataset. /// \param[in] parameter_schema Schema of the parameters (if any). + /// \param[in] is_update Whether this is an update query. PreparedStatement(FlightSqlClient* client, std::string handle, std::shared_ptr dataset_schema, - std::shared_ptr parameter_schema); + std::shared_ptr parameter_schema, + bool is_update = false); /// \brief Default destructor for the PreparedStatement class. /// The destructor will call the Close method from the class in order, @@ -472,6 +474,9 @@ class ARROW_FLIGHT_SQL_EXPORT PreparedStatement { /// \return The ResultSet schema from the query. const std::shared_ptr& dataset_schema() const; + /// \brief Check if the prepared statement represents an update query. + bool is_update() const; + /// \brief Set a RecordBatch that contains the parameters that will be bound. Status SetParameters(std::shared_ptr parameter_binding); @@ -499,6 +504,7 @@ class ARROW_FLIGHT_SQL_EXPORT PreparedStatement { std::shared_ptr parameter_schema_; std::shared_ptr parameter_binding_; bool is_closed_; + bool is_update_; }; /// \brief A handle for a server-side savepoint. diff --git a/cpp/src/arrow/flight/sql/example/sqlite_server.cc b/cpp/src/arrow/flight/sql/example/sqlite_server.cc index 94fef95afc9e..dd4f758cf577 100644 --- a/cpp/src/arrow/flight/sql/example/sqlite_server.cc +++ b/cpp/src/arrow/flight/sql/example/sqlite_server.cc @@ -479,8 +479,9 @@ class SQLiteFlightSqlServer::Impl { } std::shared_ptr parameter_schema = arrow::schema(parameter_fields); + const bool is_update = sqlite3_stmt_readonly(stmt) == 0; return ActionCreatePreparedStatementResult{ - std::move(dataset_schema), std::move(parameter_schema), std::move(handle)}; + std::move(dataset_schema), std::move(parameter_schema), std::move(handle), is_update}; } Status ClosePreparedStatement(const ServerCallContext& context, diff --git a/cpp/src/arrow/flight/sql/server.cc b/cpp/src/arrow/flight/sql/server.cc index 10dd073b6227..291f818d36db 100644 --- a/cpp/src/arrow/flight/sql/server.cc +++ b/cpp/src/arrow/flight/sql/server.cc @@ -498,6 +498,7 @@ ARROW_UNSUPPRESS_DEPRECATION_WARNING arrow::Result PackActionResult(ActionCreatePreparedStatementResult result) { pb::sql::ActionCreatePreparedStatementResult pb_result; pb_result.set_prepared_statement_handle(std::move(result.prepared_statement_handle)); + pb_result.set_is_update(result.is_update); if (result.dataset_schema != nullptr) { ARROW_ASSIGN_OR_RAISE(std::shared_ptr serialized, ipc::SerializeSchema(*result.dataset_schema)); diff --git a/cpp/src/arrow/flight/sql/server.h b/cpp/src/arrow/flight/sql/server.h index 7130e96987b8..95f4d6d54f66 100644 --- a/cpp/src/arrow/flight/sql/server.h +++ b/cpp/src/arrow/flight/sql/server.h @@ -241,6 +241,8 @@ struct ARROW_FLIGHT_SQL_EXPORT ActionCreatePreparedStatementResult { std::shared_ptr parameter_schema; /// \brief The server-generated opaque identifier for the statement. std::string prepared_statement_handle; + /// \brief Whether this prepared statement represents an update query. + bool is_update = false; }; /// @} diff --git a/cpp/src/arrow/flight/sql/server_test.cc b/cpp/src/arrow/flight/sql/server_test.cc index 42e0c7d7a1f1..9725d1678293 100644 --- a/cpp/src/arrow/flight/sql/server_test.cc +++ b/cpp/src/arrow/flight/sql/server_test.cc @@ -448,6 +448,7 @@ TEST_F(TestFlightSqlServer, TestCommandStatementUpdate) { TEST_F(TestFlightSqlServer, TestCommandPreparedStatementQuery) { ASSERT_OK_AND_ASSIGN(auto prepared_statement, sql_client->Prepare({}, "SELECT * FROM intTable")); + ASSERT_FALSE(prepared_statement->is_update()); ASSERT_OK_AND_ASSIGN(auto flight_info, prepared_statement->Execute()); @@ -587,6 +588,7 @@ TEST_F(TestFlightSqlServer, TestCommandPreparedStatementUpdate) { auto prepared_statement, sql_client->Prepare( {}, "INSERT INTO INTTABLE (keyName, value) VALUES ('new_value', 999)")); + ASSERT_TRUE(prepared_statement->is_update()); ASSERT_OK_AND_EQ(5, ExecuteCountQuery("SELECT COUNT(*) FROM intTable")); ASSERT_OK_AND_EQ(1, prepared_statement->ExecuteUpdate());