github.com/Redstoneguy129/cli@v0.0.0-20230211220159-15dca4e91917/internal/utils/connect.go (about)

     1  package utils
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/url"
     7  	"os"
     8  
     9  	"github.com/Redstoneguy129/cli/internal/debug"
    10  	"github.com/jackc/pgconn"
    11  	"github.com/jackc/pgx/v4"
    12  	"github.com/spf13/viper"
    13  )
    14  
    15  // Connnect to remote Postgres with optimised settings. The caller is responsible for closing the connection returned.
    16  func ConnectRemotePostgres(ctx context.Context, username, password, database, host string, options ...func(*pgx.ConnConfig)) (*pgx.Conn, error) {
    17  	// Use port 6543 for connection pooling
    18  	pgUrl := fmt.Sprintf(
    19  		"postgresql://%s@%s:6543/%s?connect_timeout=10",
    20  		url.UserPassword(username, password),
    21  		host,
    22  		url.PathEscape(database),
    23  	)
    24  	// Simple protocol is preferred over pgx default Parse -> Bind flow because
    25  	//   1. Using a single command for each query reduces RTT over an Internet connection.
    26  	//   2. Performance gains from using the alternate binary protocol is negligible because
    27  	//      we are only selecting from migrations table. Large reads are handled by PostgREST.
    28  	//   3. Any prepared statements are cleared server side upon closing the TCP connection.
    29  	//      Since CLI workloads are one-off scripts, we don't use connection pooling and hence
    30  	//      don't benefit from per connection server side cache.
    31  	opts := append(options, func(cc *pgx.ConnConfig) {
    32  		cc.PreferSimpleProtocol = true
    33  	})
    34  	conn, err := ConnectByUrl(ctx, pgUrl, opts...)
    35  	if !pgconn.Timeout(err) {
    36  		return conn, err
    37  	}
    38  	// Fallback to postgres when pgbouncer is unavailable
    39  	config := conn.Config()
    40  	config.Port = 5432
    41  	fmt.Fprintln(os.Stderr, "Retrying...", config.Host, config.Port)
    42  	return pgx.ConnectConfig(ctx, config)
    43  }
    44  
    45  // Connnect to local Postgres with optimised settings. The caller is responsible for closing the connection returned.
    46  func ConnectLocalPostgres(ctx context.Context, host string, port uint, database string, options ...func(*pgx.ConnConfig)) (*pgx.Conn, error) {
    47  	url := fmt.Sprintf("postgresql://postgres:postgres@%s:%d/%s?connect_timeout=2", host, port, database)
    48  	return ConnectByUrl(ctx, url, options...)
    49  }
    50  
    51  func ConnectByUrl(ctx context.Context, url string, options ...func(*pgx.ConnConfig)) (*pgx.Conn, error) {
    52  	// Parse connection url
    53  	config, err := pgx.ParseConfig(url)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	// Apply config overrides
    58  	for _, op := range options {
    59  		op(config)
    60  	}
    61  	if viper.GetBool("DEBUG") {
    62  		debug.SetupPGX(config)
    63  	}
    64  	// Connect to database
    65  	return pgx.ConnectConfig(ctx, config)
    66  }