From 1c71a0b85978d6cd6ccf584739d3a340bd325b28 Mon Sep 17 00:00:00 2001
From: satish <satish@traceable.ai>
Date: Wed, 29 Nov 2023 11:56:56 +0530
Subject: [PATCH] draft APIs schema for async query execution

---
 .../src/main/proto/query-service.proto        | 10 +++++
 .../src/main/proto/request.proto              | 39 ++++++++++++++++
 .../src/main/proto/response.proto             | 45 +++++++++++++++++++
 3 files changed, 94 insertions(+)

diff --git a/query-service-api/src/main/proto/query-service.proto b/query-service-api/src/main/proto/query-service.proto
index 1990a20d..7762581c 100644
--- a/query-service-api/src/main/proto/query-service.proto
+++ b/query-service-api/src/main/proto/query-service.proto
@@ -13,4 +13,14 @@ import "response.proto";
 service QueryService {
   rpc execute (QueryRequest) returns (stream ResultSetChunk) {
   }
+
+  rpc submitAsyncQuery (SubmitAsyncQueryRequest) returns (SubmitAsyncQueryResponse) {}
+
+  rpc submitAsyncGroupQuery (SubmitAsyncQueryGroupRequest) returns (SubmitAsyncQueryGroupResponse) {}
+
+  rpc getAsyncQueries (GetAsyncQueriesRequest) returns (GetAsyncQueriesResponse) {}
+
+  rpc getAsyncQueryResult (GetAsyncQueryResultRequest) returns (GetAsyncQueryResultResponse) {}
+
+  rpc getAsyncQueryGroupResult (GetAsyncQueryGroupResultRequest) returns (GetAsyncQueryGroupResultResponse) {}
 }
diff --git a/query-service-api/src/main/proto/request.proto b/query-service-api/src/main/proto/request.proto
index 5e4962f7..d852dbde 100644
--- a/query-service-api/src/main/proto/request.proto
+++ b/query-service-api/src/main/proto/request.proto
@@ -17,6 +17,45 @@ message Expression {
     AttributeExpression attribute_expression = 5;
   }
 }
+
+message SubmitAsyncQueryRequest {
+  // TODO: verify if all fields of QueryRequest are applicable
+  QueryRequest query_request = 1;
+}
+
+message SubmitAsyncQueryGroupRequest {
+  // TODO: verify if all fields of QueryRequest are applicable
+  QueryRequest query_request = 1;
+  // If multiple async queries are a part of same group, UI can fire each async
+  // query separately. In this case, each request needs to contain group details
+  AsyncQueryGroupDetails group = 2;
+}
+
+message AsyncQueryGroupDetails {
+  // UUID created by UI
+  string group_id = 1;
+  // unique name for each query within group
+  string query_name = 2;
+}
+
+message GetAsyncQueriesRequest {
+  // Note that these filters are different than filters within QueryRequest.
+  // These are filters for async queries listing
+  // TODO: do we need to define a different filter message?
+  Filter filter = 1;
+}
+
+message GetAsyncQueryResultRequest {
+  string async_query_id = 1;
+}
+
+message GetAsyncQueryGroupResultRequest {
+  string async_query_id = 1;
+  // In case of group query, UI can ask for results of multiple queries that
+  // belong to the group (referenced by query_name within group)
+  repeated string query_names = 2;
+}
+
 message QueryRequest {
 
   repeated string source = 1;
diff --git a/query-service-api/src/main/proto/response.proto b/query-service-api/src/main/proto/response.proto
index 9e3e684e..e1cd96c4 100644
--- a/query-service-api/src/main/proto/response.proto
+++ b/query-service-api/src/main/proto/response.proto
@@ -7,7 +7,52 @@ option java_outer_classname = "QueryResponseProto";
 package org.hypertrace.core.query.service;
 
 import "value.proto";
+import "google/protobuf/timestamp.proto";
 
+message SubmitAsyncQueryResponse {
+  AsyncQueryInfo query_info = 1;
+}
+
+message SubmitAsyncQueryGroupResponse {
+  // TODO: Reusing AsyncQueryInfo in case of query belonging to a group as well.
+  // The async_query_id will be at group level in this case.
+  AsyncQueryInfo query_info = 1;
+}
+
+message GetAsyncQueriesResponse {
+  repeated AsyncQueryInfo queries_info = 1;
+}
+
+message GetAsyncQueryResultResponse {
+  AsyncQueryResult query_result = 1;
+}
+
+message GetAsyncQueryGroupResultResponse {
+  // map of query_name and its result in case of group
+  map<string, AsyncQueryResult> group_result = 1;
+}
+
+message AsyncQueryResult {
+  // map of query_name and its result
+  ResultSetMetadata result_set_metadata = 1;
+  repeated Row row = 2;
+}
+
+message AsyncQueryInfo {
+  string async_query_id = 1;
+  AsyncQueryState state = 2;
+  google.protobuf.Timestamp created_timestamp = 3;
+  string author = 4;
+}
+
+enum AsyncQueryState {
+  ASYNC_QUERY_STATE_UNSPECIFIED = 0;
+
+  ASYNC_QUERY_STATE_IN_PROGRESS = 1;
+  ASYNC_QUERY_STATE_SUCCEEDED = 2;
+  ASYNC_QUERY_STATE_FAILED = 3;
+  // TODO: how to power partially completed status?
+}
 message ColumnMetadata {
   string column_name = 1;
   ValueType value_type = 2;