Add get_os_pid/1 function

epgsql · Mar 3, 2022 · 26ff983 · 26ff983
1 parent 0b798b0
commit 26ff983
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 4 deletions.
diff --git a/src/epgsql.erl b/src/epgsql.erl
@@ -12,6 +12,7 @@
          get_parameter/2,
          set_notice_receiver/2,
          get_cmd_status/1,
+         get_os_pid/1,
          squery/2,
          equery/2, equery/3, equery/4,
          prepared_query/3,
@@ -260,6 +261,13 @@ set_notice_receiver(C, PidOrName) ->
 get_cmd_status(C) ->
     epgsql_sock:get_cmd_status(C).
 
+%% @doc Returns the OS pid of PostgreSQL backend OS process that serves this connection.
+%%
+%% Similar to `SELECT pg_get_pid()', but does not need network roundtrips.
+-spec get_os_pid(connection()) -> integer().
+get_os_pid(C) ->
+    epgsql_sock:get_os_pid(C).
+
 -spec squery(connection(), sql_query()) -> epgsql_cmd_squery:response() | epgsql_sock:error().
 %% @doc runs simple `SqlQuery' via given `Connection'
 %% @see epgsql_cmd_squery

diff --git a/src/epgsql_sock.erl b/src/epgsql_sock.erl
@@ -52,7 +52,8 @@
          get_cmd_status/1,
          cancel/1,
          copy_send_rows/3,
-         standby_status_update/3]).
+         standby_status_update/3,
+         get_os_pid/1]).
 
 -export([handle_call/3, handle_cast/2, handle_info/2, format_status/2]).
 -export([init/1, code_change/3, terminate/2]).
@@ -151,6 +152,9 @@ copy_send_rows(C, Rows, Timeout) ->
 standby_status_update(C, FlushedLSN, AppliedLSN) ->
     gen_server:call(C, {standby_status_update, FlushedLSN, AppliedLSN}).
 
+-spec get_os_pid(epgsql:connection()) -> integer().
+get_os_pid(C) ->
+    gen_server:call(C, get_cmd_status, get_os_pid).
 
 %% -- command APIs --
 
@@ -230,6 +234,9 @@ handle_call({set_async_receiver, PidOrName}, _From, #state{async = Previous} = S
 handle_call(get_cmd_status, _From, #state{complete_status = Status} = State) ->
     {reply, {ok, Status}, State};
 
+handle_call(get_os_pid, _From, #state{backend = {Pid, _Key}} = State) ->
+    {reply, Pid, State};
+
 handle_call({standby_status_update, FlushedLSN, AppliedLSN}, _From,
             #state{handler = on_replication,
                    subproto_state = #repl{last_received_lsn = ReceivedLSN} = Repl} = State) ->

diff --git a/src/epgsqla.erl b/src/epgsqla.erl
@@ -15,6 +15,7 @@
          get_parameter/2,
          set_notice_receiver/2,
          get_cmd_status/1,
+         get_os_pid/1,
          squery/2,
          equery/2, equery/3,
          prepared_query/3,
@@ -84,6 +85,10 @@ set_notice_receiver(C, PidOrName) ->
 get_cmd_status(C) ->
     epgsql_sock:get_cmd_status(C).
 
+-spec get_os_pid(epgsql:connection()) -> integer().
+get_os_pid(C) ->
+    epgsql_sock:get_os_pid(C).
+
 -spec squery(epgsql:connection(), epgsql:sql_query()) -> reference().
 squery(C, Sql) ->
     cast(C, epgsql_cmd_squery, Sql).

diff --git a/src/epgsqli.erl b/src/epgsqli.erl
@@ -15,6 +15,7 @@
          get_parameter/2,
          set_notice_receiver/2,
          get_cmd_status/1,
+         get_os_pid/1,
          squery/2,
          equery/2, equery/3,
          prepared_query/3,
@@ -81,6 +82,10 @@ set_notice_receiver(C, PidOrName) ->
 get_cmd_status(C) ->
     epgsql_sock:get_cmd_status(C).
 
+-spec get_os_pid(epgsql:connection()) -> integer().
+get_os_pid(C) ->
+    epgsql_sock:get_os_pid(C).
+
 -spec squery(epgsql:connection(), epgsql:sql_query()) -> reference().
 squery(C, Sql) ->
     incremental(C, epgsql_cmd_squery, Sql).

diff --git a/test/epgsql_SUITE.erl b/test/epgsql_SUITE.erl
@@ -139,7 +139,8 @@ groups() ->
         listen_notify,
         listen_notify_payload,
         set_notice_receiver,
-        get_cmd_status
+        get_cmd_status,
+        get_os_pid
     ],
     SubGroups ++
         [{epgsql, [], [{group, generic} | Tests]},
@@ -1419,6 +1420,14 @@ get_cmd_status(Config) ->
         ?assertEqual({ok, 'commit'}, Module:get_cmd_status(C))
     end).
 
+get_os_pid(Config) ->
+    Module = ?config(module, Config),
+    epgsql_ct:with_connection(Config, fun(C) ->
+        {ok, [#column{}], [{PidBin}]} = Module:squery(C, "SELECT pg_backend_pid()"),
+        Pid = Module:get_os_pid(C),
+        ?assertEqual(PidBin, integer_to_binary(Pid))
+    end).
+
 range_type(Config) ->
     epgsql_ct:with_min_version(Config, [9, 2], fun(_C) ->
         check_type(Config, int4range, "int4range(10, 20)", {10, 20}, [

diff --git a/test/epgsql_cast.erl b/test/epgsql_cast.erl
@@ -6,7 +6,8 @@
 -module(epgsql_cast).
 
 -export([connect/1, connect/2, connect/3, connect/4, close/1]).
--export([get_parameter/2, set_notice_receiver/2, get_cmd_status/1, squery/2, equery/2, equery/3]).
+-export([get_parameter/2, set_notice_receiver/2, get_cmd_status/1, get_os_pid/1,
+         squery/2, equery/2, equery/3]).
 -export([prepared_query/3]).
 -export([parse/2, parse/3, parse/4, describe/2, describe/3]).
 -export([bind/3, bind/4, execute/2, execute/3, execute/4, execute_batch/2, execute_batch/3]).
@@ -58,6 +59,9 @@ set_notice_receiver(C, PidOrName) ->
 get_cmd_status(C) ->
     epgsqla:get_cmd_status(C).
 
+get_os_pid(C) ->
+    epgsqla:get_os_pid(C).
+
 squery(C, Sql) ->
     Ref = epgsqla:squery(C, Sql),
     receive_result(C, Ref).

diff --git a/test/epgsql_incremental.erl b/test/epgsql_incremental.erl
@@ -6,7 +6,8 @@
 -module(epgsql_incremental).
 
 -export([connect/1, connect/2, connect/3, connect/4, close/1]).
--export([get_parameter/2, set_notice_receiver/2, get_cmd_status/1, squery/2, equery/2, equery/3]).
+-export([get_parameter/2, set_notice_receiver/2, get_cmd_status/1, get_os_pid/1,
+         squery/2, equery/2, equery/3]).
 -export([prepared_query/3]).
 -export([parse/2, parse/3, parse/4, describe/2, describe/3]).
 -export([bind/3, bind/4, execute/2, execute/3, execute/4, execute_batch/2, execute_batch/3]).
@@ -56,6 +57,9 @@ set_notice_receiver(C, PidOrName) ->
 get_cmd_status(C) ->
     epgsqli:get_cmd_status(C).
 
+get_os_pid(C) ->
+    epgsqli:get_os_pid(C).
+
 squery(C, Sql) ->
     Ref = epgsqli:squery(C, Sql),
     case receive_results(C, Ref, []) of