pgruntime: add ConnectionTarget.ApplicationName attribute

1 files changed, 33 insertions, 5 deletions

diff --git a/pgruntime/target.go b/pgruntime/target.go
index a03c8c5..f1f4729 100644
--- a/pgruntime/target.go
+++ b/pgruntime/target.go
@@ -8,6 +8,7 @@ import (
 	"maps"
 	"net"
 	"net/url"
+	"os"
 	"strings"
 )
 
@@ -18,6 +19,9 @@ import (
 //     Its value must be formatted in the syntax accepted by [url.ParseQuery].
 //   - The ExtraConnectionOptions field is intended for options coming in via code.
 //   - Both ConnectionOptions fields accept all query parameters that are allowed in [libpq-style connection URIs].
+//   - ApplicationName is an alternative way of setting ExtraConnectionOptions["fallback_application_name"].
+//     If filled, pgruntime will also try to automatically append the hostname if available.
+//     An explicit "application_name" setting in either ConnectionOptions or ExtraConnectionOptions takes precedence over this setting.
 //
 // This type is intended to be retrieved from environment variables, for example:
 //
@@ -39,6 +43,7 @@ type ConnectionTarget struct {
 	DatabaseName           string
 	ConnectionOptions      string
 	ExtraConnectionOptions url.Values
+	ApplicationName        string
 }
 
 const (
@@ -47,7 +52,7 @@ const (
 )
 
 // ParseConnectionTargetFromURL parses a [libpq-style connection URI] into a [ConnectionTarget] instance.
-// If there are connection options, they will be placed in the ConnectionOptions field, not in ExtraConnectionOptions.
+// If there are connection options, they will be placed in the ConnectionOptions field, not in ExtraConnectionOptions or ApplicationName.
 //
 // The special libpq syntax with multiple host:port pairs is not supported by this function and will result in an error.
 //
@@ -96,9 +101,20 @@ func ParseConnectionTargetFromURL(u *url.URL) (ConnectionTarget, error) {
 	return result, nil
 }
 
+// dependency injection slot to insert a double during tests
+var osHostname = os.Hostname
+
 // IntoURL formats the ConnectionTarget as a libpq-style connection URI.
 func (t ConnectionTarget) IntoURL() (*url.URL, error) {
-	rawQuery, err := mergeConnectionOptions(t.ConnectionOptions, t.ExtraConnectionOptions)
+	appName := t.ApplicationName
+	if appName != "" {
+		hostname, err := osHostname()
+		if err == nil {
+			appName += "@" + hostname
+		}
+	}
+
+	rawQuery, err := mergeConnectionOptions(t.ConnectionOptions, t.ExtraConnectionOptions, appName)
 	if err != nil {
 		return nil, fmt.Errorf("in ConnectionTarget.IntoURL: malformed connection options: %w", err)
 	}
@@ -127,12 +143,21 @@ func (t ConnectionTarget) buildHostPort() string {
 	}
 }
 
-func mergeConnectionOptions(opts string, extraOpts url.Values) (string, error) {
-	if len(extraOpts) == 0 {
+func mergeConnectionOptions(opts string, extraOpts url.Values, appName string) (string, error) {
+	if len(extraOpts) == 0 && appName == "" {
 		return opts, nil
 	}
 	if opts == "" {
-		return extraOpts.Encode(), nil
+		if appName == "" {
+			return extraOpts.Encode(), nil
+		} else {
+			cloned := maps.Clone(extraOpts)
+			if cloned == nil {
+				cloned = make(url.Values)
+			}
+			cloned.Set("fallback_application_name", appName)
+			return cloned.Encode(), nil
+		}
 	}
 
 	values, err := url.ParseQuery(opts)
@@ -140,5 +165,8 @@ func mergeConnectionOptions(opts string, extraOpts url.Values) (string, error) {
 		return "", err
 	}
 	maps.Copy(values, extraOpts)
+	if appName != "" {
+		values.Set("fallback_application_name", appName)
+	}
 	return values.Encode(), nil
 }