From b02ed4e6ec0fe97c4bc04186d663560ffc113648 Mon Sep 17 00:00:00 2001
From: Kjell Winblad <kjellwinblad@gmail.com>
Date: Tue, 23 Apr 2024 11:09:08 +0200
Subject: [PATCH] feat: add stop after render and after render trace to pgsql
 action

---
 apps/emqx/src/emqx_trace/emqx_trace.erl      |  7 +++++--
 apps/emqx_postgresql/src/emqx_postgresql.erl | 14 +++++++++++---
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/apps/emqx/src/emqx_trace/emqx_trace.erl b/apps/emqx/src/emqx_trace/emqx_trace.erl
index 7bbe59b2b..4152fdbaa 100644
--- a/apps/emqx/src/emqx_trace/emqx_trace.erl
+++ b/apps/emqx/src/emqx_trace/emqx_trace.erl
@@ -87,7 +87,7 @@ unsubscribe(<<"$SYS/", _/binary>>, _SubOpts) ->
 unsubscribe(Topic, SubOpts) ->
     ?TRACE("UNSUBSCRIBE", "unsubscribe", #{topic => Topic, sub_opts => SubOpts}).
 
-rendered_action_template(ActionID, RenderResult) ->
+rendered_action_template(ActionID, RenderResult) when is_binary(ActionID) ->
     TraceResult = ?TRACE(
         "QUERY_RENDER",
         "action_template_rendered",
@@ -111,7 +111,10 @@ rendered_action_template(ActionID, RenderResult) ->
         _ ->
             ok
     end,
-    TraceResult.
+    TraceResult;
+rendered_action_template(_ActionID, _RenderResult) ->
+    %% We do nothing if we don't get a valid Action ID
+    ok.
 
 log(List, Msg, Meta) ->
     log(debug, List, Msg, Meta).
diff --git a/apps/emqx_postgresql/src/emqx_postgresql.erl b/apps/emqx_postgresql/src/emqx_postgresql.erl
index f27ec8615..761c9c0f6 100644
--- a/apps/emqx_postgresql/src/emqx_postgresql.erl
+++ b/apps/emqx_postgresql/src/emqx_postgresql.erl
@@ -304,7 +304,7 @@ on_query(
     }),
     Type = pgsql_query_type(TypeOrKey),
     {NameOrSQL2, Data} = proc_sql_params(TypeOrKey, NameOrSQL, Params, State),
-    Res = on_sql_query(InstId, PoolName, Type, NameOrSQL2, Data),
+    Res = on_sql_query(TypeOrKey, InstId, PoolName, Type, NameOrSQL2, Data),
     ?tp(postgres_bridge_connector_on_query_return, #{instance_id => InstId, result => Res}),
     handle_result(Res).
 
@@ -337,7 +337,7 @@ on_batch_query(
         {_Statement, RowTemplate} ->
             PrepStatement = get_prepared_statement(BinKey, State),
             Rows = [render_prepare_sql_row(RowTemplate, Data) || {_Key, Data} <- BatchReq],
-            case on_sql_query(InstId, PoolName, execute_batch, PrepStatement, Rows) of
+            case on_sql_query(Key, InstId, PoolName, execute_batch, PrepStatement, Rows) of
                 {error, _Error} = Result ->
                     handle_result(Result);
                 {_Column, Results} ->
@@ -386,7 +386,15 @@ get_prepared_statement(Key, #{prepares := PrepStatements}) ->
     BinKey = to_bin(Key),
     maps:get(BinKey, PrepStatements).
 
-on_sql_query(InstId, PoolName, Type, NameOrSQL, Data) ->
+on_sql_query(Key, InstId, PoolName, Type, NameOrSQL, Data) ->
+    emqx_trace:rendered_action_template(
+        Key,
+        #{
+            statement_type => Type,
+            statement_or_name => NameOrSQL,
+            data => Data
+        }
+    ),
     try ecpool:pick_and_do(PoolName, {?MODULE, Type, [NameOrSQL, Data]}, no_handover) of
         {error, Reason} ->
             ?tp(