From 408e23adb383c102f2c49719aad5a888fb8df3a9 Mon Sep 17 00:00:00 2001
From: Dan Sosedoff <dan.sosedoff@gmail.com>
Date: Fri, 12 Jan 2024 21:17:14 -0600
Subject: [PATCH] Allow setting readonly mode in bookmarks (#707)

---
 pkg/bookmarks/bookmarks.go      | 16 +++++++++-------
 pkg/bookmarks/bookmarks_test.go | 18 +++++++++++-------
 pkg/client/client.go            | 14 ++++++++++++--
 pkg/client/client_test.go       |  8 ++++++++
 4 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/pkg/bookmarks/bookmarks.go b/pkg/bookmarks/bookmarks.go
index 7b324a1..1071c31 100644
--- a/pkg/bookmarks/bookmarks.go
+++ b/pkg/bookmarks/bookmarks.go
@@ -20,6 +20,7 @@ type Bookmark struct {
 	Database    string          // Database name
 	SSLMode     string          // Connection SSL mode
 	SSH         *shared.SSHInfo // SSH tunnel config
+	ReadOnly    bool            // Enable read-only transaction mode
 }
 
 // SSHInfoIsEmpty returns true if ssh configuration is not provided
@@ -40,12 +41,13 @@ func (b Bookmark) ConvertToOptions() command.Options {
 	}
 
 	return command.Options{
-		URL:     b.URL,
-		Host:    b.Host,
-		Port:    b.Port,
-		User:    user,
-		Pass:    pass,
-		DbName:  b.Database,
-		SSLMode: b.SSLMode,
+		URL:      b.URL,
+		Host:     b.Host,
+		Port:     b.Port,
+		User:     user,
+		Pass:     pass,
+		DbName:   b.Database,
+		SSLMode:  b.SSLMode,
+		ReadOnly: b.ReadOnly,
 	}
 }
diff --git a/pkg/bookmarks/bookmarks_test.go b/pkg/bookmarks/bookmarks_test.go
index b41bbb0..2f5b480 100644
--- a/pkg/bookmarks/bookmarks_test.go
+++ b/pkg/bookmarks/bookmarks_test.go
@@ -68,6 +68,7 @@ func TestBookmarkWithVarsConvertToOptions(t *testing.T) {
 
 		t.Setenv("DB_USER", "user123")
 		t.Setenv("DB_PASSWORD", "password123")
+
 		opt := b.ConvertToOptions()
 		assert.Equal(t, expOpt, opt)
 	})
@@ -87,6 +88,7 @@ func TestBookmarkWithVarsConvertToOptions(t *testing.T) {
 
 		t.Setenv("DB_USER", "user123")
 		t.Setenv("DB_PASSWORD", "password123")
+
 		opt := b.ConvertToOptions()
 		assert.Equal(t, expOpt, opt)
 	})
@@ -101,16 +103,18 @@ func TestBookmarkConvertToOptions(t *testing.T) {
 		Password: "password",
 		Database: "mydatabase",
 		SSLMode:  "disable",
+		ReadOnly: true,
 	}
 
 	expOpt := command.Options{
-		URL:     "postgres://username:password@host:port/database?sslmode=disable",
-		Host:    "localhost",
-		Port:    5432,
-		User:    "postgres",
-		Pass:    "password",
-		DbName:  "mydatabase",
-		SSLMode: "disable",
+		URL:      "postgres://username:password@host:port/database?sslmode=disable",
+		Host:     "localhost",
+		Port:     5432,
+		User:     "postgres",
+		Pass:     "password",
+		DbName:   "mydatabase",
+		SSLMode:  "disable",
+		ReadOnly: true,
 	}
 
 	opt := b.ConvertToOptions()
diff --git a/pkg/client/client.go b/pkg/client/client.go
index 00ff7d1..5bb4d44 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -41,6 +41,7 @@ type Client struct {
 	serverType       string
 	lastQueryTime    time.Time
 	queryTimeout     time.Duration
+	readonly         bool
 	closed           bool
 	External         bool             `json:"external"`
 	History          []history.Record `json:"history"`
@@ -166,7 +167,16 @@ func NewFromBookmark(bookmark *bookmarks.Bookmark) (*Client, error) {
 		sshInfo = bookmark.SSH
 	}
 
-	return NewFromUrl(connStr, sshInfo)
+	client, err := NewFromUrl(connStr, sshInfo)
+	if err != nil {
+		return nil, err
+	}
+
+	if bookmark.ReadOnly {
+		client.readonly = true
+	}
+
+	return client, nil
 }
 
 func (client *Client) init() {
@@ -476,7 +486,7 @@ func (client *Client) query(query string, args ...interface{}) (*Result, error)
 
 	// We're going to force-set transaction mode on every query.
 	// This is needed so that default mode could not be changed by user.
-	if command.Opts.ReadOnly {
+	if command.Opts.ReadOnly || client.readonly {
 		if err := client.SetReadOnlyMode(); err != nil {
 			return nil, err
 		}
diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go
index f2488b9..04c3afe 100644
--- a/pkg/client/client_test.go
+++ b/pkg/client/client_test.go
@@ -679,6 +679,14 @@ func testReadOnlyMode(t *testing.T) {
 
 	_, err = client.Query("/* CREATE TABLE foobar(id integer); */ SELECT 'foo';")
 	assert.NoError(t, err)
+
+	t.Run("with local readonly flag", func(t *testing.T) {
+		command.Opts.ReadOnly = false
+		client.readonly = true
+
+		_, err := client.Query("INSERT INTO foobar(id) VALUES(1)")
+		assert.Error(t, err, "query contains keywords not allowed in read-only mode")
+	})
 }
 
 func testTablesStats(t *testing.T) {