diff --git a/cmd/project/executor.go b/cmd/project/executor.go index d09fdd7d..33f831f5 100644 --- a/cmd/project/executor.go +++ b/cmd/project/executor.go @@ -1,6 +1,8 @@ package project import ( + "database/sql" + "github.com/spf13/cobra" "github.com/shopware/shopware-cli/internal/executor" @@ -21,3 +23,36 @@ func resolveExecutor(cmd *cobra.Command, projectRoot string) (executor.Executor, return executor.New(projectRoot, envCfg, cfg) } + +// resolveProjectDatabaseConnection resolves the database credentials of the +// current environment through its executor. +func resolveProjectDatabaseConnection(cmd *cobra.Command) (*executor.DatabaseConnection, error) { + projectRoot, err := findClosestShopwareProject() + if err != nil { + return nil, err + } + + cmdExecutor, err := resolveExecutor(cmd, projectRoot) + if err != nil { + return nil, err + } + + return cmdExecutor.DatabaseConnection(cmd.Context()) +} + +// connectProjectDatabase resolves the database of the current environment and +// opens a single dedicated connection to it. The returned cleanup closes +// connection and pool. +func connectProjectDatabase(cmd *cobra.Command) (*sql.Conn, *executor.DatabaseConnection, func(), error) { + dbConn, err := resolveProjectDatabaseConnection(cmd) + if err != nil { + return nil, nil, nil, err + } + + conn, cleanup, err := dbConn.Open(cmd.Context()) + if err != nil { + return nil, nil, nil, err + } + + return conn, dbConn, cleanup, nil +} diff --git a/cmd/project/project_dump.go b/cmd/project/project_dump.go index ae654c6f..f7c56d94 100644 --- a/cmd/project/project_dump.go +++ b/cmd/project/project_dump.go @@ -2,23 +2,19 @@ package project import ( "compress/gzip" - "context" "database/sql" "errors" "fmt" "io" - "net" - "net/url" "os" "strings" - "time" "github.com/charmbracelet/x/term" "github.com/go-sql-driver/mysql" "github.com/klauspost/compress/zstd" "github.com/spf13/cobra" - "github.com/shopware/shopware-cli/internal/envfile" + "github.com/shopware/shopware-cli/internal/executor" "github.com/shopware/shopware-cli/internal/mysqldump" "github.com/shopware/shopware-cli/internal/shop" "github.com/shopware/shopware-cli/internal/system" @@ -143,22 +139,9 @@ var projectDatabaseDumpCmd = &cobra.Command{ } func assembleConnectionURI(cmd *cobra.Command) (*mysql.Config, error) { - cfg := &mysql.Config{ - Loc: time.UTC, - Net: "tcp", - ParseTime: false, - AllowNativePasswords: true, - CheckConnLiveness: true, - User: "root", - Passwd: "root", - Addr: "127.0.0.1:3306", - DBName: "shopware", - } - - if projectRoot, err := findClosestShopwareProject(); err == nil { - if err := loadDatabaseURLIntoConnection(cmd.Context(), projectRoot, cfg); err != nil { - return nil, err - } + dbConn, err := resolveDumpDatabaseConnection(cmd) + if err != nil { + return nil, err } host, _ := cmd.Flags().GetString("host") @@ -168,20 +151,20 @@ func assembleConnectionURI(cmd *cobra.Command) (*mysql.Config, error) { db, _ := cmd.Flags().GetString("database") if host != "" { - if port != "" { - cfg.Addr = fmt.Sprintf("%s:%s", host, port) - } else { - cfg.Addr = host - } + dbConn.Host = host + } + + if port != "" { + dbConn.Port = port } if db != "" { - cfg.DBName = db + dbConn.Database = db } if username != "" { - cfg.User = username - cfg.Passwd = "" + dbConn.Username = username + dbConn.Password = "" } if cmd.Flags().Changed("password") { @@ -202,57 +185,24 @@ func assembleConnectionURI(cmd *cobra.Command) (*mysql.Config, error) { return nil, fmt.Errorf("could not read password: %w", err) } - cfg.Passwd = string(pass) + dbConn.Password = string(pass) } else { - cfg.Passwd = password + dbConn.Password = password } } - return cfg, nil + return dbConn.MySQLConfig(), nil } -func loadDatabaseURLIntoConnection(ctx context.Context, projectRoot string, cfg *mysql.Config) error { - if err := envfile.LoadSymfonyEnvFile(projectRoot); err != nil { - return err - } - - databaseUrl := os.Getenv("DATABASE_URL") - - if databaseUrl == "" { - return nil - } - - logging.FromContext(ctx).Info("Using DATABASE_URL env as default connection string. options can override specific parts (--username=foo)") - - parsedUri, err := url.Parse(databaseUrl) - if err != nil { - return fmt.Errorf("could not parse DATABASE_URL: %w", err) - } - - if parsedUri.User != nil { - cfg.User = parsedUri.User.Username() - - if password, ok := parsedUri.User.Password(); ok { - cfg.Passwd = password - } else { - // Reset password if it is not set - cfg.Passwd = "" - } - } - - if parsedUri.Host != "" { - cfg.Addr = parsedUri.Host - - if parsedUri.Port() == "" { - cfg.Addr = net.JoinHostPort(parsedUri.Host, "3306") - } - } - - if parsedUri.Path != "" { - cfg.DBName = strings.Trim(parsedUri.Path, "/") +// resolveDumpDatabaseConnection resolves credentials like the other database +// commands, but keeps dump usable outside a Shopware project: there the +// process environment and the connection flags are all that is needed. +func resolveDumpDatabaseConnection(cmd *cobra.Command) (*executor.DatabaseConnection, error) { + if _, err := findClosestShopwareProject(); err != nil { + return executor.NewLocal("").DatabaseConnection(cmd.Context()) } - return nil + return resolveProjectDatabaseConnection(cmd) } func init() { diff --git a/cmd/project/project_dump_test.go b/cmd/project/project_dump_test.go new file mode 100644 index 00000000..e94f2eae --- /dev/null +++ b/cmd/project/project_dump_test.go @@ -0,0 +1,102 @@ +package project + +import ( + "os" + "path/filepath" + "testing" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newDumpFlagCommand(t *testing.T, flags map[string]string) *cobra.Command { + t.Helper() + + cmd := &cobra.Command{} + cmd.Flags().String("host", "", "") + cmd.Flags().String("port", "", "") + cmd.Flags().String("username", "", "") + cmd.Flags().String("password", "", "") + cmd.Flags().Lookup("password").NoOptDefVal = passwordFlagPrompt + cmd.Flags().String("database", "", "") + cmd.SetContext(t.Context()) + + for name, value := range flags { + require.NoError(t, cmd.Flags().Set(name, value)) + } + + return cmd +} + +// chdirOutsideProject moves into an empty directory so the project lookup +// fails and the environment-only fallback is used. +func chdirOutsideProject(t *testing.T) { + t.Helper() + t.Setenv("PROJECT_ROOT", "") + t.Chdir(t.TempDir()) +} + +func TestAssembleConnectionURIDefaults(t *testing.T) { + chdirOutsideProject(t) + t.Setenv("DATABASE_URL", "") + + cfg, err := assembleConnectionURI(newDumpFlagCommand(t, nil)) + require.NoError(t, err) + + assert.Equal(t, "127.0.0.1:3306", cfg.Addr) + assert.Equal(t, "root", cfg.User) + assert.Equal(t, "root", cfg.Passwd) + assert.Equal(t, "shopware", cfg.DBName) +} + +func TestAssembleConnectionURIFlagOverrides(t *testing.T) { + chdirOutsideProject(t) + t.Setenv("DATABASE_URL", "") + + cfg, err := assembleConnectionURI(newDumpFlagCommand(t, map[string]string{ + "host": "db.internal", + "port": "3307", + "username": "backup", + "password": "secret", + "database": "shop_prod", + })) + require.NoError(t, err) + + assert.Equal(t, "db.internal:3307", cfg.Addr) + assert.Equal(t, "backup", cfg.User) + assert.Equal(t, "secret", cfg.Passwd) + assert.Equal(t, "shop_prod", cfg.DBName) +} + +func TestAssembleConnectionURIUsernameClearsPassword(t *testing.T) { + chdirOutsideProject(t) + t.Setenv("DATABASE_URL", "") + + cfg, err := assembleConnectionURI(newDumpFlagCommand(t, map[string]string{"username": "backup"})) + require.NoError(t, err) + + assert.Equal(t, "backup", cfg.User) + assert.Empty(t, cfg.Passwd) +} + +func TestAssembleConnectionURIDatabaseURLInsideProject(t *testing.T) { + projectRoot := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(projectRoot, "bin"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(projectRoot, "bin", "console"), nil, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(projectRoot, "composer.json"), []byte(`{"require": {"shopware/core": "6.6.0"}}`), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(projectRoot, ".env"), []byte("DATABASE_URL=mysql://app:secret@db.example.com:3307/shop\n"), 0o644)) + + t.Setenv("PROJECT_ROOT", "") + t.Setenv("DATABASE_URL", "") + t.Setenv("SHOPWARE_CLI_NO_SYMFONY_CLI", "1") + t.Chdir(projectRoot) + + cfg, err := assembleConnectionURI(newDumpFlagCommand(t, map[string]string{"database": "other"})) + require.NoError(t, err) + + assert.Equal(t, "db.example.com:3307", cfg.Addr) + assert.Equal(t, "app", cfg.User) + assert.Equal(t, "secret", cfg.Passwd) + assert.Equal(t, "other", cfg.DBName, "flag overrides the URL database") +} diff --git a/cmd/project/project_restore.go b/cmd/project/project_restore.go new file mode 100644 index 00000000..ede49998 --- /dev/null +++ b/cmd/project/project_restore.go @@ -0,0 +1,119 @@ +package project + +import ( + "errors" + "fmt" + "io" + "os" + "time" + + "charm.land/huh/v2" + "github.com/spf13/cobra" + + "github.com/shopware/shopware-cli/internal/sqlshell" + "github.com/shopware/shopware-cli/internal/system" + "github.com/shopware/shopware-cli/logging" +) + +var projectDatabaseRestoreCmd = &cobra.Command{ + Use: "restore [file]", + Aliases: []string{"import"}, + Short: "Restores a SQL dump into the Shopware database", + Long: "Imports a plain, gzip- or zstd-compressed SQL file into the project database " + + "using the connection details of the current environment (local, docker, ...). " + + "The compression is detected from the file content, pass - to read from stdin.", + Example: ` shopware-cli project restore dump.sql + shopware-cli project restore dump.sql.gz + shopware-cli project restore dump.sql.zst + curl -s https://example.com/dump.sql.gz | shopware-cli project restore -`, + Args: cobra.ExactArgs(1), + SilenceUsage: true, + RunE: func(cmd *cobra.Command, args []string) error { + input := cmd.InOrStdin() + var totalSize int64 + + if args[0] != "-" { + file, err := os.Open(args[0]) + if err != nil { + return err + } + // The file is only read, a close error carries no information. + defer func() { _ = file.Close() }() + + if info, err := file.Stat(); err == nil { + totalSize = info.Size() + } + + input = file + } + + counting := &system.CountingReader{Reader: input} + + reader, err := system.DecompressReader(counting) + if err != nil { + return fmt.Errorf("could not open dump: %w", err) + } + + // Closing releases decompressor resources; read errors already + // surface through ExecuteStream. + if closer, ok := reader.(io.Closer); ok { + defer func() { _ = closer.Close() }() + } + + conn, dbConn, cleanup, err := connectProjectDatabase(cmd) + if err != nil { + return err + } + defer cleanup() + + force, _ := cmd.Flags().GetBool("force") + + // No prompt when reading the dump from stdin: it is not a terminal + // then, and the prompt would consume dump bytes. + if !force && system.IsInteractionEnabled(cmd.Context()) && isTerminalStream(cmd.InOrStdin()) { + confirmed := false + + if err := huh.NewConfirm(). + Title(fmt.Sprintf("Restore into database %q at %s?", dbConn.Database, dbConn.Addr())). + Description("Existing data will be overwritten by the dump."). + Value(&confirmed). + Run(); err != nil { + return err + } + + if !confirmed { + return errors.New("restore cancelled") + } + } + + logger := logging.FromContext(cmd.Context()) + logger.Infof("Restoring into database %q at %s", dbConn.Database, dbConn.Addr()) + + start := time.Now() + lastProgress := 0 + + statements, err := sqlshell.ExecuteStream(cmd.Context(), conn, reader, func(int) { + if totalSize <= 0 { + return + } + + progress := int(counting.BytesRead() * 100 / totalSize) + if progress >= lastProgress+10 { + lastProgress = progress + logger.Infof("Restore progress: %d%%", progress) + } + }) + if err != nil { + return fmt.Errorf("restore failed (%d statements executed): %w", statements, err) + } + + logger.Infof("Restored %d statements in %s", statements, time.Since(start).Round(time.Millisecond)) + + return nil + }, +} + +func init() { + projectRootCmd.AddCommand(projectDatabaseRestoreCmd) + projectDatabaseRestoreCmd.Flags().BoolP("force", "f", false, "skip the confirmation prompt") +} diff --git a/cmd/project/project_sql.go b/cmd/project/project_sql.go new file mode 100644 index 00000000..419f4116 --- /dev/null +++ b/cmd/project/project_sql.go @@ -0,0 +1,89 @@ +package project + +import ( + "errors" + "io" + "os" + "strings" + + "github.com/charmbracelet/x/term" + "github.com/spf13/cobra" + + "github.com/shopware/shopware-cli/internal/sqlshell" + "github.com/shopware/shopware-cli/internal/system" + "github.com/shopware/shopware-cli/logging" +) + +var projectSQLCmd = &cobra.Command{ + Use: "sql [query]", + Short: "Run SQL queries against the project database", + Long: "Connects to the project database using the connection details of the current environment (local, docker, ...), " + + "so you don't need to know the host or credentials. " + + "Without arguments an interactive SQL shell is opened; a query can be passed as argument or a script piped via stdin.", + Example: ` shopware-cli project sql "SELECT id, tax_rate FROM tax" + shopware-cli project sql --format json "SELECT * FROM sales_channel" | jq + shopware-cli project sql < script.sql + shopware-cli project sql`, + Args: cobra.ArbitraryArgs, + SilenceUsage: true, + RunE: func(cmd *cobra.Command, args []string) error { + format, err := resolveSQLFormat(cmd) + if err != nil { + return err + } + + conn, dbConn, cleanup, err := connectProjectDatabase(cmd) + if err != nil { + return err + } + defer cleanup() + + if len(args) > 0 { + return sqlshell.Run(cmd.Context(), conn, strings.Join(args, " "), cmd.OutOrStdout(), format) + } + + if !isTerminalStream(cmd.InOrStdin()) { + script, err := io.ReadAll(cmd.InOrStdin()) + if err != nil { + return err + } + + return sqlshell.Run(cmd.Context(), conn, string(script), cmd.OutOrStdout(), format) + } + + if !system.IsInteractionEnabled(cmd.Context()) { + return errors.New("no query given and interaction is disabled, pass a query as argument or pipe a script via stdin") + } + + logging.FromContext(cmd.Context()).Infof("Connected to database %q at %s. Type \"exit\" or press Ctrl+D to quit.", dbConn.Database, dbConn.Addr()) + + return sqlshell.InteractiveShell(cmd.Context(), conn, format) + }, +} + +func resolveSQLFormat(cmd *cobra.Command) (sqlshell.Format, error) { + formatFlag, _ := cmd.Flags().GetString("format") + + if formatFlag == "" { + if isTerminalStream(cmd.OutOrStdout()) { + return sqlshell.FormatTable, nil + } + + return sqlshell.FormatTSV, nil + } + + return sqlshell.ParseFormat(formatFlag) +} + +// isTerminalStream reports whether a Cobra in/out stream is an interactive +// terminal. Streams replaced via cmd.SetIn/cmd.SetOut are never terminals. +func isTerminalStream(stream any) bool { + file, ok := stream.(*os.File) + + return ok && term.IsTerminal(file.Fd()) +} + +func init() { + projectRootCmd.AddCommand(projectSQLCmd) + projectSQLCmd.Flags().String("format", "", "output format: table, tsv, json (default: table when stdout is a terminal, tsv otherwise)") +} diff --git a/cmd/project/project_sql_test.go b/cmd/project/project_sql_test.go new file mode 100644 index 00000000..5c578f8b --- /dev/null +++ b/cmd/project/project_sql_test.go @@ -0,0 +1,61 @@ +package project + +import ( + "bytes" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/shopware/shopware-cli/internal/sqlshell" +) + +func newSQLFormatCommand(t *testing.T, formatFlag string) *cobra.Command { + t.Helper() + + cmd := &cobra.Command{} + cmd.Flags().String("format", "", "") + cmd.SetOut(&bytes.Buffer{}) + + if formatFlag != "" { + require.NoError(t, cmd.Flags().Set("format", formatFlag)) + } + + return cmd +} + +func TestResolveSQLFormatDefaultsToTSVWithoutTerminal(t *testing.T) { + format, err := resolveSQLFormat(newSQLFormatCommand(t, "")) + require.NoError(t, err) + + assert.Equal(t, sqlshell.FormatTSV, format) +} + +func TestResolveSQLFormatExplicit(t *testing.T) { + format, err := resolveSQLFormat(newSQLFormatCommand(t, "json")) + require.NoError(t, err) + + assert.Equal(t, sqlshell.FormatJSON, format) +} + +func TestResolveSQLFormatInvalid(t *testing.T) { + _, err := resolveSQLFormat(newSQLFormatCommand(t, "xml")) + assert.ErrorContains(t, err, "unknown format") +} + +func TestIsTerminalStream(t *testing.T) { + assert.False(t, isTerminalStream(strings.NewReader("not a file"))) + + path := filepath.Join(t.TempDir(), "plain.txt") + require.NoError(t, os.WriteFile(path, []byte("x"), 0o644)) + + file, err := os.Open(path) + require.NoError(t, err) + defer func() { _ = file.Close() }() + + assert.False(t, isTerminalStream(file), "a regular file is not a terminal") +} diff --git a/internal/docker/compose.go b/internal/docker/compose.go index b871f663..c07f8bac 100644 --- a/internal/docker/compose.go +++ b/internal/docker/compose.go @@ -218,6 +218,9 @@ func buildCompose(hasAMQP, hasElasticsearch bool, opts *ComposeOptions) yaml.Nod database := newMappingNode() addKeyValue(database, "image", "mariadb:11.8") + // Publish the database on a random loopback port so host-side tools + // (e.g. `shopware-cli project sql`) can reach it without port conflicts. + addKeyValueNode(database, "ports", newSequenceNode("127.0.0.1::3306")) addKeyValueNode(database, "environment", dbEnv) addKeyValueNode(database, "volumes", newSequenceNode("db-data:/var/lib/mysql:rw")) addKeyValueNode(database, "command", newSequenceNode( diff --git a/internal/docker/compose_test.go b/internal/docker/compose_test.go index f611bcbd..ba6871da 100644 --- a/internal/docker/compose_test.go +++ b/internal/docker/compose_test.go @@ -52,6 +52,7 @@ func TestGenerateComposeFile(t *testing.T) { assert.Contains(t, compose, "db-data:") assert.Contains(t, compose, "ghcr.io/shopware/docker-dev:php8.3-node24-caddy") assert.Contains(t, compose, "mariadb:11.8") + assert.Contains(t, compose, "127.0.0.1::3306") assert.Contains(t, compose, "mailpit") assert.NotContains(t, compose, "lavinmq") assert.NotContains(t, compose, "opensearch") diff --git a/internal/executor/database.go b/internal/executor/database.go new file mode 100644 index 00000000..76a0dddb --- /dev/null +++ b/internal/executor/database.go @@ -0,0 +1,156 @@ +package executor + +import ( + "context" + "database/sql" + "errors" + "fmt" + "net" + "net/url" + "os" + "strings" + "time" + + "github.com/go-sql-driver/mysql" + + "github.com/shopware/shopware-cli/internal/envfile" +) + +// DatabaseConnection describes how to reach the project database from the +// host machine. +type DatabaseConnection struct { + Host string + Port string + Username string + Password string + Database string +} + +// Addr returns the host:port address of the database. +func (c *DatabaseConnection) Addr() string { + return net.JoinHostPort(c.Host, c.Port) +} + +// MySQLConfig translates the credentials into a driver configuration. +func (c *DatabaseConnection) MySQLConfig() *mysql.Config { + cfg := mysql.NewConfig() + cfg.Net = "tcp" + cfg.Addr = c.Addr() + cfg.User = c.Username + cfg.Passwd = c.Password + cfg.DBName = c.Database + cfg.Loc = time.UTC + // 0 makes the driver fetch max_allowed_packet from the server on connect + // (see go-sql-driver's connector.go), so statements from large dumps are + // neither rejected client-side nor sent beyond what the server accepts. + cfg.MaxAllowedPacket = 0 + + return cfg +} + +// Open opens a single dedicated connection to the database, so session state +// (SET, USE, ...) survives across statements. The returned cleanup closes the +// connection and its pool. +func (c *DatabaseConnection) Open(ctx context.Context) (*sql.Conn, func(), error) { + db, err := sql.Open("mysql", c.MySQLConfig().FormatDSN()) + if err != nil { + return nil, nil, err + } + + connectCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + conn, err := db.Conn(connectCtx) + if err != nil { + _ = db.Close() + return nil, nil, fmt.Errorf("could not connect to database %q at %s: %w", c.Database, c.Addr(), err) + } + + cleanup := func() { + _ = conn.Close() + _ = db.Close() + } + + return conn, cleanup, nil +} + +func defaultDatabaseConnection() *DatabaseConnection { + return &DatabaseConnection{ + Host: "127.0.0.1", + Port: "3306", + Username: "root", + Password: "root", + Database: "shopware", + } +} + +// applyDatabaseURL merges a Symfony DATABASE_URL into conn. Parts missing in +// the URL keep their current value, except the password which is cleared when +// a user without password is given. +func applyDatabaseURL(conn *DatabaseConnection, databaseURL string) error { + parsed, err := url.Parse(databaseURL) + if err != nil { + return fmt.Errorf("could not parse DATABASE_URL: %w", err) + } + + // A bare word like "shopware" parses as a path-only URL. Silently keeping + // the defaults would send commands to the wrong database, so reject it. + if parsed.Scheme == "" || parsed.Hostname() == "" { + return errors.New("invalid DATABASE_URL: expected a URL like mysql://user:password@host:3306/dbname") + } + + if parsed.User != nil { + conn.Username = parsed.User.Username() + + if password, ok := parsed.User.Password(); ok { + conn.Password = password + } else { + conn.Password = "" + } + } + + if host := parsed.Hostname(); host != "" { + conn.Host = host + } + + if port := parsed.Port(); port != "" { + conn.Port = port + } + + if dbName := strings.Trim(parsed.Path, "/"); dbName != "" { + conn.Database = dbName + } + + return nil +} + +// databaseConnectionFromEnv resolves the connection for executors that run on +// the host. Precedence: executor env overrides > real environment variables > +// Symfony env files, matching how the spawned processes see DATABASE_URL. +func databaseConnectionFromEnv(projectRoot string, extraEnv map[string]string) (*DatabaseConnection, error) { + conn := defaultDatabaseConnection() + + databaseURL := extraEnv["DATABASE_URL"] + + if databaseURL == "" { + databaseURL = os.Getenv("DATABASE_URL") + } + + if databaseURL == "" { + fileValue, err := envfile.ReadValue(projectRoot, "DATABASE_URL") + if err != nil { + return nil, err + } + databaseURL = fileValue + } + + if databaseURL == "" { + return conn, nil + } + + if err := applyDatabaseURL(conn, databaseURL); err != nil { + return nil, err + } + + return conn, nil +} diff --git a/internal/executor/database_docker_test.go b/internal/executor/database_docker_test.go new file mode 100644 index 00000000..fe0aac03 --- /dev/null +++ b/internal/executor/database_docker_test.go @@ -0,0 +1,122 @@ +package executor + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeFakeDocker puts a docker stub on PATH that answers +// `docker compose exec ... printenv DATABASE_URL` with execOutput (or fails +// when empty) and `docker compose port ` with portScript. +func writeFakeDocker(t *testing.T, execOutput, portScript string) { + t.Helper() + + if runtime.GOOS == "windows" { + t.Skip("fake docker binary requires a POSIX shell") + } + + shPath, err := exec.LookPath("sh") + require.NoError(t, err) + + execBranch := "exit 1" + if execOutput != "" { + execBranch = fmt.Sprintf("echo %q", execOutput) + } + + script := fmt.Sprintf(`#!%s +if [ "$1" = "compose" ] && [ "$2" = "exec" ]; then + %s +elif [ "$1" = "compose" ] && [ "$2" = "port" ]; then + %s +else + exit 1 +fi +`, shPath, execBranch, portScript) + + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "docker"), []byte(script), 0o755)) + t.Setenv("PATH", dir) +} + +func TestDockerDatabaseConnection(t *testing.T) { + writeFakeDocker(t, "mysql://app:secret@database/shop", `echo "0.0.0.0:55001"`) + + dockerExec := &DockerExecutor{projectRoot: t.TempDir()} + + conn, err := dockerExec.DatabaseConnection(t.Context()) + require.NoError(t, err) + + assert.Equal(t, "127.0.0.1:55001", conn.Addr()) + assert.Equal(t, "app", conn.Username) + assert.Equal(t, "secret", conn.Password) + assert.Equal(t, "shop", conn.Database) +} + +func TestDockerDatabaseConnectionEnvOverrideSkipsContainerLookup(t *testing.T) { + // The exec branch fails, so passing proves the container env is not read. + writeFakeDocker(t, "", `echo "[::]:56001"`) + + dockerExec := &DockerExecutor{ + projectRoot: t.TempDir(), + env: map[string]string{"DATABASE_URL": "mysql://root:root@database:3306/override"}, + } + + conn, err := dockerExec.DatabaseConnection(t.Context()) + require.NoError(t, err) + + assert.Equal(t, "127.0.0.1", conn.Host) + assert.Equal(t, "56001", conn.Port) + assert.Equal(t, "override", conn.Database) +} + +func TestDockerDatabaseConnectionExternalHostKept(t *testing.T) { + writeFakeDocker(t, "mysql://app:pw@db.example.com:3307/prod", `echo "no such service: db.example.com" >&2; exit 1`) + + dockerExec := &DockerExecutor{projectRoot: t.TempDir()} + + conn, err := dockerExec.DatabaseConnection(t.Context()) + require.NoError(t, err) + + assert.Equal(t, "db.example.com:3307", conn.Addr()) + assert.Equal(t, "prod", conn.Database) +} + +func TestDockerDatabaseConnectionUnpublishedPort(t *testing.T) { + writeFakeDocker(t, "mysql://root:root@database/shopware", `echo ""`) + + dockerExec := &DockerExecutor{projectRoot: t.TempDir()} + + _, err := dockerExec.DatabaseConnection(t.Context()) + require.Error(t, err) + + assert.Contains(t, err.Error(), "does not publish port") +} + +func TestDockerDatabaseConnectionPortLookupFailure(t *testing.T) { + writeFakeDocker(t, "mysql://root:root@database/shopware", `echo "daemon not reachable" >&2; exit 1`) + + dockerExec := &DockerExecutor{projectRoot: t.TempDir()} + + _, err := dockerExec.DatabaseConnection(t.Context()) + require.Error(t, err) + + assert.Contains(t, err.Error(), "daemon not reachable") +} + +func TestDockerDatabaseConnectionEnvironmentNotRunning(t *testing.T) { + writeFakeDocker(t, "", "exit 1") + + dockerExec := &DockerExecutor{projectRoot: t.TempDir()} + + _, err := dockerExec.DatabaseConnection(t.Context()) + require.Error(t, err) + + assert.Contains(t, err.Error(), "could not read DATABASE_URL") +} diff --git a/internal/executor/database_test.go b/internal/executor/database_test.go new file mode 100644 index 00000000..02650885 --- /dev/null +++ b/internal/executor/database_test.go @@ -0,0 +1,198 @@ +package executor + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDatabaseConnectionDefaults(t *testing.T) { + t.Setenv("DATABASE_URL", "") + + conn, err := databaseConnectionFromEnv(t.TempDir(), nil) + require.NoError(t, err) + + assert.Equal(t, "127.0.0.1", conn.Host) + assert.Equal(t, "3306", conn.Port) + assert.Equal(t, "root", conn.Username) + assert.Equal(t, "root", conn.Password) + assert.Equal(t, "shopware", conn.Database) + assert.Equal(t, "127.0.0.1:3306", conn.Addr()) +} + +func TestDatabaseConnectionFromEnvFile(t *testing.T) { + t.Setenv("DATABASE_URL", "") + + projectRoot := t.TempDir() + writeFile(t, filepath.Join(projectRoot, ".env"), "DATABASE_URL=mysql://app:secret@db.example.com:3307/shop?sslmode=disable\n") + + conn, err := databaseConnectionFromEnv(projectRoot, nil) + require.NoError(t, err) + + assert.Equal(t, "db.example.com", conn.Host) + assert.Equal(t, "3307", conn.Port) + assert.Equal(t, "app", conn.Username) + assert.Equal(t, "secret", conn.Password) + assert.Equal(t, "shop", conn.Database) +} + +func TestDatabaseConnectionEnvLocalOverridesEnv(t *testing.T) { + t.Setenv("DATABASE_URL", "") + + projectRoot := t.TempDir() + writeFile(t, filepath.Join(projectRoot, ".env"), "DATABASE_URL=mysql://a:a@one/first\n") + writeFile(t, filepath.Join(projectRoot, ".env.local"), "DATABASE_URL=mysql://b:b@two/second\n") + + conn, err := databaseConnectionFromEnv(projectRoot, nil) + require.NoError(t, err) + + assert.Equal(t, "two", conn.Host) + assert.Equal(t, "second", conn.Database) +} + +func TestDatabaseConnectionRealEnvWinsOverFile(t *testing.T) { + t.Setenv("DATABASE_URL", "mysql://real:env@realhost/realdb") + + projectRoot := t.TempDir() + writeFile(t, filepath.Join(projectRoot, ".env"), "DATABASE_URL=mysql://file:file@filehost/filedb\n") + + conn, err := databaseConnectionFromEnv(projectRoot, nil) + require.NoError(t, err) + + assert.Equal(t, "realhost", conn.Host) + assert.Equal(t, "3306", conn.Port) + assert.Equal(t, "realdb", conn.Database) +} + +func TestDatabaseConnectionExecutorEnvWins(t *testing.T) { + t.Setenv("DATABASE_URL", "mysql://real:env@realhost/realdb") + + conn, err := databaseConnectionFromEnv(t.TempDir(), map[string]string{ + "DATABASE_URL": "mysql://extra:extra@extrahost/extradb", + }) + require.NoError(t, err) + + assert.Equal(t, "extrahost", conn.Host) + assert.Equal(t, "extradb", conn.Database) +} + +func TestApplyDatabaseURL(t *testing.T) { + cases := []struct { + name string + url string + expected DatabaseConnection + }{ + { + name: "full url", + url: "mysql://user:pass@host:3307/db", + expected: DatabaseConnection{Host: "host", Port: "3307", Username: "user", Password: "pass", Database: "db"}, + }, + { + name: "no port keeps default", + url: "mysql://user:pass@host/db", + expected: DatabaseConnection{Host: "host", Port: "3306", Username: "user", Password: "pass", Database: "db"}, + }, + { + name: "user without password clears default password", + url: "mysql://user@host/db", + expected: DatabaseConnection{Host: "host", Port: "3306", Username: "user", Password: "", Database: "db"}, + }, + { + name: "url encoded credentials", + url: "mysql://us%40er:p%40ss%2Fword@host/db", + expected: DatabaseConnection{Host: "host", Port: "3306", Username: "us@er", Password: "p@ss/word", Database: "db"}, + }, + { + name: "query parameters ignored", + url: "mysql://user:pass@host/db?serverVersion=8.0&charset=utf8mb4", + expected: DatabaseConnection{Host: "host", Port: "3306", Username: "user", Password: "pass", Database: "db"}, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + conn := defaultDatabaseConnection() + require.NoError(t, applyDatabaseURL(conn, tc.url)) + assert.Equal(t, tc.expected, *conn) + }) + } +} + +func TestApplyDatabaseURLInvalid(t *testing.T) { + cases := map[string]string{ + "space in password": "mysql://user:pa ss@host/db", + "bare word": "shopware", + "missing host": "mysql:///shopware", + "scheme-less address": "root:root@localhost/shopware", + } + + for name, databaseURL := range cases { + t.Run(name, func(t *testing.T) { + conn := defaultDatabaseConnection() + assert.Error(t, applyDatabaseURL(conn, databaseURL)) + }) + } +} + +func TestDatabaseConnectionMySQLConfig(t *testing.T) { + conn := &DatabaseConnection{Host: "db.internal", Port: "3307", Username: "app", Password: "secret", Database: "shop"} + + cfg := conn.MySQLConfig() + + assert.Equal(t, "tcp", cfg.Net) + assert.Equal(t, "db.internal:3307", cfg.Addr) + assert.Equal(t, "app", cfg.User) + assert.Equal(t, "secret", cfg.Passwd) + assert.Equal(t, "shop", cfg.DBName) + assert.Equal(t, 0, cfg.MaxAllowedPacket, "must fetch max_allowed_packet from the server") +} + +func TestDatabaseConnectionOpenUnreachable(t *testing.T) { + // Port 1 on loopback refuses connections immediately. + conn := &DatabaseConnection{Host: "127.0.0.1", Port: "1", Username: "root", Password: "root", Database: "shop"} + + _, _, err := conn.Open(t.Context()) + require.Error(t, err) + + assert.Contains(t, err.Error(), `could not connect to database "shop" at 127.0.0.1:1`) +} + +func TestDatabaseConnectionUnreadableEnvFile(t *testing.T) { + t.Setenv("DATABASE_URL", "") + + projectRoot := t.TempDir() + // A directory named .env makes the env file layer fail to read. + if err := os.Mkdir(filepath.Join(projectRoot, ".env"), 0o755); err != nil { + t.Fatal(err) + } + + _, err := databaseConnectionFromEnv(projectRoot, nil) + assert.Error(t, err) +} + +func TestLocalExecutorDatabaseConnection(t *testing.T) { + t.Setenv("DATABASE_URL", "") + + projectRoot := t.TempDir() + writeFile(t, filepath.Join(projectRoot, ".env"), "DATABASE_URL=mysql://shopware:shopware@127.0.0.1:13306/dev\n") + + exec := &LocalExecutor{projectRoot: projectRoot} + + conn, err := exec.DatabaseConnection(t.Context()) + require.NoError(t, err) + + assert.Equal(t, "127.0.0.1:13306", conn.Addr()) + assert.Equal(t, "shopware", conn.Username) + assert.Equal(t, "dev", conn.Database) +} + +func writeFile(t *testing.T, path, content string) { + t.Helper() + + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatal(err) + } +} diff --git a/internal/executor/docker.go b/internal/executor/docker.go index 326de01c..2e101c59 100644 --- a/internal/executor/docker.go +++ b/internal/executor/docker.go @@ -3,6 +3,7 @@ package executor import ( "context" "fmt" + "net" "os/exec" "path/filepath" "strings" @@ -104,6 +105,87 @@ func (d *DockerExecutor) AdminAPIClient(ctx context.Context) (*adminSdk.Client, return adminAPIClient(ctx, d.shopCfg, d.envCfg) } +// DatabaseConnection resolves the database credentials as seen inside the +// compose network and translates the service host to the port published on +// the host machine. +func (d *DockerExecutor) DatabaseConnection(ctx context.Context) (*DatabaseConnection, error) { + conn := defaultDatabaseConnection() + conn.Host = "database" + + databaseURL := d.env["DATABASE_URL"] + + if databaseURL == "" { + cmd := exec.CommandContext(ctx, "docker", "compose", "exec", "-T", "web", "printenv", "DATABASE_URL") + cmd.Dir = d.projectRoot + logCmd(ctx, cmd) + + var stdout, stderr strings.Builder + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + return nil, fmt.Errorf("could not read DATABASE_URL from the web container, is the environment running?: %w\n%s", err, stderr.String()) + } + + databaseURL = strings.TrimSpace(stdout.String()) + } + + if databaseURL != "" { + if err := applyDatabaseURL(conn, databaseURL); err != nil { + return nil, err + } + } + + // The host part of DATABASE_URL is only resolvable inside the compose + // network when it names a compose service. Swap it for the address the + // port is published on. + if err := d.resolvePublishedPort(ctx, conn); err != nil { + return nil, err + } + + return conn, nil +} + +// resolvePublishedPort rewrites conn's address to the host-published mapping +// of the compose service it points at. When the host is not a compose service +// (external database), the address is kept untouched. +func (d *DockerExecutor) resolvePublishedPort(ctx context.Context, conn *DatabaseConnection) error { + cmd := exec.CommandContext(ctx, "docker", "compose", "port", conn.Host, conn.Port) + cmd.Dir = d.projectRoot + logCmd(ctx, cmd) + + var stdout, stderr strings.Builder + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + if strings.Contains(stderr.String(), "no such service") { + return nil + } + + return fmt.Errorf("could not resolve published port of service %q: %w\n%s", conn.Host, err, stderr.String()) + } + + published := strings.TrimSpace(stdout.String()) + if line, _, found := strings.Cut(published, "\n"); found { + published = strings.TrimSpace(line) + } + + host, port, err := net.SplitHostPort(published) + if err != nil || port == "0" { + return fmt.Errorf("service %q does not publish port %s to the host, regenerate the compose file by restarting the environment (shopware-cli project dev)", conn.Host, conn.Port) + } + + if host == "0.0.0.0" || host == "::" || host == "" { + host = "127.0.0.1" + } + + conn.Host = host + conn.Port = port + + return nil +} + func (d *DockerExecutor) containerWorkdir() string { if d.relDir == "" { return "/var/www/html" diff --git a/internal/executor/executor.go b/internal/executor/executor.go index b9bbb80f..3476005a 100644 --- a/internal/executor/executor.go +++ b/internal/executor/executor.go @@ -35,6 +35,9 @@ type Executor interface { StopEnvironment(ctx context.Context) error EnvironmentStatus(ctx context.Context) (bool, error) AdminAPIClient(ctx context.Context) (*adminSdk.Client, error) + // DatabaseConnection returns credentials to reach the project database + // from the host machine. + DatabaseConnection(ctx context.Context) (*DatabaseConnection, error) } func adminAPIClient(ctx context.Context, cfg *shop.Config, envCfg *shop.EnvironmentConfig) (*adminSdk.Client, error) { diff --git a/internal/executor/local.go b/internal/executor/local.go index 076c8f7c..f6305d82 100644 --- a/internal/executor/local.go +++ b/internal/executor/local.go @@ -139,6 +139,10 @@ func (l *LocalExecutor) AdminAPIClient(ctx context.Context) (*adminSdk.Client, e return adminAPIClient(ctx, l.shopCfg, l.envCfg) } +func (l *LocalExecutor) DatabaseConnection(_ context.Context) (*DatabaseConnection, error) { + return databaseConnectionFromEnv(l.projectRoot, l.env) +} + func (l *LocalExecutor) StartEnvironment(_ context.Context) error { return ErrNotSupported } diff --git a/internal/executor/symfony_cli.go b/internal/executor/symfony_cli.go index 39f12316..3a2fbff6 100644 --- a/internal/executor/symfony_cli.go +++ b/internal/executor/symfony_cli.go @@ -75,6 +75,10 @@ func (s *SymfonyCLIExecutor) AdminAPIClient(ctx context.Context) (*adminSdk.Clie return adminAPIClient(ctx, s.shopCfg, s.envCfg) } +func (s *SymfonyCLIExecutor) DatabaseConnection(_ context.Context) (*DatabaseConnection, error) { + return databaseConnectionFromEnv(s.projectRoot, s.env) +} + func (s *SymfonyCLIExecutor) StartEnvironment(_ context.Context) error { return ErrNotSupported } diff --git a/internal/shop/pluginmigrate/pluginmigrate_test.go b/internal/shop/pluginmigrate/pluginmigrate_test.go index b2a7167c..339d70a5 100644 --- a/internal/shop/pluginmigrate/pluginmigrate_test.go +++ b/internal/shop/pluginmigrate/pluginmigrate_test.go @@ -46,11 +46,14 @@ func (f *fakeExecutor) NPMCommand(ctx context.Context, args ...string) *executor return shellProcess(ctx, "true") } -func (f *fakeExecutor) NormalizePath(hostPath string) string { return hostPath } -func (f *fakeExecutor) Type() string { return executor.TypeLocal } -func (f *fakeExecutor) WithEnv(map[string]string) executor.Executor { return f } -func (f *fakeExecutor) WithRelDir(string) executor.Executor { return f } -func (f *fakeExecutor) StartEnvironment(context.Context) error { return nil } +func (f *fakeExecutor) NormalizePath(hostPath string) string { return hostPath } +func (f *fakeExecutor) Type() string { return executor.TypeLocal } +func (f *fakeExecutor) WithEnv(map[string]string) executor.Executor { return f } +func (f *fakeExecutor) WithRelDir(string) executor.Executor { return f } +func (f *fakeExecutor) StartEnvironment(context.Context) error { return nil } +func (f *fakeExecutor) DatabaseConnection(context.Context) (*executor.DatabaseConnection, error) { + return nil, executor.ErrNotSupported +} func (f *fakeExecutor) StopEnvironment(context.Context) error { return nil } func (f *fakeExecutor) EnvironmentStatus(context.Context) (bool, error) { return true, nil } func (f *fakeExecutor) AdminAPIClient(context.Context) (*adminSdk.Client, error) { diff --git a/internal/shop/upgrade/run_test.go b/internal/shop/upgrade/run_test.go index 5aa7fcc1..8c569769 100644 --- a/internal/shop/upgrade/run_test.go +++ b/internal/shop/upgrade/run_test.go @@ -51,8 +51,12 @@ func (f *fakeExecutor) WithEnv(env map[string]string) executor.Executor { f.env = env return f } -func (f *fakeExecutor) WithRelDir(string) executor.Executor { return f } -func (f *fakeExecutor) StartEnvironment(context.Context) error { return nil } +func (f *fakeExecutor) WithRelDir(string) executor.Executor { return f } +func (f *fakeExecutor) StartEnvironment(context.Context) error { return nil } + +func (f *fakeExecutor) DatabaseConnection(context.Context) (*executor.DatabaseConnection, error) { + return nil, executor.ErrNotSupported +} func (f *fakeExecutor) StopEnvironment(context.Context) error { return nil } func (f *fakeExecutor) EnvironmentStatus(context.Context) (bool, error) { return true, nil } func (f *fakeExecutor) AdminAPIClient(context.Context) (*adminSdk.Client, error) { diff --git a/internal/sqlshell/fakedb_test.go b/internal/sqlshell/fakedb_test.go new file mode 100644 index 00000000..5be44981 --- /dev/null +++ b/internal/sqlshell/fakedb_test.go @@ -0,0 +1,112 @@ +package sqlshell + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "io" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +// The fake driver serves canned result sets keyed by the exact query text +// and lets the rendering pipeline run on real *sql.Rows without a server. +// Queries containing "boom" fail, exec statements containing "one_row" +// report one affected row, all others three. + +type fakeTable struct { + cols []string + types []string + rows [][]driver.Value +} + +var fakeQueries sync.Map // query string -> fakeTable + +type fakeDriver struct{} + +func (fakeDriver) Open(string) (driver.Conn, error) { return &fakeDBConn{}, nil } + +type fakeDBConn struct{} + +func (*fakeDBConn) Prepare(string) (driver.Stmt, error) { + return nil, errors.New("prepare not supported") +} + +func (*fakeDBConn) Close() error { return nil } + +func (*fakeDBConn) Begin() (driver.Tx, error) { return nil, errors.New("tx not supported") } + +func (*fakeDBConn) QueryContext(ctx context.Context, query string, _ []driver.NamedValue) (driver.Rows, error) { + if strings.Contains(query, "boom") { + return nil, errors.New("query exploded") + } + + // Simulates a long-running query: blocks until the context is cancelled. + if strings.Contains(query, "block") { + <-ctx.Done() + return nil, ctx.Err() + } + + table, ok := fakeQueries.Load(query) + if !ok { + return nil, fmt.Errorf("unexpected query: %s", query) + } + + return &fakeDriverRows{table: table.(fakeTable)}, nil +} + +func (*fakeDBConn) ExecContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Result, error) { + if strings.Contains(query, "boom") { + return nil, errors.New("exec exploded") + } + + if strings.Contains(query, "one_row") { + return driver.RowsAffected(1), nil + } + + return driver.RowsAffected(3), nil +} + +type fakeDriverRows struct { + table fakeTable + pos int +} + +func (r *fakeDriverRows) Columns() []string { return r.table.cols } + +func (r *fakeDriverRows) Close() error { return nil } + +func (r *fakeDriverRows) Next(dest []driver.Value) error { + if r.pos >= len(r.table.rows) { + return io.EOF + } + + copy(dest, r.table.rows[r.pos]) + r.pos++ + + return nil +} + +func (r *fakeDriverRows) ColumnTypeDatabaseTypeName(index int) string { + return r.table.types[index] +} + +var registerFakeDriver = sync.OnceFunc(func() { + sql.Register("sqlshell-fake", fakeDriver{}) +}) + +func openFakeDB(t *testing.T) *sql.DB { + t.Helper() + registerFakeDriver() + + db, err := sql.Open("sqlshell-fake", "") + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + return db +} diff --git a/internal/sqlshell/interactive.go b/internal/sqlshell/interactive.go new file mode 100644 index 00000000..b20c5747 --- /dev/null +++ b/internal/sqlshell/interactive.go @@ -0,0 +1,284 @@ +package sqlshell + +import ( + "bytes" + "context" + "errors" + "fmt" + "strings" + + "charm.land/bubbles/v2/textinput" + tea "charm.land/bubbletea/v2" +) + +// InteractiveShell runs the read-eval-print loop as an inline bubbletea +// program, providing real line editing (word deletion via alt+backspace, +// ctrl+w and alt+d, cursor movement) and session history on the arrow keys. +// Results are printed into the normal terminal scrollback. +func InteractiveShell(ctx context.Context, db Conn, format Format) error { + _, err := tea.NewProgram(newInteractiveModel(ctx, db, format), tea.WithContext(ctx)).Run() + + return filterInteractiveErr(err) +} + +// filterInteractiveErr treats context cancellation as a normal shell exit; +// bubbletea wraps it in tea.ErrProgramKilled. +func filterInteractiveErr(err error) error { + if errors.Is(err, context.Canceled) { + return nil + } + + return err +} + +func newInteractiveModel(ctx context.Context, db Conn, format Format) *interactiveModel { + input := textinput.New() + input.Prompt = promptMain + input.Focus() + + return &interactiveModel{ + newExecution: func(statements []string, quit bool) (context.CancelFunc, tea.Cmd) { + runCtx, cancel := context.WithCancel(ctx) + + return cancel, func() tea.Msg { + var out bytes.Buffer + + for _, stmt := range statements { + err := RunStatement(runCtx, db, stmt, &out, format) + + if runCtx.Err() != nil { + // Cancelling closes the driver connection mid-query, + // the session connection may not survive it. + _, _ = fmt.Fprintln(&out, "Query cancelled. Restart the shell if further statements fail.") + break + } + + if err != nil { + _, _ = fmt.Fprintf(&out, "ERROR: %s\n", err) + } + } + + return resultMsg{output: strings.TrimRight(out.String(), "\n"), quit: quit} + } + }, + input: input, + delimiter: DefaultDelimiter, + } +} + +// resultMsg carries the rendered output of executed statements back into the +// update loop. +type resultMsg struct { + output string + quit bool +} + +type interactiveModel struct { + // newExecution prepares one asynchronous statement run against the + // session connection; it carries the command context in its closure and + // returns the cancel function for that run. + newExecution func(statements []string, quit bool) (context.CancelFunc, tea.Cmd) + + input textinput.Model + buffer string + delimiter string + + history []string + histPos int + stashed string + + running bool + cancelling bool + cancelRun context.CancelFunc + done bool +} + +func (m *interactiveModel) Init() tea.Cmd { + return textinput.Blink +} + +func (m *interactiveModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.input.SetWidth(max(msg.Width-len(promptContinuation)-1, 20)) + return m, nil + + case resultMsg: + m.running = false + m.cancelling = false + + if m.cancelRun != nil { + m.cancelRun() + m.cancelRun = nil + } + + var cmds []tea.Cmd + if msg.output != "" { + cmds = append(cmds, tea.Println(msg.output)) + } + + if msg.quit { + m.done = true + cmds = append(cmds, tea.Quit) + } + + return m, tea.Sequence(cmds...) + + case tea.KeyPressMsg: + if m.running { + // The terminal is in raw mode, ctrl+c raises no SIGINT: cancel + // the running statement here instead. + if msg.String() == "ctrl+c" && m.cancelRun != nil && !m.cancelling { + m.cancelling = true + m.cancelRun() + } + + return m, nil + } + + return m.handleKey(msg) + } + + return m, nil +} + +func (m *interactiveModel) handleKey(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + switch msg.String() { + case "enter": + return m.submitLine() + + case "up": + if m.histPos > 0 { + if m.histPos == len(m.history) { + m.stashed = m.input.Value() + } + m.histPos-- + m.input.SetValue(m.history[m.histPos]) + m.input.CursorEnd() + } + return m, nil + + case "down": + if m.histPos < len(m.history) { + m.histPos++ + if m.histPos == len(m.history) { + m.input.SetValue(m.stashed) + } else { + m.input.SetValue(m.history[m.histPos]) + } + m.input.CursorEnd() + } + return m, nil + + case "ctrl+c": + // Drop the pending input like the mysql client; quit when idle. + if m.input.Value() == "" && m.buffer == "" { + m.done = true + return m, tea.Quit + } + + echo := tea.Println(m.prompt() + m.input.Value() + "^C") + m.reset() + return m, echo + + case "ctrl+d": + // End of input on an empty line: run what is pending, then quit. + if m.input.Value() == "" { + return m.finish() + } + } + + var cmd tea.Cmd + m.input, cmd = m.input.Update(msg) + + return m, cmd +} + +// submitLine feeds the entered line into the statement buffer and executes +// every completed statement. +func (m *interactiveModel) submitLine() (tea.Model, tea.Cmd) { + line := m.input.Value() + echo := tea.Println(m.prompt() + line) + + if strings.TrimSpace(line) != "" { + m.history = append(m.history, line) + } + m.histPos = len(m.history) + m.stashed = "" + m.input.SetValue("") + + if m.buffer == "" && isExitCommand(line) { + m.done = true + return m, tea.Sequence(echo, tea.Quit) + } + + m.buffer += line + "\n" + + var statements []string + statements, m.buffer, m.delimiter = SplitStatementsWithDelimiter(m.buffer, m.delimiter) + + m.input.Prompt = m.prompt() + + if len(statements) == 0 { + return m, echo + } + + return m, tea.Sequence(echo, m.startExecution(statements, false)) +} + +// finish runs a trailing statement without semicolon (matching Run and +// ExecuteStream) and quits. +func (m *interactiveModel) finish() (tea.Model, tea.Cmd) { + if rest := strings.TrimSpace(m.buffer); rest != "" { + m.buffer = "" + + return m, m.startExecution([]string{rest}, true) + } + + m.done = true + + return m, tea.Quit +} + +// startExecution kicks off an asynchronous, cancellable statement run. +// Statement errors are printed, they do not end the session. +func (m *interactiveModel) startExecution(statements []string, quit bool) tea.Cmd { + cancel, cmd := m.newExecution(statements, quit) + + m.cancelRun = cancel + m.running = true + + return cmd +} + +func (m *interactiveModel) prompt() string { + if m.buffer != "" { + return promptContinuation + } + + return promptMain +} + +func (m *interactiveModel) reset() { + m.buffer = "" + m.histPos = len(m.history) + m.stashed = "" + m.input.SetValue("") + m.input.Prompt = promptMain +} + +func (m *interactiveModel) View() tea.View { + if m.done { + return tea.NewView("") + } + + if m.running { + if m.cancelling { + return tea.NewView("Cancelling query...") + } + + return tea.NewView("Executing... press ctrl+c to cancel") + } + + return tea.NewView(m.input.View()) +} diff --git a/internal/sqlshell/interactive_test.go b/internal/sqlshell/interactive_test.go new file mode 100644 index 00000000..ee8f3cfe --- /dev/null +++ b/internal/sqlshell/interactive_test.go @@ -0,0 +1,215 @@ +package sqlshell + +import ( + "context" + "database/sql/driver" + "errors" + "fmt" + "strings" + "testing" + + tea "charm.land/bubbletea/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTestInteractiveModel(t *testing.T) *interactiveModel { + t.Helper() + + return newInteractiveModel(t.Context(), openFakeDB(t), FormatTSV) +} + +func typeText(m *interactiveModel, text string) { + for _, r := range text { + m.Update(tea.KeyPressMsg(tea.Key{Code: r, Text: string(r)})) + } +} + +func pressKey(m *interactiveModel, key tea.Key) tea.Cmd { + _, cmd := m.Update(tea.KeyPressMsg(key)) + return cmd +} + +func TestInteractiveDeleteWordBackward(t *testing.T) { + m := newTestInteractiveModel(t) + + typeText(m, "SELECT id FROM product") + + pressKey(m, tea.Key{Code: tea.KeyBackspace, Mod: tea.ModAlt}) + assert.Equal(t, "SELECT id FROM ", m.input.Value()) + + pressKey(m, tea.Key{Code: 'w', Mod: tea.ModCtrl}) + assert.Equal(t, "SELECT id ", m.input.Value()) +} + +func TestInteractiveDeleteWordForward(t *testing.T) { + m := newTestInteractiveModel(t) + + typeText(m, "SELECT id FROM product") + m.input.CursorStart() + + pressKey(m, tea.Key{Code: 'd', Mod: tea.ModAlt}) + assert.Equal(t, " id FROM product", m.input.Value()) +} + +func TestInteractiveHistoryNavigation(t *testing.T) { + m := newTestInteractiveModel(t) + + typeText(m, "SELECT 1;") + pressKey(m, tea.Key{Code: tea.KeyEnter}) + m.Update(resultMsg{output: ""}) + + typeText(m, "SEL") + pressKey(m, tea.Key{Code: tea.KeyUp}) + assert.Equal(t, "SELECT 1;", m.input.Value(), "up recalls the previous line") + + pressKey(m, tea.Key{Code: tea.KeyDown}) + assert.Equal(t, "SEL", m.input.Value(), "down restores the stashed live input") +} + +func TestInteractiveContinuationAndExecution(t *testing.T) { + m := newTestInteractiveModel(t) + + typeText(m, "SELECT 1") + pressKey(m, tea.Key{Code: tea.KeyEnter}) + + assert.False(t, m.running, "incomplete statement must not execute") + assert.Equal(t, promptContinuation, m.prompt()) + + typeText(m, ";") + pressKey(m, tea.Key{Code: tea.KeyEnter}) + + assert.True(t, m.running) + assert.Empty(t, m.buffer) +} + +func TestInteractiveExitCommand(t *testing.T) { + m := newTestInteractiveModel(t) + + typeText(m, "exit") + pressKey(m, tea.Key{Code: tea.KeyEnter}) + + assert.True(t, m.done) +} + +func TestInteractiveCtrlCClearsPendingInput(t *testing.T) { + m := newTestInteractiveModel(t) + + typeText(m, "SELECT 1") + pressKey(m, tea.Key{Code: tea.KeyEnter}) + typeText(m, "FROM t") + + pressKey(m, tea.Key{Code: 'c', Mod: tea.ModCtrl}) + + assert.False(t, m.done) + assert.Empty(t, m.buffer) + assert.Empty(t, m.input.Value()) + assert.Equal(t, promptMain, m.prompt()) +} + +func TestInteractiveCtrlCQuitsWhenIdle(t *testing.T) { + m := newTestInteractiveModel(t) + + pressKey(m, tea.Key{Code: 'c', Mod: tea.ModCtrl}) + + assert.True(t, m.done) +} + +func TestInteractiveCtrlDRunsTrailingStatement(t *testing.T) { + m := newTestInteractiveModel(t) + + typeText(m, "UPDATE t SET one_row = 1") + pressKey(m, tea.Key{Code: tea.KeyEnter}) + + pressKey(m, tea.Key{Code: 'd', Mod: tea.ModCtrl}) + + assert.True(t, m.running, "pending statement must run on EOF") + assert.Empty(t, m.buffer) +} + +func TestInteractiveExecuteRendersResults(t *testing.T) { + fakeQueries.Store("SELECT 3", fakeTable{ + cols: []string{"3"}, + types: []string{"BIGINT"}, + rows: [][]driver.Value{{[]byte("3")}}, + }) + + m := newTestInteractiveModel(t) + + cmd := m.startExecution([]string{"SELECT 3", "SELECT boom"}, true) + require.True(t, m.running) + + result, ok := cmd().(resultMsg) + require.True(t, ok) + + assert.True(t, result.quit) + assert.Contains(t, result.output, "3\n3") + assert.Contains(t, result.output, "ERROR: query exploded") +} + +func TestInteractiveCtrlCCancelsRunningQuery(t *testing.T) { + m := newTestInteractiveModel(t) + + cmd := m.startExecution([]string{"SELECT block"}, false) + require.True(t, m.running) + + // The statement blocks until its context is cancelled. + results := make(chan tea.Msg, 1) + go func() { results <- cmd() }() + + pressKey(m, tea.Key{Code: 'c', Mod: tea.ModCtrl}) + assert.True(t, m.cancelling) + + result, ok := (<-results).(resultMsg) + require.True(t, ok) + assert.Contains(t, result.output, "Query cancelled") + + m.Update(result) + assert.False(t, m.running) + assert.False(t, m.cancelling) + assert.Nil(t, m.cancelRun) +} + +func TestFilterInteractiveErr(t *testing.T) { + assert.NoError(t, filterInteractiveErr(nil)) + assert.NoError(t, filterInteractiveErr(fmt.Errorf("%w: %w", tea.ErrProgramKilled, context.Canceled))) + + other := errors.New("terminal broke") + assert.Equal(t, other, filterInteractiveErr(other)) +} + +func TestInteractiveIgnoresKeysWhileRunning(t *testing.T) { + m := newTestInteractiveModel(t) + m.running = true + + typeText(m, "SELECT 1") + + assert.Empty(t, m.input.Value()) +} + +func TestInteractiveResultQuits(t *testing.T) { + m := newTestInteractiveModel(t) + m.running = true + + _, cmd := m.Update(resultMsg{output: "done", quit: true}) + + assert.False(t, m.running) + assert.True(t, m.done) + require.NotNil(t, cmd) +} + +func TestInteractiveViewShowsRunningState(t *testing.T) { + m := newTestInteractiveModel(t) + + assert.Contains(t, m.View().Content, promptMain) + + m.running = true + assert.Contains(t, m.View().Content, "Executing") + assert.Contains(t, m.View().Content, "ctrl+c") + + m.cancelling = true + assert.Contains(t, m.View().Content, "Cancelling") + + m.done = true + assert.Equal(t, "", strings.TrimSpace(m.View().Content)) +} diff --git a/internal/sqlshell/run.go b/internal/sqlshell/run.go new file mode 100644 index 00000000..f8db49c3 --- /dev/null +++ b/internal/sqlshell/run.go @@ -0,0 +1,489 @@ +package sqlshell + +import ( + "context" + "database/sql" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "strings" + "time" + "unicode/utf8" +) + +// Execer executes statements that do not return rows. +type Execer interface { + ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) +} + +// Conn is the database handle the shell operates on. It is satisfied by +// *sql.Conn (preferred, since session state like SET or USE sticks to a +// single connection) and *sql.DB. +type Conn interface { + Execer + QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) +} + +// Format controls how result sets are rendered. +type Format string + +const ( + // FormatTable renders mysql-client style ASCII tables. + FormatTable Format = "table" + // FormatTSV renders tab-separated values with a header line. + FormatTSV Format = "tsv" + // FormatJSON renders each result set as a JSON array of row objects. + FormatJSON Format = "json" +) + +// ParseFormat validates a format name given on the command line. +func ParseFormat(name string) (Format, error) { + switch Format(name) { + case FormatTable, FormatTSV, FormatJSON: + return Format(name), nil + } + + return "", fmt.Errorf("unknown format %q, allowed values: table, tsv, json", name) +} + +// Run executes all statements in input against db and renders the results to out. +func Run(ctx context.Context, db Conn, input string, out io.Writer, format Format) error { + statements, rest, _ := SplitStatementsWithDelimiter(input, DefaultDelimiter) + if rest != "" { + // The trailing semicolon is optional. + statements = append(statements, rest) + } + + multiple := len(statements) > 1 + + for _, stmt := range statements { + if err := RunStatement(ctx, db, stmt, out, format); err != nil { + if multiple { + return fmt.Errorf("%q: %w", summarizeStatement(stmt), err) + } + + return err + } + } + + return nil +} + +// RunStatement executes a single statement and renders its result to out. +// Comment-only statements and stray DELIMITER directives (client commands, +// not SQL) are skipped. +func RunStatement(ctx context.Context, db Conn, stmt string, out io.Writer, format Format) error { + keyword := firstKeyword(stmt) + if keyword == "" || keyword == "delimiter" { + return nil + } + + start := time.Now() + + if !returnsResultSet(keyword) && !hasReturningClause(keyword, stmt) { + result, err := db.ExecContext(ctx, stmt) + if err != nil { + return err + } + + affected, err := result.RowsAffected() + if err != nil { + affected = 0 + } + + if format == FormatJSON { + _, err := fmt.Fprintf(out, "{\"rows_affected\":%d}\n", affected) + return err + } + + _, err = fmt.Fprintf(out, "Query OK, %d %s affected (%s)\n", affected, pluralRows(affected), formatDuration(time.Since(start))) + + return err + } + + rows, err := db.QueryContext(ctx, stmt) + if err != nil { + return err + } + // Read errors surface through rows.Err below, the close error is redundant. + defer func() { _ = rows.Close() }() + + // A statement like CALL can produce more than one result set. + for { + if err := renderResultSet(rows, time.Since(start), out, format); err != nil { + return err + } + + if !rows.NextResultSet() { + break + } + } + + return rows.Err() +} + +// firstKeyword returns the lowercased first word of the statement, skipping +// leading comments and parentheses. MySQL executable comments (/*!40014 ...) +// count as statement content. It returns "" for comment-only input. +func firstKeyword(stmt string) string { + for { + stmt = strings.TrimLeft(stmt, " \t\r\n(") + + switch { + case strings.HasPrefix(stmt, "/*!"): + // Executable comment: the server runs its content, so the + // keyword follows the marker and optional version number. + stmt = strings.TrimLeft(stmt[3:], "0123456789") + case strings.HasPrefix(stmt, "/*"): + _, after, found := strings.Cut(stmt[2:], "*/") + if !found { + return "" + } + stmt = after + case strings.HasPrefix(stmt, "#"), strings.HasPrefix(stmt, "--"): + _, after, found := strings.Cut(stmt, "\n") + if !found { + return "" + } + stmt = after + default: + end := strings.IndexFunc(stmt, func(r rune) bool { + return r == ' ' || r == '\t' || r == '\r' || r == '\n' || r == '(' || r == ';' + }) + if end == -1 { + return strings.ToLower(stmt) + } + + return strings.ToLower(stmt[:end]) + } + } +} + +// returnsResultSet reports whether a statement starting with keyword is +// expected to return rows and should go through Query instead of Exec. +func returnsResultSet(keyword string) bool { + switch keyword { + case "select", "show", "describe", "desc", "explain", "with", "values", "table", "call", "check", "checksum", "analyze", "optimize", "repair", "help": + return true + } + + return false +} + +// hasReturningClause reports whether a DML statement carries a MariaDB +// RETURNING clause and therefore produces a result set. The keyword scan +// ignores quoted sections and comments, follows MariaDB's rule that "--" is +// only a comment when followed by whitespace, and treats completed comments +// as token separators. +func hasReturningClause(keyword, stmt string) bool { + switch keyword { + case "insert", "replace", "delete": + default: + return false + } + + const returning = "RETURNING" + + state := stateNormal + // boundary tracks whether the previous position separates tokens, so the + // keyword is not matched inside identifiers like returning_col. + boundary := true + + for i := 0; i < len(stmt); i++ { + if state != stateNormal { + before := state + state, i = consumeQuoteOrComment(state, stmt, i) + + if state == stateNormal { + // A closed comment separates tokens, a closed literal does not. + boundary = before == stateLineComment || before == stateBlockComment + } + + continue + } + + c := stmt[i] + + switch { + case c == '\'': + state = stateSingleQuote + boundary = false + case c == '"': + state = stateDoubleQuote + boundary = false + case c == '`': + state = stateBacktick + boundary = false + case c == '#': + state = stateLineComment + case c == '-' && strings.HasPrefix(stmt[i:], "--") && (i+2 >= len(stmt) || isSQLSpace(stmt[i+2])): + state = stateLineComment + i++ + case c == '/' && strings.HasPrefix(stmt[i:], "/*"): + state = stateBlockComment + i++ + case isSQLSpace(c) || c == '(' || c == ')' || c == ',': + boundary = true + case (c == 'r' || c == 'R') && boundary: + if i+len(returning) <= len(stmt) && strings.EqualFold(stmt[i:i+len(returning)], returning) && boundaryAfterKeyword(stmt, i+len(returning)) { + return true + } + + boundary = false + default: + boundary = false + } + } + + return false +} + +func isSQLSpace(c byte) bool { + return c == ' ' || c == '\t' || c == '\r' || c == '\n' +} + +// boundaryAfterKeyword reports whether the keyword ending right before pos is +// followed by a token separator: end of input, whitespace, punctuation or a +// block comment. +func boundaryAfterKeyword(stmt string, pos int) bool { + if pos >= len(stmt) { + return true + } + + if isSQLSpace(stmt[pos]) || stmt[pos] == '(' || stmt[pos] == '*' { + return true + } + + return strings.HasPrefix(stmt[pos:], "/*") +} + +func summarizeStatement(stmt string) string { + stmt = strings.Join(strings.Fields(stmt), " ") + + const maxLen = 60 + if utf8.RuneCountInString(stmt) > maxLen { + return string([]rune(stmt)[:maxLen]) + "…" + } + + return stmt +} + +func renderResultSet(rows *sql.Rows, elapsed time.Duration, out io.Writer, format Format) error { + columns, err := rows.Columns() + if err != nil { + return err + } + + // Non-query statements executed through Query (e.g. the OK packet at the + // end of a CALL) produce a result set without columns. + if len(columns) == 0 { + return nil + } + + columnTypes, err := rows.ColumnTypes() + if err != nil { + return err + } + + binary := make([]bool, len(columnTypes)) + numeric := make([]bool, len(columnTypes)) + for i, columnType := range columnTypes { + binary[i] = isBinaryType(columnType.DatabaseTypeName()) + numeric[i] = isNumericType(columnType.DatabaseTypeName()) + } + + var data [][]any + + values := make([]sql.RawBytes, len(columns)) + pointers := make([]any, len(columns)) + for i := range values { + pointers[i] = &values[i] + } + + for rows.Next() { + if err := rows.Scan(pointers...); err != nil { + return err + } + + row := make([]any, len(columns)) + for i, value := range values { + switch { + case value == nil: + row[i] = nil + case binary[i]: + row[i] = "0x" + strings.ToUpper(hex.EncodeToString(value)) + default: + row[i] = string(value) + } + } + + data = append(data, row) + } + + if err := rows.Err(); err != nil { + return err + } + + if format == FormatJSON { + return renderJSON(columns, numeric, data, out) + } + + if format == FormatTSV { + return renderTSV(columns, data, out) + } + + return renderTable(columns, data, elapsed, out) +} + +func isBinaryType(databaseType string) bool { + switch strings.ToUpper(databaseType) { + case "BINARY", "VARBINARY", "BLOB", "TINYBLOB", "MEDIUMBLOB", "LONGBLOB", "BIT", "GEOMETRY": + return true + } + + return false +} + +func isNumericType(databaseType string) bool { + switch strings.ToUpper(databaseType) { + case "TINYINT", "SMALLINT", "MEDIUMINT", "INT", "BIGINT", "UNSIGNED TINYINT", "UNSIGNED SMALLINT", "UNSIGNED MEDIUMINT", "UNSIGNED INT", "UNSIGNED BIGINT", "DECIMAL", "FLOAT", "DOUBLE", "YEAR": + return true + } + + return false +} + +func renderTable(columns []string, data [][]any, elapsed time.Duration, out io.Writer) error { + widths := make([]int, len(columns)) + for i, column := range columns { + widths[i] = utf8.RuneCountInString(column) + } + + cells := make([][]string, len(data)) + for r, row := range data { + cells[r] = make([]string, len(columns)) + for i, value := range row { + cell := "NULL" + if value != nil { + cell = value.(string) + } + + cells[r][i] = cell + if width := utf8.RuneCountInString(cell); width > widths[i] { + widths[i] = width + } + } + } + + var builder strings.Builder + + writeSeparator := func() { + for _, width := range widths { + builder.WriteString("+") + builder.WriteString(strings.Repeat("-", width+2)) + } + builder.WriteString("+\n") + } + + writeRow := func(row []string) { + for i, cell := range row { + builder.WriteString("| ") + builder.WriteString(cell) + builder.WriteString(strings.Repeat(" ", widths[i]-utf8.RuneCountInString(cell)+1)) + } + builder.WriteString("|\n") + } + + if len(data) == 0 { + _, err := fmt.Fprintf(out, "Empty set (%s)\n", formatDuration(elapsed)) + return err + } + + writeSeparator() + writeRow(columns) + writeSeparator() + for _, row := range cells { + writeRow(row) + } + writeSeparator() + + if _, err := io.WriteString(out, builder.String()); err != nil { + return err + } + + _, err := fmt.Fprintf(out, "%d %s in set (%s)\n", len(data), pluralRows(int64(len(data))), formatDuration(elapsed)) + + return err +} + +func pluralRows(count int64) string { + if count == 1 { + return "row" + } + + return "rows" +} + +func renderTSV(columns []string, data [][]any, out io.Writer) error { + escape := strings.NewReplacer("\\", "\\\\", "\t", "\\t", "\n", "\\n", "\r", "\\r") + + writeLine := func(cells []string) error { + _, err := fmt.Fprintln(out, strings.Join(cells, "\t")) + return err + } + + header := make([]string, len(columns)) + for i, column := range columns { + header[i] = escape.Replace(column) + } + + if err := writeLine(header); err != nil { + return err + } + + for _, row := range data { + cells := make([]string, len(row)) + for i, value := range row { + if value == nil { + cells[i] = "NULL" + continue + } + + cells[i] = escape.Replace(value.(string)) + } + + if err := writeLine(cells); err != nil { + return err + } + } + + return nil +} + +func renderJSON(columns []string, numeric []bool, data [][]any, out io.Writer) error { + result := make([]map[string]any, len(data)) + + for r, row := range data { + object := make(map[string]any, len(columns)) + for i, column := range columns { + value := row[i] + + if str, ok := value.(string); ok && numeric[i] && json.Valid([]byte(str)) { + value = json.RawMessage(str) + } + + object[column] = value + } + + result[r] = object + } + + encoder := json.NewEncoder(out) + + return encoder.Encode(result) +} + +func formatDuration(elapsed time.Duration) string { + return fmt.Sprintf("%.3f sec", elapsed.Seconds()) +} diff --git a/internal/sqlshell/run_db_test.go b/internal/sqlshell/run_db_test.go new file mode 100644 index 00000000..33f9a1f0 --- /dev/null +++ b/internal/sqlshell/run_db_test.go @@ -0,0 +1,182 @@ +package sqlshell + +import ( + "context" + "database/sql/driver" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRunQueryAndExec(t *testing.T) { + fakeQueries.Store("SELECT id, name FROM tax", fakeTable{ + cols: []string{"id", "name"}, + types: []string{"BIGINT", "VARCHAR"}, + rows: [][]driver.Value{ + {[]byte("1"), []byte("Standard")}, + {[]byte("2"), nil}, + }, + }) + + var out strings.Builder + + err := Run(t.Context(), openFakeDB(t), "SELECT id, name FROM tax; UPDATE tax SET name = 'x';", &out, FormatTSV) + require.NoError(t, err) + + assert.Contains(t, out.String(), "id\tname\n1\tStandard\n2\tNULL\n") + assert.Contains(t, out.String(), "Query OK, 3 rows affected") +} + +func TestRunTrailingStatementWithoutSemicolon(t *testing.T) { + var out strings.Builder + + err := Run(t.Context(), openFakeDB(t), "UPDATE t SET one_row = 1", &out, FormatTSV) + require.NoError(t, err) + + assert.Contains(t, out.String(), "Query OK, 1 row affected") +} + +func TestRunTableFormat(t *testing.T) { + fakeQueries.Store("SELECT v FROM t", fakeTable{ + cols: []string{"v"}, + types: []string{"VARCHAR"}, + rows: [][]driver.Value{{[]byte("x")}}, + }) + + var out strings.Builder + + require.NoError(t, Run(t.Context(), openFakeDB(t), "SELECT v FROM t", &out, FormatTable)) + + assert.Contains(t, out.String(), "| v |") + assert.Contains(t, out.String(), "| x |") + assert.Contains(t, out.String(), "1 row in set") +} + +func TestRunStatementBinaryHexAndNumericJSON(t *testing.T) { + fakeQueries.Store("SELECT id, price, name FROM product", fakeTable{ + cols: []string{"id", "price", "name"}, + types: []string{"BINARY", "DECIMAL", "VARCHAR"}, + rows: [][]driver.Value{{[]byte{0xab, 0xcd}, []byte("19.99"), []byte("Fancy")}}, + }) + + var out strings.Builder + + require.NoError(t, RunStatement(t.Context(), openFakeDB(t), "SELECT id, price, name FROM product", &out, FormatJSON)) + + assert.JSONEq(t, `[{"id":"0xABCD","price":19.99,"name":"Fancy"}]`, out.String()) +} + +func TestRunStatementExecJSON(t *testing.T) { + var out strings.Builder + + require.NoError(t, RunStatement(t.Context(), openFakeDB(t), "DELETE FROM t", &out, FormatJSON)) + + assert.JSONEq(t, `{"rows_affected":3}`, out.String()) +} + +func TestRunStatementReturningUsesQueryPath(t *testing.T) { + fakeQueries.Store("INSERT INTO t (v) VALUES ('a') RETURNING id", fakeTable{ + cols: []string{"id"}, + types: []string{"BIGINT"}, + rows: [][]driver.Value{{[]byte("7")}}, + }) + + var out strings.Builder + + require.NoError(t, RunStatement(t.Context(), openFakeDB(t), "INSERT INTO t (v) VALUES ('a') RETURNING id", &out, FormatTSV)) + + assert.Equal(t, "id\n7\n", out.String()) +} + +func TestRunStatementSkipsCommentsAndDelimiter(t *testing.T) { + var out strings.Builder + + require.NoError(t, RunStatement(t.Context(), openFakeDB(t), "-- nothing here", &out, FormatTSV)) + require.NoError(t, RunStatement(t.Context(), openFakeDB(t), "DELIMITER ;", &out, FormatTSV)) + + assert.Empty(t, out.String()) +} + +func TestRunStatementEmptyResultSetColumns(t *testing.T) { + fakeQueries.Store("CALL cleanup", fakeTable{}) + + var out strings.Builder + + require.NoError(t, RunStatement(t.Context(), openFakeDB(t), "CALL cleanup", &out, FormatTSV)) + assert.Empty(t, out.String()) +} + +func TestRunMultiStatementErrorIncludesSummary(t *testing.T) { + fakeQueries.Store("SELECT 1", fakeTable{ + cols: []string{"1"}, + types: []string{"BIGINT"}, + rows: [][]driver.Value{{[]byte("1")}}, + }) + + var out strings.Builder + + err := Run(t.Context(), openFakeDB(t), "SELECT 1; SELECT boom; SELECT 1;", &out, FormatTSV) + require.Error(t, err) + + assert.Contains(t, err.Error(), "SELECT boom") + assert.Contains(t, err.Error(), "query exploded") +} + +func TestRunSingleStatementErrorHasNoSummary(t *testing.T) { + var out strings.Builder + + err := Run(t.Context(), openFakeDB(t), "SELECT boom", &out, FormatTSV) + require.Error(t, err) + + assert.Equal(t, "query exploded", err.Error()) +} + +func TestShellExecutesAndRecoversFromErrors(t *testing.T) { + fakeQueries.Store("SELECT 1", fakeTable{ + cols: []string{"1"}, + types: []string{"BIGINT"}, + rows: [][]driver.Value{{[]byte("1")}}, + }) + fakeQueries.Store("SELECT\n2", fakeTable{ + cols: []string{"2"}, + types: []string{"BIGINT"}, + rows: [][]driver.Value{{[]byte("2")}}, + }) + + input := "SELECT 1;\n" + // simple statement + "SELECT boom;\n" + // error must not end the session + "SELECT\n2;\n" + // statement spanning two lines + "UPDATE t SET one_row = 1" // no semicolon: runs on EOF + + var out, errOut strings.Builder + + err := Shell(t.Context(), openFakeDB(t), strings.NewReader(input), &out, &errOut, FormatTSV) + require.NoError(t, err) + + assert.Contains(t, out.String(), "1\n1\n") + assert.Contains(t, out.String(), "2\n2\n") + assert.Contains(t, out.String(), "Query OK, 1 row affected") + assert.Contains(t, out.String(), promptContinuation) + assert.Contains(t, errOut.String(), "query exploded") +} + +func TestShellStopsOnCancelledContext(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + var out, errOut strings.Builder + + err := Shell(ctx, openFakeDB(t), strings.NewReader("SELECT 1;\n"), &out, &errOut, FormatTSV) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestShellExitCommandStopsSession(t *testing.T) { + var out, errOut strings.Builder + + err := Shell(t.Context(), openFakeDB(t), strings.NewReader("exit\nSELECT boom;\n"), &out, &errOut, FormatTSV) + require.NoError(t, err) + + assert.Empty(t, errOut.String(), "statements after exit must not run") +} diff --git a/internal/sqlshell/run_test.go b/internal/sqlshell/run_test.go new file mode 100644 index 00000000..4230882e --- /dev/null +++ b/internal/sqlshell/run_test.go @@ -0,0 +1,178 @@ +package sqlshell + +import ( + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseFormat(t *testing.T) { + for _, valid := range []string{"table", "tsv", "json"} { + format, err := ParseFormat(valid) + assert.NoError(t, err) + assert.Equal(t, Format(valid), format) + } + + _, err := ParseFormat("xml") + assert.Error(t, err) +} + +func TestFirstKeyword(t *testing.T) { + cases := map[string]string{ + "SELECT 1": "select", + " select * from t": "select", + "(SELECT 1) UNION (SELECT 2)": "select", + "/* leading */ SHOW TABLES": "show", + "-- comment\nUPDATE t SET a = 1": "update", + "# comment\nDELETE FROM t": "delete", + "WITH x AS (SELECT 1) SELECT 1": "with", + "INSERT INTO t VALUES (1)": "insert", + "EXPLAIN(FORMAT=JSON) SELECT 1": "explain", + "/* only a comment, no keyword*/": "", + "-- only a comment": "", + + // MySQL executable comments are run by the server. + "/*!40014 SET FOREIGN_KEY_CHECKS=0 */": "set", + "/*!40000 ALTER TABLE `t` DISABLE KEYS */": "alter", + "/*!50001 CREATE VIEW v AS SELECT 1 FROM t */": "create", + } + + for input, expected := range cases { + assert.Equal(t, expected, firstKeyword(input), "input: %q", input) + } +} + +func TestReturnsResultSet(t *testing.T) { + for _, keyword := range []string{"select", "show", "describe", "desc", "explain", "with", "values", "table", "call"} { + assert.True(t, returnsResultSet(keyword), keyword) + } + + for _, keyword := range []string{"insert", "update", "delete", "replace", "create", "drop", "alter", "truncate", "set", "grant"} { + assert.False(t, returnsResultSet(keyword), keyword) + } +} + +func TestHasReturningClause(t *testing.T) { + positive := []string{ + "INSERT INTO t (a) VALUES (1) RETURNING id", + "insert into t values (1) returning *", + "DELETE FROM t WHERE id = 1 RETURNING id, name", + "REPLACE INTO t VALUES (1)\nRETURNING id", + "INSERT INTO t VALUES (1--1) RETURNING id", // -- without whitespace is subtraction, not a comment + "INSERT INTO t VALUES (1) /* note */RETURNING id", // a closed comment separates tokens + "DELETE FROM t WHERE id = 1 RETURNING/* note */ id", // also after the keyword + } + + for _, stmt := range positive { + assert.True(t, hasReturningClause(firstKeyword(stmt), stmt), stmt) + } + + negative := []string{ + "INSERT INTO t (note) VALUES ('use RETURNING here')", // inside string literal + `INSERT INTO t (note) VALUES ("RETURNING too")`, // inside double quotes + "INSERT INTO t (a) VALUES (1) -- RETURNING id", // inside line comment + "INSERT INTO t (a) VALUES (1) # RETURNING id", // inside hash comment + "INSERT INTO t /* RETURNING */ VALUES (1)", // inside block comment + "INSERT INTO `returning` VALUES (1)", // quoted identifier + "INSERT INTO t SET a = 1 - 1", // lone dash is not a comment + "UPDATE t SET a = 1", // keyword without clause + "SELECT returning FROM t", // not a DML keyword + "INSERT INTO t VALUES (returning_col)", // no boundary after word + "INSERT INTO t VALUES (1)", + } + + for _, stmt := range negative { + assert.False(t, hasReturningClause(firstKeyword(stmt), stmt), stmt) + } +} + +func TestSummarizeStatement(t *testing.T) { + assert.Equal(t, "SELECT 1", summarizeStatement("SELECT\n\t1")) + + long := "SELECT " + strings.Repeat("a", 100) + summary := summarizeStatement(long) + assert.True(t, strings.HasSuffix(summary, "…")) + assert.Less(t, len([]rune(summary)), 65) +} + +func TestRenderTable(t *testing.T) { + var out strings.Builder + + err := renderTable( + []string{"id", "name"}, + [][]any{ + {"1", "Standard rate"}, + {"2", nil}, + }, + time.Millisecond, + &out, + ) + require.NoError(t, err) + + expected := "+----+---------------+\n" + + "| id | name |\n" + + "+----+---------------+\n" + + "| 1 | Standard rate |\n" + + "| 2 | NULL |\n" + + "+----+---------------+\n" + + "2 rows in set (0.001 sec)\n" + + assert.Equal(t, expected, out.String()) +} + +func TestRenderTableEmptySet(t *testing.T) { + var out strings.Builder + + require.NoError(t, renderTable([]string{"id"}, nil, time.Millisecond, &out)) + assert.Equal(t, "Empty set (0.001 sec)\n", out.String()) +} + +func TestRenderTSV(t *testing.T) { + var out strings.Builder + + err := renderTSV( + []string{"id", "name"}, + [][]any{ + {"1", "with\ttab and\nnewline"}, + {"2", nil}, + }, + &out, + ) + require.NoError(t, err) + + expected := "id\tname\n" + + "1\twith\\ttab and\\nnewline\n" + + "2\tNULL\n" + + assert.Equal(t, expected, out.String()) +} + +func TestRenderJSON(t *testing.T) { + var out strings.Builder + + err := renderJSON( + []string{"id", "rate", "name"}, + []bool{true, true, false}, + [][]any{ + {"1", "19.00", "Standard"}, + {"2", nil, nil}, + }, + &out, + ) + require.NoError(t, err) + + assert.JSONEq(t, `[{"id":1,"rate":19.00,"name":"Standard"},{"id":2,"rate":null,"name":null}]`, out.String()) +} + +func TestIsExitCommand(t *testing.T) { + for _, line := range []string{"exit", "quit", "EXIT", "exit;", `\q`, " quit ; "} { + assert.True(t, isExitCommand(line), line) + } + + for _, line := range []string{"exits", "select 1", ""} { + assert.False(t, isExitCommand(line), line) + } +} diff --git a/internal/sqlshell/shell.go b/internal/sqlshell/shell.go new file mode 100644 index 00000000..473f410c --- /dev/null +++ b/internal/sqlshell/shell.go @@ -0,0 +1,87 @@ +package sqlshell + +import ( + "bufio" + "context" + "fmt" + "io" + "strings" +) + +const ( + promptMain = "sql> " + promptContinuation = " -> " +) + +// Shell runs an interactive read-eval-print loop until EOF or an exit +// command (exit, quit, \q). Statement errors are printed and do not end the +// session. +func Shell(ctx context.Context, db Conn, in io.Reader, out, errOut io.Writer, format Format) error { + scanner := bufio.NewScanner(in) + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + + var buffer string + delimiter := DefaultDelimiter + + if _, err := fmt.Fprint(out, promptMain); err != nil { + return err + } + + for scanner.Scan() { + if ctx.Err() != nil { + return ctx.Err() + } + + line := scanner.Text() + + if buffer == "" && isExitCommand(line) { + return nil + } + + buffer += line + "\n" + + var statements []string + // The remainder keeps its trailing newline: the splitter only trims + // the leading whitespace. + statements, buffer, delimiter = SplitStatementsWithDelimiter(buffer, delimiter) + for _, stmt := range statements { + if err := RunStatement(ctx, db, stmt, out, format); err != nil { + // Best-effort: the session continues even when stderr is gone. + _, _ = fmt.Fprintf(errOut, "ERROR: %s\n", err) + } + } + + prompt := promptMain + if buffer != "" { + prompt = promptContinuation + } + + if _, err := fmt.Fprint(out, prompt); err != nil { + return err + } + } + + // Execute what is left in the buffer on EOF, so a statement without a + // trailing semicolon before Ctrl+D is not silently discarded. This + // matches the trailing statement handling of Run and ExecuteStream. + if rest := strings.TrimSpace(buffer); rest != "" { + if err := RunStatement(ctx, db, rest, out, format); err != nil { + // Best-effort: the session ends anyway. + _, _ = fmt.Fprintf(errOut, "ERROR: %s\n", err) + } + } + + // Cosmetic newline after EOF. + _, _ = fmt.Fprintln(out) + + return scanner.Err() +} + +func isExitCommand(line string) bool { + switch strings.ToLower(strings.TrimRight(strings.TrimSpace(line), "; \t")) { + case "exit", "quit", `\q`: + return true + } + + return false +} diff --git a/internal/sqlshell/split.go b/internal/sqlshell/split.go new file mode 100644 index 00000000..35b21610 --- /dev/null +++ b/internal/sqlshell/split.go @@ -0,0 +1,189 @@ +// Package sqlshell implements a small SQL shell used by `project sql` and +// `project restore`: splitting scripts into statements, executing them and +// rendering results. +package sqlshell + +import "strings" + +// DefaultDelimiter is the statement terminator used unless a DELIMITER +// directive changes it. +const DefaultDelimiter = ";" + +const ( + stateNormal = iota + stateSingleQuote + stateDoubleQuote + stateBacktick + stateLineComment + stateBlockComment +) + +// SplitStatements splits input into complete SQL statements and returns the +// trailing incomplete remainder. It starts with the default ";" delimiter; +// use SplitStatementsWithDelimiter to carry delimiter state across calls. +func SplitStatements(input string) ([]string, string) { + statements, rest, _ := SplitStatementsWithDelimiter(input, DefaultDelimiter) + return statements, rest +} + +// SplitStatementsWithDelimiter splits input into complete SQL statements, +// terminated by delimiter, and returns the trailing incomplete remainder plus +// the delimiter active after the input. Terminators inside single-quoted, +// double-quoted or backtick-quoted sections, line comments (-- , #) and block +// comments (/* */) do not end a statement. +// +// DELIMITER directives (a client feature used by dumps around triggers and +// routines) are consumed, not returned as statements: they change the active +// delimiter from their line onwards. +// +// Only leading whitespace is trimmed from the remainder: its trailing +// whitespace is significant when more input is appended later (streaming). +func SplitStatementsWithDelimiter(input, delimiter string) ([]string, string, string) { + if delimiter == "" { + delimiter = DefaultDelimiter + } + + var statements []string + + state := stateNormal + start := 0 + // blankStatement tracks whether only whitespace and comments were seen + // since the last statement boundary; a DELIMITER directive is only + // recognized there. + blankStatement := true + + for i := 0; i < len(input); i++ { + c := input[i] + + if state != stateNormal { + state, i = consumeQuoteOrComment(state, input, i) + continue + } + + if blankStatement && isDelimiterDirective(input[i:]) { + lineEnd := strings.IndexByte(input[i:], '\n') + if lineEnd == -1 { + // The directive line may continue in the next chunk. + return statements, strings.TrimLeft(input[i:], " \t\r\n"), delimiter + } + + if token := delimiterToken(input[i : i+lineEnd]); token != "" { + delimiter = token + } + + start = i + lineEnd + 1 + i += lineEnd + + continue + } + + if strings.HasPrefix(input[i:], delimiter) { + if stmt := strings.TrimSpace(input[start:i]); stmt != "" { + statements = append(statements, stmt) + } + + start = i + len(delimiter) + i = start - 1 + blankStatement = true + + continue + } + + switch c { + case '\'': + state = stateSingleQuote + blankStatement = false + case '"': + state = stateDoubleQuote + blankStatement = false + case '`': + state = stateBacktick + blankStatement = false + case '#': + state = stateLineComment + case '-': + // MySQL only treats "--" as a comment when followed by + // whitespace or the end of the line. + if strings.HasPrefix(input[i:], "--") && (i+2 >= len(input) || input[i+2] == ' ' || input[i+2] == '\t' || input[i+2] == '\n' || input[i+2] == '\r') { + state = stateLineComment + i++ + } else { + blankStatement = false + } + case '/': + if strings.HasPrefix(input[i:], "/*") { + state = stateBlockComment + i++ + } else { + blankStatement = false + } + case ' ', '\t', '\r', '\n': + // whitespace keeps the statement blank + default: + blankStatement = false + } + } + + return statements, strings.TrimLeft(input[start:], " \t\r\n"), delimiter +} + +// consumeQuoteOrComment advances the scanner while inside a quoted section +// or comment. It processes input[i] and returns the resulting state and the +// last consumed index. +func consumeQuoteOrComment(state int, input string, i int) (int, int) { + c := input[i] + + switch state { + case stateSingleQuote: + switch c { + case '\\': + i++ + case '\'': + state = stateNormal + } + case stateDoubleQuote: + switch c { + case '\\': + i++ + case '"': + state = stateNormal + } + case stateBacktick: + if c == '`' { + state = stateNormal + } + case stateLineComment: + if c == '\n' { + state = stateNormal + } + case stateBlockComment: + if strings.HasPrefix(input[i:], "*/") { + state = stateNormal + i++ + } + } + + return state, i +} + +// isDelimiterDirective reports whether s starts with a DELIMITER keyword +// followed by its argument. +func isDelimiterDirective(s string) bool { + const keyword = "delimiter" + + if len(s) <= len(keyword) || !strings.EqualFold(s[:len(keyword)], keyword) { + return false + } + + return s[len(keyword)] == ' ' || s[len(keyword)] == '\t' +} + +// delimiterToken extracts the new delimiter from a DELIMITER directive line. +func delimiterToken(line string) string { + fields := strings.Fields(line) + if len(fields) < 2 { + return "" + } + + return fields[1] +} diff --git a/internal/sqlshell/split_test.go b/internal/sqlshell/split_test.go new file mode 100644 index 00000000..79534ed8 --- /dev/null +++ b/internal/sqlshell/split_test.go @@ -0,0 +1,149 @@ +package sqlshell + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestSplitStatements(t *testing.T) { + cases := []struct { + name string + input string + statements []string + rest string + }{ + { + name: "single statement", + input: "SELECT 1;", + statements: []string{"SELECT 1"}, + }, + { + name: "multiple statements", + input: "SELECT 1; SELECT 2;\nSELECT 3;", + statements: []string{"SELECT 1", "SELECT 2", "SELECT 3"}, + }, + { + name: "incomplete statement", + input: "SELECT * FROM product", + rest: "SELECT * FROM product", + }, + { + name: "complete and incomplete", + input: "SELECT 1; SELECT 2", + statements: []string{"SELECT 1"}, + rest: "SELECT 2", + }, + { + name: "semicolon in single quotes", + input: "SELECT 'a;b';", + statements: []string{"SELECT 'a;b'"}, + }, + { + name: "semicolon in double quotes", + input: `SELECT "a;b";`, + statements: []string{`SELECT "a;b"`}, + }, + { + name: "semicolon in backticks", + input: "SELECT `a;b` FROM t;", + statements: []string{"SELECT `a;b` FROM t"}, + }, + { + name: "escaped quote in string", + input: `SELECT 'it\'s;ok';`, + statements: []string{`SELECT 'it\'s;ok'`}, + }, + { + name: "doubled quote in string", + input: "SELECT 'it''s;ok';", + statements: []string{"SELECT 'it''s;ok'"}, + }, + { + name: "semicolon in line comment", + input: "SELECT 1 -- comment;\n;", + statements: []string{"SELECT 1 -- comment;"}, + }, + { + name: "semicolon in hash comment", + input: "SELECT 1 # comment;\n;", + statements: []string{"SELECT 1 # comment;"}, + }, + { + name: "semicolon in block comment", + input: "SELECT /* ; */ 1;", + statements: []string{"SELECT /* ; */ 1"}, + }, + { + name: "double dash without space is no comment", + input: "SELECT 1--2;", + statements: []string{"SELECT 1--2"}, + }, + { + name: "empty statements are dropped", + input: ";; ;SELECT 1;", + statements: []string{"SELECT 1"}, + }, + { + name: "unterminated string keeps rest", + input: "SELECT 'abc", + rest: "SELECT 'abc", + }, + { + name: "delimiter directive for trigger", + input: "DROP TRIGGER IF EXISTS `t`;\n" + + "DELIMITER //\n" + + "CREATE TRIGGER t BEFORE INSERT ON x FOR EACH ROW BEGIN SET @a = 1; SET @b = 2; END//\n" + + "DELIMITER ;\n" + + "SELECT 1;", + statements: []string{ + "DROP TRIGGER IF EXISTS `t`", + "CREATE TRIGGER t BEFORE INSERT ON x FOR EACH ROW BEGIN SET @a = 1; SET @b = 2; END", + "SELECT 1", + }, + }, + { + name: "delimiter directive is case-insensitive", + input: "delimiter $$\nSELECT 1$$SELECT 2$$", + statements: []string{"SELECT 1", "SELECT 2"}, + }, + { + name: "delimiter word inside statement is content", + input: "SELECT 'DELIMITER //';", + statements: []string{"SELECT 'DELIMITER //'"}, + }, + { + name: "incomplete delimiter directive stays in rest", + input: "SELECT 1;\nDELIMITER /", + statements: []string{ + "SELECT 1", + }, + rest: "DELIMITER /", + }, + { + name: "delimiter directive without token keeps delimiter", + input: "DELIMITER \nSELECT 1;", + statements: []string{"SELECT 1"}, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + statements, rest := SplitStatements(tc.input) + assert.Equal(t, tc.statements, statements) + assert.Equal(t, tc.rest, rest) + }) + } +} + +func TestSplitStatementsCarriesDelimiterAcrossCalls(t *testing.T) { + statements, rest, delimiter := SplitStatementsWithDelimiter("DELIMITER //\nCREATE TRIGGER t BEGIN SET @a = 1; ", DefaultDelimiter) + assert.Empty(t, statements) + assert.Equal(t, "CREATE TRIGGER t BEGIN SET @a = 1; ", rest) + assert.Equal(t, "//", delimiter) + + statements, rest, delimiter = SplitStatementsWithDelimiter(rest+"END//\nDELIMITER ;\nSELECT 1;", delimiter) + assert.Equal(t, []string{"CREATE TRIGGER t BEGIN SET @a = 1; END", "SELECT 1"}, statements) + assert.Empty(t, rest) + assert.Equal(t, ";", delimiter) +} diff --git a/internal/sqlshell/stream.go b/internal/sqlshell/stream.go new file mode 100644 index 00000000..8b55854e --- /dev/null +++ b/internal/sqlshell/stream.go @@ -0,0 +1,74 @@ +package sqlshell + +import ( + "context" + "fmt" + "io" + "strings" +) + +// ExecuteStream reads SQL statements from r and executes them one by one, +// without rendering results. It is meant for restoring dumps, where the input +// can be far larger than memory. onStatement, when non-nil, is called after +// every executed statement with the running count. The number of executed +// statements is returned, also when an error aborts the run. +func ExecuteStream(ctx context.Context, db Execer, r io.Reader, onStatement func(count int)) (int, error) { + // No bufio wrapper: reads of len(chunk) would bypass its internal buffer + // anyway, it would only allocate a second unused megabyte. + chunk := make([]byte, 1<<20) + + var buffer string + delimiter := DefaultDelimiter + count := 0 + + execute := func(stmt string) error { + if keyword := firstKeyword(stmt); keyword == "" || keyword == "delimiter" { + return nil + } + + if _, err := db.ExecContext(ctx, stmt); err != nil { + return fmt.Errorf("%q: %w", summarizeStatement(stmt), err) + } + + count++ + + if onStatement != nil { + onStatement(count) + } + + return nil + } + + for { + n, readErr := r.Read(chunk) + + if n > 0 { + buffer += string(chunk[:n]) + + var statements []string + statements, buffer, delimiter = SplitStatementsWithDelimiter(buffer, delimiter) + for _, stmt := range statements { + if err := execute(stmt); err != nil { + return count, err + } + } + } + + if readErr == io.EOF { + break + } + + if readErr != nil { + return count, readErr + } + } + + // The trailing semicolon is optional. + if strings.TrimSpace(buffer) != "" { + if err := execute(buffer); err != nil { + return count, err + } + } + + return count, nil +} diff --git a/internal/sqlshell/stream_test.go b/internal/sqlshell/stream_test.go new file mode 100644 index 00000000..ee4e9102 --- /dev/null +++ b/internal/sqlshell/stream_test.go @@ -0,0 +1,139 @@ +package sqlshell + +import ( + "context" + "database/sql" + "errors" + "io" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type fakeExecer struct { + statements []string + failOn string +} + +type fakeResult struct{} + +func (fakeResult) LastInsertId() (int64, error) { return 0, nil } +func (fakeResult) RowsAffected() (int64, error) { return 0, nil } + +func (f *fakeExecer) ExecContext(_ context.Context, query string, _ ...any) (sql.Result, error) { + if f.failOn != "" && strings.Contains(query, f.failOn) { + return nil, errors.New("boom") + } + + f.statements = append(f.statements, query) + + return fakeResult{}, nil +} + +// slowReader yields the input in tiny chunks so statements span reads. +type slowReader struct { + data string + pos int +} + +func (s *slowReader) Read(p []byte) (int, error) { + if s.pos >= len(s.data) { + return 0, io.EOF + } + + n := copy(p[:min(len(p), 3)], s.data[s.pos:]) + s.pos += n + + return n, nil +} + +func TestExecuteStream(t *testing.T) { + script := "SET NAMES utf8mb4;\n" + + "-- structure of table t;\n" + + "CREATE TABLE t (name VARCHAR(20));\n" + + "/*!40014 SET FOREIGN_KEY_CHECKS=0 */;\n" + + "INSERT INTO t VALUES ('a;b'), ('c');\n" + + "/* plain comment */;\n" + + "UPDATE t SET name = 'x' WHERE name = 'c'" + + db := &fakeExecer{} + + var counts []int + count, err := ExecuteStream(t.Context(), db, &slowReader{data: script}, func(n int) { + counts = append(counts, n) + }) + require.NoError(t, err) + + assert.Equal(t, 5, count) + assert.Equal(t, []int{1, 2, 3, 4, 5}, counts) + assert.Equal(t, []string{ + "SET NAMES utf8mb4", + "-- structure of table t;\nCREATE TABLE t (name VARCHAR(20))", + "/*!40014 SET FOREIGN_KEY_CHECKS=0 */", + "INSERT INTO t VALUES ('a;b'), ('c')", + "UPDATE t SET name = 'x' WHERE name = 'c'", + }, db.statements) +} + +func TestExecuteStreamTriggerDump(t *testing.T) { + // Mirrors the trigger section emitted by internal/mysqldump, fed in + // 3-byte chunks so the delimiter state must survive chunk boundaries. + script := "DROP TRIGGER IF EXISTS `order_update`;\n" + + "DELIMITER //\n" + + "CREATE TRIGGER order_update BEFORE UPDATE ON `order` FOR EACH ROW BEGIN\n" + + " SET NEW.updated_at = NOW();\n" + + " SET @counter = @counter + 1;\n" + + "END//\n" + + "DELIMITER ;\n" + + "INSERT INTO t VALUES (1);\n" + + db := &fakeExecer{} + + count, err := ExecuteStream(t.Context(), db, &slowReader{data: script}, nil) + require.NoError(t, err) + + assert.Equal(t, 3, count) + assert.Equal(t, []string{ + "DROP TRIGGER IF EXISTS `order_update`", + "CREATE TRIGGER order_update BEFORE UPDATE ON `order` FOR EACH ROW BEGIN\n" + + " SET NEW.updated_at = NOW();\n" + + " SET @counter = @counter + 1;\n" + + "END", + "INSERT INTO t VALUES (1)", + }, db.statements) +} + +func TestExecuteStreamTrailingDelimiterDirective(t *testing.T) { + db := &fakeExecer{} + + // No trailing newline after the final directive: it must not be sent to + // the server as SQL. + count, err := ExecuteStream(t.Context(), db, strings.NewReader("SELECT 1;\nDELIMITER ;"), nil) + require.NoError(t, err) + + assert.Equal(t, 1, count) + assert.Equal(t, []string{"SELECT 1"}, db.statements) +} + +func TestExecuteStreamStopsOnError(t *testing.T) { + db := &fakeExecer{failOn: "two"} + + count, err := ExecuteStream(t.Context(), db, strings.NewReader("SELECT one; SELECT two; SELECT three;"), nil) + require.Error(t, err) + + assert.Equal(t, 1, count) + assert.Contains(t, err.Error(), "SELECT two") + assert.Equal(t, []string{"SELECT one"}, db.statements) +} + +func TestExecuteStreamEmptyInput(t *testing.T) { + db := &fakeExecer{} + + count, err := ExecuteStream(t.Context(), db, strings.NewReader("\n-- nothing to do\n"), nil) + require.NoError(t, err) + + assert.Equal(t, 0, count) + assert.Empty(t, db.statements) +} diff --git a/internal/system/decompress.go b/internal/system/decompress.go new file mode 100644 index 00000000..94c7580c --- /dev/null +++ b/internal/system/decompress.go @@ -0,0 +1,42 @@ +package system + +import ( + "bufio" + "bytes" + "compress/gzip" + "errors" + "io" + + "github.com/klauspost/compress/zstd" +) + +var ( + gzipMagic = []byte{0x1f, 0x8b} + zstdMagic = []byte{0x28, 0xb5, 0x2f, 0xfd} +) + +// DecompressReader detects gzip and zstd compression from the leading magic +// bytes and returns a transparently decompressing reader. Close the returned +// reader when it implements io.Closer to release decompressor resources. +func DecompressReader(r io.Reader) (io.Reader, error) { + buffered := bufio.NewReaderSize(r, 1<<16) + + magic, err := buffered.Peek(4) + if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) { + return nil, err + } + + switch { + case bytes.HasPrefix(magic, gzipMagic): + return gzip.NewReader(buffered) + case bytes.HasPrefix(magic, zstdMagic): + reader, err := zstd.NewReader(buffered) + if err != nil { + return nil, err + } + + return reader.IOReadCloser(), nil + } + + return buffered, nil +} diff --git a/internal/system/decompress_test.go b/internal/system/decompress_test.go new file mode 100644 index 00000000..8d0ba9a4 --- /dev/null +++ b/internal/system/decompress_test.go @@ -0,0 +1,81 @@ +package system + +import ( + "bytes" + "compress/gzip" + "errors" + "io" + "strings" + "testing" + + "github.com/klauspost/compress/zstd" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDecompressReaderPlain(t *testing.T) { + reader, err := DecompressReader(strings.NewReader("SELECT 1;")) + require.NoError(t, err) + + content, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, "SELECT 1;", string(content)) +} + +func TestDecompressReaderGzip(t *testing.T) { + var buf bytes.Buffer + writer := gzip.NewWriter(&buf) + _, err := writer.Write([]byte("SELECT 'gz';")) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + reader, err := DecompressReader(&buf) + require.NoError(t, err) + + content, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, "SELECT 'gz';", string(content)) +} + +func TestDecompressReaderZstd(t *testing.T) { + var buf bytes.Buffer + writer, err := zstd.NewWriter(&buf) + require.NoError(t, err) + _, err = writer.Write([]byte("SELECT 'zst';")) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + reader, err := DecompressReader(&buf) + require.NoError(t, err) + + content, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, "SELECT 'zst';", string(content)) +} + +func TestDecompressReaderEmpty(t *testing.T) { + reader, err := DecompressReader(strings.NewReader("")) + require.NoError(t, err) + + content, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Empty(t, content) +} + +func TestDecompressReaderShortPlainInput(t *testing.T) { + reader, err := DecompressReader(strings.NewReader(";")) + require.NoError(t, err) + + content, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, ";", string(content)) +} + +type failingReader struct{} + +func (failingReader) Read([]byte) (int, error) { return 0, errors.New("disk error") } + +func TestDecompressReaderPropagatesReadError(t *testing.T) { + _, err := DecompressReader(failingReader{}) + assert.ErrorContains(t, err, "disk error") +} diff --git a/internal/system/reader.go b/internal/system/reader.go new file mode 100644 index 00000000..8acef9a2 --- /dev/null +++ b/internal/system/reader.go @@ -0,0 +1,23 @@ +package system + +import "io" + +// CountingReader wraps a reader and counts the bytes read through it, e.g. +// for progress reporting against a known total size. +type CountingReader struct { + Reader io.Reader + + count int64 +} + +func (c *CountingReader) Read(p []byte) (int, error) { + n, err := c.Reader.Read(p) + c.count += int64(n) + + return n, err +} + +// BytesRead returns the number of bytes read so far. +func (c *CountingReader) BytesRead() int64 { + return c.count +} diff --git a/internal/system/reader_test.go b/internal/system/reader_test.go new file mode 100644 index 00000000..7488cc07 --- /dev/null +++ b/internal/system/reader_test.go @@ -0,0 +1,20 @@ +package system + +import ( + "io" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCountingReader(t *testing.T) { + counting := &CountingReader{Reader: strings.NewReader("hello world")} + + content, err := io.ReadAll(counting) + require.NoError(t, err) + + assert.Equal(t, "hello world", string(content)) + assert.Equal(t, int64(11), counting.BytesRead()) +}