Repository: supabase/cli
Pull request #322 (2022-08-01 05:02:59)

chore: adds debug flag for printing postgres messages
## What kind of change does this PR introduce?

a global debug flag

## What is the current behavior?

no debug logging

## What is the new behavior?

- adds `SUPABASE_DEBUG=1` env var and `--debug` flag
- prints postgres messages to stderr when debugging

Example
```bash
$ supabase --debug db remote set "postgresql://postgres:password@localhost:5432/postgres"
...
2022/08/01 13:21:14 PG Recv: {"Type":"DataRow","Values":[{"text":"20220729102259"}]}
2022/08/01 13:21:14 PG Recv: {"Type":"CommandComplete","CommandTag":"SELECT 1"}
2022/08/01 13:21:14 PG Recv: {"Type":"ReadyForQuery","TxStatus":"I"}
Error: The remote database was set up with a different Supabase CLI project.
```

## Additional context

Add any other context or screenshots.


coveralls: ## Pull Request Test Coverage Report for [Build 2772395667](https://coveralls.io/builds/51298842)

* **6** of **128**   **(4.69%)**  changed or added relevant lines in **4** files are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage decreased (**-0.4%**) to **20.074%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [cmd/db.go](https://coveralls.io/builds/51298842/source?filename=cmd%2Fdb.go#L98) | 0 | 3 | 0.0%
| [cmd/root.go](https://coveralls.io/builds/51298842/source?filename=cmd%2Froot.go#L28) | 6 | 9 | 66.67%
| [internal/db/commit/migra.go](https://coveralls.io/builds/51298842/source?filename=internal%2Fdb%2Fcommit%2Fmigra.go#L51) | 0 | 5 | 0.0%
| [internal/debug/postgres.go](https://coveralls.io/builds/51298842/source?filename=internal%2Fdebug%2Fpostgres.go#L21) | 0 | 111 | 0.0%
<!-- | **Total:** | **6** | **128** | **4.69%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/51298842/badge)](https://coveralls.io/builds/51298842) |
| :-- | --: |
| Change from base [Build 2772188662](https://coveralls.io/builds/51297936): |  -0.4% |
| Covered Lines: | 1029 |
| Relevant Lines: | 5126 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/cmd/db.go b/cmd/db.go
index c4f0ea8b..08c231e8 100644
--- a/cmd/db.go
+++ b/cmd/db.go
@@ -3,6 +3,7 @@ package cmd
 import (
 	"github.com/spf13/afero"
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
 	"github.com/supabase/cli/internal/db/branch/create"
 	"github.com/supabase/cli/internal/db/branch/delete"
 	"github.com/supabase/cli/internal/db/branch/list"
@@ -14,6 +15,7 @@ import (
 	"github.com/supabase/cli/internal/db/remote/set"
 	"github.com/supabase/cli/internal/db/reset"
 	"github.com/supabase/cli/internal/db/switch_"
+	"github.com/supabase/cli/internal/debug"
 )
 
 var (
@@ -93,6 +95,9 @@ var (
 		Short: "Set the remote database to push migrations to.",
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
+			if viper.GetBool("DEBUG") {
+				return set.Run(args[0], afero.NewOsFs(), debug.SetupPGX)
+			}
 			return set.Run(args[0], afero.NewOsFs())
 		},
 	}
diff --git a/cmd/root.go b/cmd/root.go
index e288bfd4..cb1d3c89 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -1,7 +1,11 @@
 package cmd
 
 import (
+	"strings"
+
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"github.com/spf13/viper"
 )
 
 // Passed from `-ldflags`: https://stackoverflow.com/q/11354518.
@@ -20,6 +24,16 @@ func Execute() {
 }
 
 func init() {
+	cobra.OnInitialize(func() {
+		viper.SetEnvPrefix("SUPABASE")
+		viper.AutomaticEnv()
+	})
+	flags := rootCmd.PersistentFlags()
+	flags.Bool("debug", false, "enable debug mode")
+	flags.VisitAll(func(f *pflag.Flag) {
+		key := strings.ReplaceAll(f.Name, "-", "_")
+		_ = viper.BindPFlag(key, flags.Lookup(f.Name))
+	})
 	rootCmd.SetVersionTemplate("{{.Version}}\n")
 }
 
diff --git a/go.mod b/go.mod
index b34639c6..105f1b58 100644
--- a/go.mod
+++ b/go.mod
@@ -85,7 +85,7 @@ require (
 	github.com/spf13/afero v1.8.2
 	github.com/spf13/cast v1.4.1 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
-	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.8.0
 	github.com/subosito/gotenv v1.2.0 // indirect
 	github.com/yuin/goldmark v1.4.11 // indirect
diff --git a/internal/db/commit/migra.go b/internal/db/commit/migra.go
index 9877f860..a0dfd9ef 100644
--- a/internal/db/commit/migra.go
+++ b/internal/db/commit/migra.go
@@ -21,6 +21,8 @@ import (
 	"github.com/docker/docker/pkg/stdcopy"
 	"github.com/jackc/pgx/v4"
 	"github.com/spf13/afero"
+	"github.com/spf13/viper"
+	"github.com/supabase/cli/internal/debug"
 	"github.com/supabase/cli/internal/utils"
 )
 
@@ -46,6 +48,11 @@ func RunMigra(name string, fsys afero.Fs) error {
 		}
 	}
 
+	var opts []func(*pgx.ConnConfig)
+	if viper.GetBool("DEBUG") {
+		opts = append(opts, debug.SetupPGX)
+	}
+
 	ctx := context.Background()
 	// Trap Ctrl+C and call cancel on the context:
 	// https://medium.com/@matryer/make-ctrl-c-cancel-the-context-context-bd006a8ad6ff
@@ -71,7 +78,7 @@ func RunMigra(name string, fsys afero.Fs) error {
 
 	fmt.Println("Initialising schema...")
 	url := fmt.Sprintf("postgresql://postgres:postgres@localhost:%d/%s", utils.Config.Db.Port, utils.ShadowDbName)
-	if err := applyMigrations(ctx, url, fsys); err != nil {
+	if err := applyMigrations(ctx, url, fsys, opts...); err != nil {
 		return err
 	}
 
diff --git a/internal/debug/postgres.go b/internal/debug/postgres.go
new file mode 100644
index 00000000..5bda31da
--- /dev/null
+++ b/internal/debug/postgres.go
@@ -0,0 +1,164 @@
+package debug
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"log"
+	"net"
+	"os"
+
+	"github.com/jackc/pgproto3/v2"
+	"github.com/jackc/pgx/v4"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+type Proxy struct {
+	dialContext func(ctx context.Context, network, addr string) (net.Conn, error)
+	errChan     chan error
+}
+
+func NewProxy() Proxy {
+	dialer := net.Dialer{}
+	return Proxy{
+		dialContext: dialer.DialContext,
+		errChan:     make(chan error, 1),
+	}
+}
+
+func SetupPGX(config *pgx.ConnConfig) {
+	proxy := Proxy{
+		dialContext: config.DialFunc,
+		errChan:     make(chan error, 1),
+	}
+	config.DialFunc = proxy.DialFunc
+	config.TLSConfig = nil
+}
+
+func (p *Proxy) DialFunc(ctx context.Context, network, addr string) (net.Conn, error) {
+	serverConn, err := p.dialContext(ctx, network, addr)
+	if err != nil {
+		return nil, err
+	}
+
+	const bufSize = 1024 * 1024
+	ln := bufconn.Listen(bufSize)
+	go func() {
+		defer serverConn.Close()
+		clientConn, err := ln.Accept()
+		if err != nil {
+			// Unreachable code as bufconn never throws, but just in case
+			panic(err)
+		}
+		defer clientConn.Close()
+
+		backend := NewBackend(clientConn)
+		frontend := NewFrontend(serverConn)
+		go backend.forward(frontend, p.errChan)
+		go frontend.forward(backend, p.errChan)
+
+		for {
+			// Since pgx closes connection first, every EOF is seen as unexpected
+			if err := <-p.errChan; err != nil && err != io.ErrUnexpectedEOF {
+				panic(err)
+			}
+		}
+	}()
+
+	return ln.DialContext(ctx)
+}
+
+type Backend struct {
+	*pgproto3.Backend
+	logger *log.Logger
+}
+
+func NewBackend(clientConn net.Conn) Backend {
+	return Backend{
+		pgproto3.NewBackend(pgproto3.NewChunkReader(clientConn), clientConn),
+		log.New(os.Stderr, "PG Recv: ", log.LstdFlags|log.Lmsgprefix),
+	}
+}
+
+func (b *Backend) forward(frontend Frontend, errChan chan error) {
+	startupMessage, err := b.ReceiveStartupMessage()
+	if err != nil {
+		errChan <- err
+		return
+	}
+
+	buf, err := json.Marshal(startupMessage)
+	if err != nil {
+		errChan <- err
+		return
+	}
+	frontend.logger.Println(string(buf))
+
+	if err = frontend.Send(startupMessage); err != nil {
+		errChan <- err
+		return
+	}
+
+	for {
+		msg, err := b.Receive()
+		if err != nil {
+			errChan <- err
+			return
+		}
+
+		buf, err := json.Marshal(msg)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		frontend.logger.Println(string(buf))
+
+		if err = frontend.Send(msg); err != nil {
+			errChan <- err
+			return
+		}
+	}
+}
+
+type Frontend struct {
+	*pgproto3.Frontend
+	logger *log.Logger
+}
+
+func NewFrontend(serverConn net.Conn) Frontend {
+	return Frontend{
+		pgproto3.NewFrontend(pgproto3.NewChunkReader(serverConn), serverConn),
+		log.New(os.Stderr, "PG Send: ", log.LstdFlags|log.Lmsgprefix),
+	}
+}
+
+func (f *Frontend) forward(backend Backend, errChan chan error) {
+	for {
+		msg, err := f.Receive()
+		if err != nil {
+			errChan <- err
+			return
+		}
+
+		buf, err := json.Marshal(msg)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		backend.logger.Println(string(buf))
+
+		if _, ok := msg.(pgproto3.AuthenticationResponseMessage); ok {
+			// Set the authentication type so the next backend.Receive() will
+			// properly decode the appropriate 'p' message.
+			if err := backend.SetAuthType(f.GetAuthType()); err != nil {
+				errChan <- err
+				return
+			}
+		}
+
+		if err := backend.Send(msg); err != nil {
+			errChan <- err
+			return
+		}
+	}
+}


Test patch:
diff --git a/internal/debug/postgres_test.go b/internal/debug/postgres_test.go
new file mode 100644
index 00000000..65f83ffc
--- /dev/null
+++ b/internal/debug/postgres_test.go
@@ -0,0 +1,31 @@
+package debug
+
+import (
+	"context"
+	"testing"
+
+	"github.com/jackc/pgx/v4"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/supabase/cli/internal/testing/pgtest"
+)
+
+func TestPostgresProxy(t *testing.T) {
+	const postgresUrl = "postgresql://postgres:password@localhost:5432/postgres"
+
+	t.Run("forwards messages between frontend and backend", func(t *testing.T) {
+		// Parse connection url
+		config, err := pgx.ParseConfig(postgresUrl)
+		require.NoError(t, err)
+		// Setup postgres mock
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Intercept(config)
+		// Run test
+		SetupPGX(config)
+		ctx := context.Background()
+		proxy, err := pgx.ConnectConfig(ctx, config)
+		assert.NoError(t, err)
+		assert.NoError(t, proxy.Close(ctx))
+	})
+}