Introduce CancelQuery override in DuckDBFlightSqlServer to release gl…

…obal mutex upon query cancellation.
voltrondata · Jul 17, 2023 · bf3cc0c · bf3cc0c
1 parent a653e7a
commit bf3cc0c
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/src/duckdb/duckdb_server.cpp b/src/duckdb/duckdb_server.cpp
@@ -241,18 +241,19 @@ namespace arrow {
 
                     Status ClosePreparedStatement(const ServerCallContext &context,
                                                   const ActionClosePreparedStatementRequest &request) {
+                        // Unlock our mutex so that other threads can run queries...
+                        global_mutex.unlock();
+
                         const std::string &prepared_statement_handle = request.prepared_statement_handle;
 
                         auto search = prepared_statements_.find(prepared_statement_handle);
+
                         if (search != prepared_statements_.end()) {
                             prepared_statements_.erase(prepared_statement_handle);
                         } else {
                             return Status::Invalid("Prepared statement not found");
                         }
 
-                        // Unlock our mutex so that other threads can run queries...
-                        global_mutex.unlock();
-
                         return Status::OK();
                     }
 
@@ -354,6 +355,14 @@ namespace arrow {
                         return std::unique_ptr<FlightInfo>(new FlightInfo(result));
                     }
 
+                    arrow::Result<CancelResult> CancelQuery(
+                            const ServerCallContext &context, const ActionCancelQueryRequest &request) {
+                        // Release the global mutex
+                        global_mutex.unlock();
+
+                        return Status::OK();
+                    }
+
                 };
 
                 DuckDBFlightSqlServer::DuckDBFlightSqlServer(std::shared_ptr<Impl> impl)
@@ -443,6 +452,10 @@ namespace arrow {
                     return impl_->DoPutPreparedStatementUpdate(context, command, reader);
                 }
 
+                arrow::Result<CancelResult> DuckDBFlightSqlServer::CancelQuery(
+                        const ServerCallContext &context, const ActionCancelQueryRequest &request) {
+                    return impl_->CancelQuery(context, request);
+                }
 
             }  // namespace sqlite
         }  // namespace sql

diff --git a/src/duckdb/duckdb_server.h b/src/duckdb/duckdb_server.h
@@ -82,6 +82,9 @@ class DuckDBFlightSqlServer : public FlightSqlServerBase {
   arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTables(
       const ServerCallContext& context, const GetTables& command) override;
 
+  arrow::Result<CancelResult> CancelQuery(
+      const ServerCallContext& context, const ActionCancelQueryRequest& request) override;
+
  private:
   class Impl;
   std::shared_ptr<Impl> impl_;