enable stream load record to audit log system table

Author: sollhui

be/src/common/config.cpp

@@ -627,6 +627,8 @@ DEFINE_mInt32(slave_replica_writer_rpc_timeout_sec, "60");
 // Whether to enable stream load record function, the default is false.
 // False: disable stream load record
 DEFINE_mBool(enable_stream_load_record, "false");
+// Whether to enable stream load record to audit log table, the default is true.
+DEFINE_mBool(enable_stream_load_record_to_audit_log_table, "true");
 // batch size of stream load record reported to FE
 DEFINE_mInt32(stream_load_record_batch_size, "50");
 // expire time of stream load record in rocksdb.

be/src/common/config.h

@@ -648,6 +648,8 @@ DECLARE_mInt32(slave_replica_writer_rpc_timeout_sec);
 // Whether to enable stream load record function, the default is false.
 // False: disable stream load record
 DECLARE_mBool(enable_stream_load_record);
+// Whether to enable stream load record to audit log table, the default is true.
+DECLARE_mBool(enable_stream_load_record_to_audit_log_table);
 // batch size of stream load record reported to FE
 DECLARE_mInt32(stream_load_record_batch_size);
 // expire time of stream load record in rocksdb.

be/src/http/action/http_stream.cpp

@@ -113,9 +113,11 @@ void HttpStreamAction::handle(HttpRequest* req) {
     // add new line at end
     str = str + '\n';
     HttpChannel::send_reply(req, str);
-    if (config::enable_stream_load_record) {
-        str = ctx->prepare_stream_load_record(str);
-        _save_stream_load_record(ctx, str);
+    if (config::enable_stream_load_record || config::enable_stream_load_record_to_audit_log_table) {
+        if (req->header(HTTP_SKIP_RECORD).empty()) {
+            str = ctx->prepare_stream_load_record(str);
+            _save_stream_load_record(ctx, str);
+        }
     }
     // update statistics
     http_stream_requests_total->increment(1);
@@ -176,9 +178,12 @@ int HttpStreamAction::on_header(HttpRequest* req) {
         // add new line at end
         str = str + '\n';
         HttpChannel::send_reply(req, str);
-        if (config::enable_stream_load_record) {
-            str = ctx->prepare_stream_load_record(str);
-            _save_stream_load_record(ctx, str);
+        if (config::enable_stream_load_record ||
+            config::enable_stream_load_record_to_audit_log_table) {
+            if (req->header(HTTP_SKIP_RECORD).empty()) {
+                str = ctx->prepare_stream_load_record(str);
+                _save_stream_load_record(ctx, str);
+            }
         }
         return -1;
     }

be/src/http/action/stream_load.cpp

@@ -131,9 +131,11 @@ void StreamLoadAction::handle(HttpRequest* req) {
     str = str + '\n';
     HttpChannel::send_reply(req, str);
 #ifndef BE_TEST
-    if (config::enable_stream_load_record) {
-        str = ctx->prepare_stream_load_record(str);
-        _save_stream_load_record(ctx, str);
+    if (config::enable_stream_load_record || config::enable_stream_load_record_to_audit_log_table) {
+        if (req->header(HTTP_SKIP_RECORD).empty()) {
+            str = ctx->prepare_stream_load_record(str);
+            _save_stream_load_record(ctx, str);
+        }
     }
 #endif
 
@@ -238,9 +240,12 @@ int StreamLoadAction::on_header(HttpRequest* req) {
         str = str + '\n';
         HttpChannel::send_reply(req, str);
 #ifndef BE_TEST
-        if (config::enable_stream_load_record) {
-            str = ctx->prepare_stream_load_record(str);
-            _save_stream_load_record(ctx, str);
+        if (config::enable_stream_load_record ||
+            config::enable_stream_load_record_to_audit_log_table) {
+            if (req->header(HTTP_SKIP_RECORD).empty()) {
+                str = ctx->prepare_stream_load_record(str);
+                _save_stream_load_record(ctx, str);
+            }
         }
 #endif
         return -1;

be/src/http/http_client.cpp

@@ -409,6 +409,12 @@ Status HttpClient::execute_post_request(const std::string& payload, std::string*
     return execute(response);
 }
 
+Status HttpClient::execute_put_request(const std::string& payload, std::string* response) {
+    set_payload(payload);
+    curl_easy_setopt(_curl, CURLOPT_CUSTOMREQUEST, "PUT");
+    return execute(response);
+}
+
 Status HttpClient::execute_delete_request(const std::string& payload, std::string* response) {
     set_method(DELETE);
     set_payload(payload);

be/src/http/http_client.h

@@ -158,6 +158,8 @@ class HttpClient {
 
     Status execute_post_request(const std::string& payload, std::string* response);
 
+    Status execute_put_request(const std::string& payload, std::string* response);
+
     Status execute_delete_request(const std::string& payload, std::string* response);
 
     // execute a simple method, and its response is saved in response argument

be/src/http/http_common.h

@@ -73,5 +73,6 @@ static const std::string HTTP_GROUP_COMMIT = "group_commit";
 static const std::string HTTP_CLOUD_CLUSTER = "cloud_cluster";
 static const std::string HTTP_EMPTY_FIELD_AS_NULL = "empty_field_as_null";
 static const std::string HTTP_COMPUTE_GROUP = "compute_group";
+static const std::string HTTP_SKIP_RECORD = "skip_record";
 
 } // namespace doris

be/src/runtime/exec_env.h

@@ -30,6 +30,7 @@
 
 #include "common/config.h"
 #include "common/status.h"
+#include "exec/schema_scanner/schema_routine_load_job_scanner.h"
 #include "io/cache/fs_file_cache_storage.h"
 #include "olap/memtable_memory_limiter.h"
 #include "olap/options.h"
@@ -105,6 +106,7 @@ class LoadStreamMgr;
 class LoadStreamMapPool;
 class StreamLoadExecutor;
 class RoutineLoadTaskExecutor;
+class StreamLoadRecorderManager;
 class SmallFileMgr;
 class BackendServiceClient;
 class TPaloBrokerServiceClient;
@@ -484,6 +486,7 @@ class ExecEnv {
 
     std::unique_ptr<StreamLoadExecutor> _stream_load_executor;
     RoutineLoadTaskExecutor* _routine_load_task_executor = nullptr;
+    StreamLoadRecorderManager* _stream_load_recorder_manager = nullptr;
     SmallFileMgr* _small_file_mgr = nullptr;
     HeartbeatFlags* _heartbeat_flags = nullptr;
     vectorized::ScannerScheduler* _scanner_scheduler = nullptr;

be/src/runtime/exec_env_init.cpp

@@ -93,6 +93,7 @@
 #include "runtime/small_file_mgr.h"
 #include "runtime/stream_load/new_load_stream_mgr.h"
 #include "runtime/stream_load/stream_load_executor.h"
+#include "runtime/stream_load/stream_load_recorder_manager.h"
 #include "runtime/thread_context.h"
 #include "runtime/user_function_cache.h"
 #include "runtime/workload_group/workload_group_manager.h"
@@ -392,6 +393,10 @@ Status ExecEnv::_init(const std::vector<StorePath>& store_paths,
         return st;
     }
 
+    // should start after storage_engine->open()
+    _stream_load_recorder_manager = new StreamLoadRecorderManager();
+    _stream_load_recorder_manager->start();
+
     // create internal workload group should be after storage_engin->open()
     RETURN_IF_ERROR(_create_internal_workload_group());
     _workload_sched_mgr = new WorkloadSchedPolicyMgr();
@@ -765,6 +770,7 @@ void ExecEnv::destroy() {
     SAFE_STOP(_group_commit_mgr);
     // _routine_load_task_executor should be stopped before _new_load_stream_mgr.
     SAFE_STOP(_routine_load_task_executor);
+    SAFE_STOP(_stream_load_recorder_manager);
     // stop workload scheduler
     SAFE_STOP(_workload_sched_mgr);
     // stop pipline step 2, cgroup execution
@@ -831,6 +837,7 @@ void ExecEnv::destroy() {
     SAFE_DELETE(_file_meta_cache);
     SAFE_DELETE(_group_commit_mgr);
     SAFE_DELETE(_routine_load_task_executor);
+    SAFE_DELETE(_stream_load_recorder_manager);
     // _stream_load_executor
     SAFE_DELETE(_function_client_cache);
     SAFE_DELETE(_streaming_client_cache);

be/src/runtime/stream_load/stream_load_recorder.cpp

@@ -100,6 +100,21 @@ Status StreamLoadRecorder::put(const std::string& key, const std::string& value)
     return Status::OK();
 }
 
+Status StreamLoadRecorder::get(const std::string& key, std::string* value) {
+    rocksdb::ColumnFamilyHandle* handle = _handles[0];
+    rocksdb::ReadOptions read_options;
+    rocksdb::Status s = _db->Get(read_options, handle, rocksdb::Slice(key), value);
+    if (s.IsNotFound()) {
+        return Status::NotFound("Key not found: {}", key);
+    }
+    if (!s.ok()) {
+        LOG(WARNING) << "rocks db get key:" << key << " failed, reason:" << s.ToString();
+        return Status::InternalError("Stream load record rocksdb get failed, reason: {}",
+                                     s.ToString());
+    }
+    return Status::OK();
+}
+
 Status StreamLoadRecorder::get_batch(const std::string& start, int batch_size,
                                      std::map<std::string, std::string>* stream_load_records) {
     rocksdb::ColumnFamilyHandle* handle = _handles[0];