Feat/query cancelation

CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ## [0.1.2] - 2026-05-20
 
+### Added
+- **Query Cancellation**: Press `Ctrl+C` during a running query to cancel it immediately via PostgreSQL's out-of-band cancel signal.
+
 ### Fixed
 - **hardcoded style**: Removed the hardcoded "monokai" style from the syntax highlighter.
 

internal/app/app.go

@@ -13,6 +13,7 @@ import (
 	"time"
 
 	tea "charm.land/bubbletea/v2"
+	"github.com/balaji01-4d/pgxspecial"
 	"github.com/balajz/pgxcli/internal/app/commands"
 	"github.com/balajz/pgxcli/internal/app/renderer"
 	"github.com/balajz/pgxcli/internal/app/ui"
@@ -22,13 +23,12 @@ import (
 	"github.com/balajz/pgxcli/internal/database"
 	"github.com/balajz/pgxcli/internal/database/result"
 	"github.com/balajz/pgxcli/internal/parser"
-	"github.com/balaji01-4d/pgxspecial"
 )
 
 // Application defines the interface for the main application logic.
 type Application interface {
 	// Start starts the main repl loop, reading input, executing commands and printing results until the user exits.
-	Start(ctx context.Context, client *database.Client) error
+	Start(ctx context.Context) error
 
 	// Close performs saving history before exiting.
 	Close() error
@@ -46,20 +46,22 @@ type pgxCLI struct {
 	config    *config.Config
 	logger    *slog.Logger
 	completer *completer.Completer
+	client    *database.Client
 }
 
-func New(cfg *config.Config, printer cliio.Printer, logger *slog.Logger, completer *completer.Completer) (Application, error) {
+func New(cfg *config.Config, printer cliio.Printer, logger *slog.Logger, completer *completer.Completer, client *database.Client) (Application, error) {
 	return &pgxCLI{
 		config:    cfg,
 		logger:    logger,
 		Printer:   printer,
 		completer: completer,
+		client:    client,
 	}, nil
 }
 
-func (p *pgxCLI) execute(ctx context.Context, client *database.Client, query string) tea.Cmd {
+func (p *pgxCLI) execute(ctx context.Context, query string) tea.Cmd {
 	promptReady := func() tea.Msg {
-		prefix := client.ParsePrompt(p.config.Main.Prompt)
+		prefix := p.client.ParsePrompt(p.config.Main.Prompt)
 		return ui.ReadyMsg{Prefix: prefix} // this is used to unblock input after executing a command
 	}
 
@@ -72,15 +74,15 @@ func (p *pgxCLI) execute(ctx context.Context, client *database.Client, query str
 	}
 
 	return func() tea.Msg {
-		metaResult, okay, err := client.ExecuteSpecial(ctx, query)
+		metaResult, okay, err := p.client.ExecuteSpecial(ctx, query)
 		if err != nil {
 			p.logger.Error("error executing special command", "error", err)
 			return ui.ExecCmdMsg{Cmd: tea.Sequence(p.printError(err), promptReady)}
 		}
 		if okay {
 			start := time.Now()
 			p.logger.Debug("special command executed", "result_kind", metaResult.ResultKind())
-			result, quit, err := p.handleSpecialCommand(ctx, metaResult, client)
+			result, quit, err := p.handleSpecialCommand(ctx, metaResult, p.client)
 			if quit {
 				p.logger.Info("REPL exiting via quit command")
 				return ui.ExecCmdMsg{Cmd: tea.Quit}
@@ -110,7 +112,7 @@ func (p *pgxCLI) execute(ctx context.Context, client *database.Client, query str
 				continue
 			}
 
-			queryResult, err := client.ExecuteQuery(ctx, stmt)
+			queryResult, err := p.client.ExecuteQuery(ctx, stmt)
 			if err != nil {
 				p.logger.Error("query execution failed", "error", err)
 				cmds = append(cmds, p.printError(err))
@@ -136,13 +138,20 @@ func (p *pgxCLI) execute(ctx context.Context, client *database.Client, query str
 	}
 }
 
-func (p *pgxCLI) Start(ctx context.Context, client *database.Client) error {
+func (p *pgxCLI) Start(ctx context.Context) error {
 	executeFunc := func(query string) tea.Cmd {
-		return p.execute(ctx, client, query)
+		return p.execute(ctx, query)
 	}
 
-	initialPrefix := client.ParsePrompt(p.config.Main.Prompt)
-	m, err := ui.New(initialPrefix, p.completer.GetKeyWords(), p.config.Main.HistoryFile, string(p.config.Main.Style), executeFunc)
+	initialPrefix := p.client.ParsePrompt(p.config.Main.Prompt)
+	m, err := ui.New(
+		initialPrefix,
+		p.completer.GetKeyWords(),
+		p.config.Main.HistoryFile,
+		string(p.config.Main.Style),
+		executeFunc,
+		p.Cancel,
+	)
 	if err != nil {
 		return fmt.Errorf("creating UI model: %w", err)
 	}
@@ -223,6 +232,13 @@ func (p *pgxCLI) handleSpecialCommand(ctx context.Context, metaResult pgxspecial
 	}
 }
 
+func (p *pgxCLI) Cancel(ctx context.Context) error {
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
+	return p.client.Cancel(ctx)
+}
+
 func (p *pgxCLI) handleQueryResult(r result.Result) (tea.Cmd, error) {
 	res, ok := r.(*result.QueryResult)
 	if !ok {

internal/app/renderer/table_renderer.go

@@ -4,8 +4,8 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/balajz/pgxcli/internal/config"
 	"github.com/balaji01-4d/pgxspecial"
+	"github.com/balajz/pgxcli/internal/config"
 )
 
 type rowsTableResult interface {

internal/app/renderer/table_renderer_test.go

@@ -4,8 +4,8 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/balajz/pgxcli/internal/config"
 	"github.com/balaji01-4d/pgxspecial"
+	"github.com/balajz/pgxcli/internal/config"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )

internal/app/ui/model.go

@@ -3,6 +3,7 @@ package ui
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"os"
 	"os/exec"
@@ -38,6 +39,10 @@ type ReadyMsg struct{ Prefix string }
 // ExecCmdMsg is used to dispatch a batch/sequence of commands.
 type ExecCmdMsg struct{ Cmd tea.Cmd }
 
+type cancel func(ctx context.Context) error
+
+type execute func(query string) tea.Cmd
+
 type Model struct {
 	input         *editline.Model
 	width, height int
@@ -47,10 +52,11 @@ type Model struct {
 	style         string
 
 	// execute executes a query passed and return as ExecCmdMsg + ReadyMsg.
-	execute func(string) tea.Cmd
+	execute execute
+	cancel  cancel
 }
 
-func New(initialPrefix string, pgKeywords []string, historyFile string, style string, executeFunc func(string) tea.Cmd) (*Model, error) {
+func New(initialPrefix string, pgKeywords []string, historyFile string, style string, executeFunc execute, cancelFunc cancel) (*Model, error) {
 	el := editline.New(0, 0)
 	el.Prompt = initialPrefix
 	if historyFile == "" || historyFile == config.Default {
@@ -66,6 +72,7 @@ func New(initialPrefix string, pgKeywords []string, historyFile string, style st
 		historyFile: historyFile,
 		style:       style,
 		execute:     executeFunc,
+		cancel:      cancelFunc,
 	}, nil
 }
 
@@ -97,13 +104,20 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		return m, nil
 
 	case tea.KeyMsg:
-		if m.executing {
-			return m, nil
-		}
-
 		switch msg.String() {
 		case "ctrl+c":
 			m.input.Reset()
+			if m.executing {
+				m.executing = false
+				cancelFn := m.cancel
+				return m, func() tea.Msg {
+					if err := cancelFn(context.Background()); err != nil {
+						return ExecCmdMsg{Cmd: PrintErrCmd(err)}
+					}
+					return nil
+				}
+			}
+
 			return m, nil
 		}
 	}
@@ -252,6 +266,33 @@ func applyEditlineConfig(el *editline.Model, historyFile string, pgKeywords []st
 	)
 	el.AutoComplete = postgresAutocomplete(pgKeywords)
 
+	// Configure multi-line input detection:
+	// A SQL query is only complete (submittable) if it is empty, starts with a backslash
+	// (special command), or ends with a semicolon ';'. Otherwise, Enter should just insert a newline.
+	el.CheckInputComplete = func(entireInput [][]rune, line, col int) bool {
+		var sb strings.Builder
+		for i, rline := range entireInput {
+			if i > 0 {
+				sb.WriteByte('\n')
+			}
+			sb.WriteString(string(rline))
+		}
+		input := strings.TrimSpace(sb.String())
+
+		// If input is empty, let it submit (no-op)
+		if input == "" {
+			return true
+		}
+
+		// If it's a special command (e.g. \clear, \d), it's single-line and complete
+		if strings.HasPrefix(input, "\\") {
+			return true
+		}
+
+		// Otherwise, a SQL statement is complete only if it ends with a semicolon
+		return strings.HasSuffix(input, ";")
+	}
+
 	entries, err := history.LoadHistory(historyFile)
 	if err != nil {
 		return fmt.Errorf("loading history: %w", err)

internal/cli/root.go

@@ -80,7 +80,7 @@ func NewRootCmd(ctx context.Context, cliCtx *CliContext) *cobra.Command {
 			if !bool(interactiveConnFlag) {
 				ui.PrintBanner(version)
 			}
-			return cliCtx.App.Start(ctx, cliCtx.Client)
+			return cliCtx.App.Start(ctx)
 		},
 
 		PersistentPostRunE: func(_ *cobra.Command, _ []string) error {
@@ -358,7 +358,7 @@ func ensureConnected(cliCtx *CliContext) error {
 // which includes setting up the logger, config and autocompleter with PostgreSQL keywords.
 func initApplication(cliCtx *CliContext) error {
 	completer := completer.New(cliCtx.Logger.Logger)
-	pgxCLI, err := app.New(cliCtx.config, cliCtx.Printer, cliCtx.Logger.Logger, completer)
+	pgxCLI, err := app.New(cliCtx.config, cliCtx.Printer, cliCtx.Logger.Logger, completer, cliCtx.Client)
 	if err != nil {
 		cliCtx.Logger.Error("Failed to initialize app", "error", err)
 		return err

internal/database/client.go

@@ -9,8 +9,8 @@ import (
 	"strings"
 	"time"
 
-	"github.com/balajz/pgxcli/internal/database/result"
 	"github.com/balaji01-4d/pgxspecial"
+	"github.com/balajz/pgxcli/internal/database/result"
 )
 
 const nilPlaceholder = "(nil)"
@@ -103,6 +103,13 @@ func (c *Client) ChangeDatabase(ctx context.Context, dbName string) error {
 	return nil
 }
 
+func (c *Client) Cancel(ctx context.Context) error {
+	if !c.IsConnected() {
+		return ErrConnectionNotEstablished
+	}
+	return c.executor.cancel(ctx)
+}
+
 // ParsePrompt resolves prompt placeholders using current connection metadata.
 func (c *Client) ParsePrompt(str string) string {
 	var user, host, shortHost, db, port string

internal/database/executor.go

@@ -2,22 +2,29 @@ package database
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"time"
 
-	"github.com/balajz/pgxcli/internal/database/result"
 	"github.com/balaji01-4d/pgxspecial"
+	"github.com/balajz/pgxcli/internal/database/result"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 )
 
+var (
+	ErrConnectionClosed         = errors.New("connection closed unexpectedly")
+	ErrConnectionNotEstablished = errors.New("connection not established")
+)
+
 type conn interface {
 	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
 	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
 	Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
 	Config() *pgx.ConnConfig
+	PgConn() *pgconn.PgConn
 	Ping(ctx context.Context) error
 	Close(ctx context.Context) error
 }
@@ -108,6 +115,13 @@ func (e *executor) executeSpecial(ctx context.Context, cmd string) (pgxspecial.S
 	return normalizedRows, ok, nil
 }
 
+func (e *executor) cancel(ctx context.Context) error {
+	if e.Conn == nil {
+		return ErrConnectionNotEstablished
+	}
+	return e.Conn.PgConn().CancelRequest(ctx)
+}
+
 func (e *executor) close(ctx context.Context) error {
 	if e.Conn != nil {
 		return e.Conn.Close(ctx)
@@ -117,7 +131,7 @@ func (e *executor) close(ctx context.Context) error {
 
 func (e *executor) ping(ctx context.Context) error {
 	if e.Conn == nil {
-		return fmt.Errorf("database not connected")
+		return ErrConnectionNotEstablished
 	}
 	return e.Conn.Ping(ctx)
 }

internal/database/executor_test.go

@@ -185,9 +185,8 @@ func TestExecutorPing(t *testing.T) {
 			withConn: true,
 		},
 		{
-			name:       "no connection",
-			wantErr:    true,
-			wantErrMsg: "database not connected",
+			name:    "no connection",
+			wantErr: true,
 		},
 	}
 

internal/database/mocks_test.go

@@ -31,6 +31,7 @@ func (mc *MockConn) Ping(ctx context.Context) error {
 func (mc *MockConn) QueryRow(_ context.Context, _ string, _ ...any) pgx.Row { return nil }
 func (mc *MockConn) Close(_ context.Context) error                          { return nil }
 func (mc *MockConn) Config() *pgx.ConnConfig                                { return nil }
+func (mc *MockConn) PgConn() *pgconn.PgConn                                 { return nil }
 
 type MockRows struct {
 	mock.Mock