Glauber/region (#970)

This code is necessary to make local regions work.

Author: glommer

internal/cmd/db.go

@@ -182,7 +182,13 @@ func closestLocation(client *turso.Client) (string, error) {
 		return closest, nil
 	}
 
-	closest, err := client.Locations.Closest()
+	settings, err := settings.ReadSettings()
+	if err != nil {
+		return "", fmt.Errorf("could not read settings file to probe for locations: %w", err)
+	}
+
+	regionUrl := settings.GetRegionURL()
+	closest, err := client.Locations.Closest(regionUrl)
 	if err != nil {
 		// We fallback to ams if we are unable to probe the closest location.
 		return "ams", err

internal/cmd/db_shell.go

@@ -20,16 +20,10 @@ import (
 	"github.com/tursodatabase/turso-cli/internal/turso"
 )
 
-var proxy string
-
 func init() {
 	dbCmd.AddCommand(shellCmd)
 	addInstanceFlag(shellCmd, "Connect to the database at the specified instance.")
 	addLocationFlag(shellCmd, "Connect to the database at the specified location.")
-	shellCmd.Flags().StringVar(&proxy, "proxy", "", "Proxy to use for the connection.")
-	shellCmd.RegisterFlagCompletionFunc("proxy", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
-		return []string{}, cobra.ShellCompDirectiveNoFileComp
-	})
 	flags.AddAttachClaims(shellCmd)
 }
 
@@ -222,7 +216,7 @@ var shellCmd = &cobra.Command{
 
 		shellConfig := shell.ShellConfig{
 			DbUri:          dbUrl,
-			Proxy:          proxy,
+			Proxy:          getTursoProxyUrl(),
 			AuthToken:      authToken,
 			InF:            cmd.InOrStdin(),
 			OutF:           cmd.OutOrStdout(),

internal/cmd/utils.go

@@ -261,6 +261,11 @@ func getTursoUrl() string {
 	return url
 }
 
+func getTursoProxyUrl() string {
+	config, _ := settings.ReadSettings() // ok to ignore, we'll fallback to default
+	return config.GetProxyURL()
+}
+
 func promptConfirmation(prompt string) (bool, error) {
 	reader := bufio.NewReader(os.Stdin)
 

internal/settings/settings.go

@@ -2,6 +2,7 @@ package settings
 
 import (
 	"fmt"
+	"net/url"
 	"os"
 	"path"
 	"path/filepath"
@@ -34,6 +35,7 @@ func ReadSettings() (*Settings, error) {
 	configPath := configdir.LocalConfig("turso")
 	viper.BindEnv("config-path", "TURSO_CONFIG_FOLDER")
 	viper.BindEnv("baseURL", "TURSO_API_BASEURL")
+	viper.BindEnv("proxyURL", "TURSO_PROXY_URL")
 
 	configPathFlag := viper.GetString("config-path")
 	if len(configPathFlag) > 0 {
@@ -143,6 +145,21 @@ func (s *Settings) GetBaseURL() string {
 	return viper.GetString("baseURL")
 }
 
+func (s *Settings) GetProxyURL() string {
+	return viper.GetString("proxyURL")
+}
+
+func (s *Settings) GetRegionURL() string {
+	baseURL := s.GetProxyURL()
+	parsedURL, err := url.Parse(baseURL)
+	if err != nil {
+		return "https://region.turso.io"
+	}
+
+	parsedURL.Host = "region." + parsedURL.Host
+	return parsedURL.String()
+}
+
 func (s *Settings) SetAutoupdate(autoupdate string) {
 	config := viper.GetStringMap("config")
 	if config == nil {

internal/turso/locations.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
-	"time"
 	"unicode/utf8"
 
 	"github.com/rodaine/table"
@@ -53,8 +52,8 @@ type ClosestLocationResponse struct {
 	Server string
 }
 
-func (c *LocationsClient) Closest() (string, error) {
-	r, err := c.client.Get("https://region.turso.io", nil)
+func (c *LocationsClient) Closest(regionUrl string) (string, error) {
+	r, err := c.client.Get(regionUrl, nil)
 	if err != nil {
 		return "", fmt.Errorf("failed to request closest: %s", err)
 	}
@@ -73,35 +72,6 @@ func (c *LocationsClient) Closest() (string, error) {
 	return data.Server, nil
 }
 
-func ProbeLocation(location string) *time.Duration {
-	client := &http.Client{Timeout: 2 * time.Second}
-	req, err := http.NewRequest("GET", "http://region.turso.io:8080/", nil)
-	if err != nil {
-		return nil
-	}
-	req.Header.Add("fly-prefer-region", location)
-
-	start := time.Now()
-	r, err := client.Do(req)
-	if err != nil {
-		return nil
-	}
-	defer r.Body.Close()
-
-	dur := time.Since(start)
-	if r.StatusCode != http.StatusOK {
-		return nil
-	}
-	data, err := unmarshal[ClosestLocationResponse](r)
-	if err != nil {
-		return nil
-	}
-	if data.Server != location {
-		return nil
-	}
-	return &dur
-}
-
 func LocationsTable(columns []interface{}) table.Table {
 	regex := regexp.MustCompile(`\x1b\[[0-9;]*m`)
 	return table.New(columns...).WithWidthFunc(func(s string) int {

internal/turso/turso.go

@@ -1,14 +1,17 @@
 package turso
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"mime/multipart"
+	"net"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
 	"os"
 	"runtime"
+	"strings"
 
 	"github.com/tursodatabase/turso-cli/internal/flags"
 )
@@ -101,7 +104,26 @@ func (t *Client) do(method, path string, body io.Reader, extraHeaders map[string
 	if err != nil {
 		return nil, err
 	}
-	resp, err := http.DefaultClient.Do(req)
+
+	var client = &http.Client{
+		// go seems to have issues resolving things like region.localhost:9090, while curl works fine.
+		// Make sure we can get it to work in those cases by resolving it manually.
+		Transport: &http.Transport{
+			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+				host, port, err := net.SplitHostPort(addr)
+				if err != nil {
+					return nil, err
+				}
+
+				if strings.HasSuffix(host, ".localhost") {
+					return (&net.Dialer{}).DialContext(ctx, network, net.JoinHostPort("127.0.0.1", port))
+				}
+
+				return (&net.Dialer{}).DialContext(ctx, network, addr)
+			},
+		},
+	}
+	resp, err := client.Do(req)
 	if err != nil {
 		return nil, err
 	}