RFC: [FlightRPC][WIP] Substrait, transaction, cancellation for Flight…

… SQL

cpp/src/arrow/flight/sql/CMakeLists.txt

@@ -61,7 +61,8 @@ add_arrow_lib(arrow_flight_sql
               STATIC_LINK_LIBS
               arrow_flight_static
               PRIVATE_INCLUDES
-              "${Protobuf_INCLUDE_DIRS}")
+              "${Protobuf_INCLUDE_DIRS}"
+              "${CMAKE_CURRENT_BINARY_DIR}/../")
 
 if(ARROW_FLIGHT_TEST_LINKAGE STREQUAL "static" AND ARROW_BUILD_STATIC)
   set(ARROW_FLIGHT_SQL_TEST_LINK_LIBS arrow_flight_sql_static)
@@ -90,6 +91,8 @@ if(ARROW_BUILD_TESTS OR ARROW_BUILD_EXAMPLES)
                  STATIC_LINK_LIBS
                  ${ARROW_FLIGHT_SQL_TEST_LINK_LIBS}
                  ${SQLite3_LIBRARIES}
+                 EXTRA_INCLUDES
+                 "${CMAKE_CURRENT_BINARY_DIR}/../"
                  LABELS
                  "arrow_flight_sql")
 

format/FlightSql.proto

@@ -18,6 +18,7 @@
 
 syntax = "proto3";
 import "google/protobuf/descriptor.proto";
+import "Flight.proto";
 
 option java_package = "org.apache.arrow.flight.sql.impl";
 package arrow.flight.protocol.sql;
@@ -89,6 +90,17 @@ enum SqlInfo {
    */
   FLIGHT_SQL_SERVER_READ_ONLY = 3;
 
+  /*
+   * Retrieves a boolean value indicating whether the Flight SQL Server supports executing Substrait plans.
+   */
+  FLIGHT_SQL_SUBSTRAIT = 4;
+
+  /*
+   * Retrieves an int32 indicating whether the Flight SQL Server supports explicit transaction RPCs.
+   *
+   * The possible values are listed in `FlightSqlTransactionSupport`.
+   */
+  FLIGHT_SQL_TRANSACTION = 5;
 
   // SQL Syntax Information [500-1000): provides information about SQL syntax supported by the Flight SQL Server.
 
@@ -760,6 +772,18 @@ enum SqlInfo {
   SQL_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED = 576;
 }
 
+// The level of support for Flight SQL transaction RPCs.
+enum FlightSqlTransactionSupport {
+  // Unknown/not indicated
+  FLIGHT_SQL_TRANSACTION_SUPPORT_UNKNOWN = 0;
+  // No supprot
+  FLIGHT_SQL_TRANSACTION_SUPPORT_NONE = 1;
+  // Transactions, but not savepoints
+  FLIGHT_SQL_TRANSACTION_SUPPORT_TRANSACTION = 2;
+  // Transactions and savepoints
+  FLIGHT_SQL_TRANSACTION_SUPPORT_SAVEPOINT = 3;
+}
+
 enum SqlSupportedCaseSensitivity {
   SQL_CASE_SENSITIVITY_UNKNOWN = 0;
   SQL_CASE_SENSITIVITY_CASE_INSENSITIVE = 1;
@@ -1405,7 +1429,7 @@ message CommandGetCrossReference {
   string fk_table = 6;
 }
 
-// SQL Execution Action Messages
+// Query Execution Action Messages
 
 /*
  * Request message for the "CreatePreparedStatement" action on a Flight SQL enabled backend.
@@ -1418,7 +1442,17 @@ message ActionCreatePreparedStatementRequest {
 }
 
 /*
- * Wrap the result of a "GetPreparedStatement" action.
+ * Request message for the "CreatePreparedSubstraitPlan" action on a Flight SQL enabled backend.
+ */
+message ActionCreatePreparedSubstraitPlan {
+  option (experimental) = true;
+
+  // The serialized substrait.Plan to create a prepared statement for.
+  bytes plan = 1;
+}
+
+/*
+ * Wrap the result of a "CreatePreparedStatement" or "CreatePreparedSubstraitPlan" action.
  *
  * The resultant PreparedStatement can be closed either:
  * - Manually, through the "ClosePreparedStatement" action;
@@ -1450,8 +1484,48 @@ message ActionClosePreparedStatementRequest {
   bytes prepared_statement_handle = 1;
 }
 
+/*
+ * Request message for the "BeginTransaction" action.
+ * Begins a transaction or creates a savepoint within a transaction.
+ */
+message ActionBeginTransactionRequest {
+  // Create a savepoint within the identified transaction. Only supported if
+  // FLIGHT_SQL_TRANSACTION is FLIGHT_SQL_TRANSACTION_SUPPORT_SAVEPOINT.
+  bytes transaction_id = 1;
+  // Name for the savepoint (if applicable).
+  string name = 2;
+}
+
+/*
+ * The result of a "BeginTransaction" action.
+ *
+ * The transaction/savepoint can be manipulated with the "EndTransaction"
+ * action, or automatically via server timeout.
+ */
+message ActionBeginTransactionResult {
+  // Opaque handle for the transaction on the server.
+  bytes transaction_id = 1;
+}
+
+/*
+ * Request message for the "EndTransaction" action.
+ * Commit or rollback a transaction, or release/rollback a savepoint within a transaction.
+ */
+message ActionEndTransactionRequest {
+  enum EndTransaction {
+    END_TRANSACTION_UNSPECIFIED = 0;
+    // Commit the transaction or release the savepoint.
+    END_TRANSACTION_COMMIT = 1;
+    // Roll back the transaction or to a savepoint.
+    END_TRANSACTION_ROLLBACK = 2;
+  }
+  // Opaque handle for the transaction on the server.
+  bytes transaction_id = 1;
+  // Whether to commit/rollback/release the given transaction/savepoint.
+  EndTransaction action = 2;
+}
 
-// SQL Execution Messages.
+// Query Execution Messages.
 
 /*
  * Represents a SQL query. Used in the command member of FlightDescriptor
@@ -1475,6 +1549,35 @@ message CommandStatementQuery {
 
   // The SQL syntax.
   string query = 1;
+  // Include the query as part of this transaction (by default queries are auto-committed).
+  bytes transaction_id = 2;
+}
+
+/*
+ * Represents a Substrait plan. Used in the command member of FlightDescriptor
+ * for the following RPC calls:
+ *  - GetSchema: return the Arrow schema of the query.
+ *    Fields on this schema may contain the following metadata:
+ *    - ARROW:FLIGHT:SQL:CATALOG_NAME      - Table's catalog name
+ *    - ARROW:FLIGHT:SQL:DB_SCHEMA_NAME    - Database schema name
+ *    - ARROW:FLIGHT:SQL:TABLE_NAME        - Table name
+ *    - ARROW:FLIGHT:SQL:TYPE_NAME         - The data source-specific name for the data type of the column.
+ *    - ARROW:FLIGHT:SQL:PRECISION         - Column precision/size
+ *    - ARROW:FLIGHT:SQL:SCALE             - Column scale/decimal digits if applicable
+ *    - ARROW:FLIGHT:SQL:IS_AUTO_INCREMENT - "1" indicates if the column is auto incremented, "0" otherwise.
+ *    - ARROW:FLIGHT:SQL:IS_CASE_SENSITIVE - "1" indicates if the column is case sensitive, "0" otherwise.
+ *    - ARROW:FLIGHT:SQL:IS_READ_ONLY      - "1" indicates if the column is read only, "0" otherwise.
+ *    - ARROW:FLIGHT:SQL:IS_SEARCHABLE     - "1" indicates if the column is searchable via WHERE clause, "0" otherwise.
+ *  - GetFlightInfo: execute the query.
+ *  - DoPut: execute the query.
+ */
+message CommandStatementSubstraitPlan {
+  option (experimental) = true;
+
+  // A serialized substrait.Plan
+  bytes plan = 1;
+  // Include the query as part of this transaction (by default queries are auto-committed).
+  bytes transaction_id = 2;
 }
 
 /**
@@ -1511,6 +1614,8 @@ message CommandPreparedStatementQuery {
 
   // Opaque handle for the prepared statement on the server.
   bytes prepared_statement_handle = 1;
+  // Include the query as part of this transaction (by default queries are auto-committed).
+  bytes transaction_id = 2;
 }
 
 /*
@@ -1522,6 +1627,8 @@ message CommandStatementUpdate {
 
   // The SQL syntax.
   string query = 1;
+  // Include the query as part of this transaction (by default queries are auto-committed).
+  bytes transaction_id = 2;
 }
 
 /*
@@ -1534,6 +1641,8 @@ message CommandPreparedStatementUpdate {
 
   // Opaque handle for the prepared statement on the server.
   bytes prepared_statement_handle = 1;
+  // Include the query as part of this transaction (by default queries are auto-committed).
+  bytes transaction_id = 2;
 }
 
 /*
@@ -1549,6 +1658,17 @@ message DoPutUpdateResult {
   int64 record_count = 1;
 }
 
+/*
+ * Explicitly cancel a running query.
+ *
+ * For distributed queries (ones that use GetFlightInfo->DoGet), this lets a
+ * single client explicitly cancel work across all clients, given server support.
+ */
+message ActionCancelQuery {
+  // The result of the GetFlightInfo RPC that initated the query.
+  FlightInfo info = 1;
+}
+
 extend google.protobuf.MessageOptions {
   bool experimental = 1000;
 }