Add pscale import d1 for Cloudflare D1 offline migration

go.mod

@@ -21,6 +21,7 @@ require (
 	github.com/google/go-cmp v0.7.0
 	github.com/hashicorp/go-cleanhttp v0.5.2
 	github.com/hashicorp/go-version v1.8.0
+	github.com/jackc/pgx/v5 v5.8.0
 	github.com/lensesio/tableprinter v0.0.0-20201125135848-89e81fc956e7
 	github.com/lib/pq v1.12.0
 	github.com/mark3labs/mcp-go v0.46.0
@@ -30,7 +31,7 @@ require (
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/muesli/termenv v0.16.0
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
-	github.com/planetscale/planetscale-go v0.171.0
+	github.com/planetscale/planetscale-go v0.173.0
 	github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4
 	github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7
 	github.com/spf13/cobra v1.10.2
@@ -78,6 +79,9 @@ require (
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
+	github.com/jackc/puddle/v2 v2.2.2 // indirect
 	github.com/kataras/tablewriter v0.0.0-20180708051242-e063d29b7c23 // indirect
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
 	github.com/klauspost/compress v1.18.2 // indirect
@@ -115,7 +119,6 @@ require (
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect
 	google.golang.org/grpc v1.79.3 // indirect
 	google.golang.org/protobuf v1.36.10 // indirect
-	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 

go.sum

@@ -112,6 +112,14 @@ github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec h1:qv2VnGeEQHchGaZ/u
 github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec/go.mod h1:Q48J4R4DvxnHolD5P8pOtXigYlRuPLGl6moFx3ulM68=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
+github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
+github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
+github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
+github.com/jackc/pgx/v5 v5.8.0 h1:TYPDoleBBme0xGSAX3/+NujXXtpZn9HBONkQC7IEZSo=
+github.com/jackc/pgx/v5 v5.8.0/go.mod h1:QVeDInX2m9VyzvNeiCJVjCkNFqzsNb43204HshNSZKw=
+github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
+github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
 github.com/kataras/tablewriter v0.0.0-20180708051242-e063d29b7c23 h1:M8exrBzuhWcU6aoHJlHWPe4qFjVKzkMGRal78f5jRRU=
 github.com/kataras/tablewriter v0.0.0-20180708051242-e063d29b7c23/go.mod h1:kBSna6b0/RzsOcOZf515vAXwSsXYusl2U7SA0XP09yI=
 github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs=
@@ -120,7 +128,6 @@ github.com/klauspost/compress v1.18.2 h1:iiPHWW0YrcFgpBYhsA6D1+fqHssJscY/Tm/y2Uq
 github.com/klauspost/compress v1.18.2/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4=
 github.com/klauspost/connect-compress/v2 v2.1.0 h1:8fM8QrVeHT69e5VVSh4yjDaQASYIvOp2uMZq7nVLj2U=
 github.com/klauspost/connect-compress/v2 v2.1.0/go.mod h1:Ayurh2wscMMx3AwdGGVL+ylSR5316WfApREDgsqHyH8=
-github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
@@ -176,8 +183,8 @@ github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjL
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
 github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e h1:MZ8D+Z3m2vvqGZLvoQfpaGg/j1fNDr4j03s3PRz4rVY=
 github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e/go.mod h1:hwAsSPQdvPa3WcfKfzTXxtEq/HlqwLjQasfO6QbGo4Q=
-github.com/planetscale/planetscale-go v0.171.0 h1:ZBltEV1SANoQDUMc0VVUmGD4urjquZqh/y9XgGxG3Pw=
-github.com/planetscale/planetscale-go v0.171.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
+github.com/planetscale/planetscale-go v0.173.0 h1:++++Fg31PMPRyl5IaR65XyW1WRd8wiv2OSS32hUCXnE=
+github.com/planetscale/planetscale-go v0.173.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 h1:Xv5pj20Rhfty1Tv0OVcidg4ez4PvGrpKvb6rvUwQgDs=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4/go.mod h1:M52h5IWxAcbdQ1hSZrLAGQC4ZXslxEsK/Wh9nu3wdWs=
 github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 h1:aRd6vdE1fyuSI4RVj7oCr8lFmgqXvpnPUmN85VbZCp8=
@@ -218,6 +225,7 @@ github.com/spf13/viper v1.21.0/go.mod h1:P0lhsswPGWD/1lZJ9ny3fYnVqxiegrlNrEmgLjb
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=

internal/cmd/importcmd/d1.go

@@ -0,0 +1,142 @@
+package importcmd
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+
+	ps "github.com/planetscale/planetscale-go/planetscale"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/import/d1"
+	"github.com/planetscale/cli/internal/printer"
+)
+
+const defaultD1Branch = "main"
+
+var d1DatabaseBranchArgs = cobra.RangeArgs(1, 2)
+
+func parseDatabaseBranch(args []string) (database, branch string) {
+	database = args[0]
+	branch = defaultD1Branch
+	if len(args) > 1 {
+		branch = args[1]
+	}
+	return database, branch
+}
+
+func d1Org(ch *cmdutil.Helper) string {
+	return ch.Config.Organization
+}
+
+func writeD1(ch *cmdutil.Helper, resp d1.Response) error {
+	if resp.Status == "error" {
+		switch ch.Printer.Format() {
+		case printer.JSON:
+			if err := ch.Printer.PrintJSON(resp); err != nil {
+				return err
+			}
+		case printer.Human:
+			d1.PrintHumanResponse(ch.Printer, resp)
+		default:
+			return fmt.Errorf(`import d1 does not support output format %q (use human or json)`, ch.Printer.Format())
+		}
+		return d1CommandError(resp)
+	}
+
+	switch ch.Printer.Format() {
+	case printer.JSON:
+		return ch.Printer.PrintJSON(resp)
+	case printer.Human:
+		d1.PrintHumanResponse(ch.Printer, resp)
+		return nil
+	default:
+		return fmt.Errorf(`import d1 does not support output format %q (use human or json)`, ch.Printer.Format())
+	}
+}
+
+func d1CommandError(resp d1.Response) error {
+	msg := "import d1 command failed"
+	if resp.Error != nil {
+		msg = resp.Error.Message
+		if resp.Error.Remediation != "" {
+			msg += "\n" + resp.Error.Remediation
+		}
+	}
+	return &cmdutil.Error{
+		Msg:      msg,
+		ExitCode: cmdutil.ActionRequestedExitCode,
+		Printed:  true,
+	}
+}
+
+func d1NotifyAPI(client *ps.Client, disabled bool) d1.NotifyAPIConfig {
+	return d1.NotifyAPIConfig{Client: client, Disabled: disabled}
+}
+
+func importTableCount(prepared *d1.ImportPrepareResult) int {
+	if prepared == nil || prepared.Plan == nil {
+		return 0
+	}
+	return countDataTables(prepared.Plan.Tables)
+}
+
+func verifyTableCount(org, database, branch, migrationID, inputPath string) int {
+	path := inputPath
+	if path == "" && migrationID != "" {
+		if state, err := d1.LoadState(org, database, branch, migrationID); err == nil {
+			path = state.InputPath
+		}
+	}
+	if path == "" {
+		return 0
+	}
+	tables, err := d1.ParseDump(path)
+	if err != nil {
+		return 0
+	}
+	n := 0
+	for _, t := range tables {
+		if !d1.IsORMMetadataTable(t.Name) {
+			n++
+		}
+	}
+	return n
+}
+
+func countDataTables(tables []d1.TablePlan) int {
+	n := 0
+	for _, table := range tables {
+		if !d1.IsORMMetadataTable(table.Name) {
+			n++
+		}
+	}
+	return n
+}
+
+// D1Cmd returns the import d1 subcommand group.
+func D1Cmd(ch *cmdutil.Helper) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "d1 <command>",
+		Short: "Import Cloudflare D1 into PlanetScale Postgres",
+		Long: `Offline import from Cloudflare D1 (SQLite) to PlanetScale Postgres.
+
+Export your D1 database with wrangler (wrangler d1 export <name> --remote --output ./d1-export.sql),
+lint the dump, then start the import (use --dry-run to preview).
+All commands support --format json for machine-readable output.
+
+Branch-scoped commands use the same positional form as other PlanetScale CLI commands:
+  pscale import d1 start <database> [branch] --input ./d1-export.sql
+Org comes from your pscale config (pscale org).`,
+	}
+
+	cmd.AddCommand(d1DoctorCmd(ch))
+	cmd.AddCommand(d1LintCmd(ch))
+	cmd.AddCommand(d1ConvertSchemaCmd(ch))
+	cmd.AddCommand(d1StartCmd(ch))
+	cmd.AddCommand(d1VerifyCmd(ch))
+	cmd.AddCommand(d1StatusCmd(ch))
+	cmd.AddCommand(d1CompleteCmd(ch))
+
+	return cmd
+}

internal/cmd/importcmd/d1_complete.go

@@ -0,0 +1,52 @@
+package importcmd
+
+import (
+	"github.com/spf13/cobra"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/import/d1"
+	"github.com/planetscale/cli/internal/printer"
+)
+
+func d1CompleteCmd(ch *cmdutil.Helper) *cobra.Command {
+	var flags struct {
+		migrationID string
+		force       bool
+		noNotify    bool
+	}
+
+	cmd := &cobra.Command{
+		Use:     "complete <database> [branch]",
+		Aliases: []string{"teardown"},
+		Short:   "Mark a D1 migration as complete in local state",
+		Args:    d1DatabaseBranchArgs,
+		Example: `  pscale import d1 complete mydb --migration-id abc123
+  pscale import d1 complete mydb --migration-id abc123 --format json`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			database, branch := parseDatabaseBranch(args)
+			if !flags.force && ch.Printer.Format() == printer.Human {
+				if err := ch.Printer.ConfirmCommand(flags.migrationID, "import d1 complete", "complete"); err != nil {
+					return err
+				}
+			}
+			client, err := ch.Client()
+			if err != nil {
+				return writeD1(ch, d1.ErrorResponse("complete", err))
+			}
+			resp, err := d1.CompleteResponse(d1Org(ch), database, branch, flags.migrationID)
+			if err != nil {
+				return writeD1(ch, d1.ErrorResponse("complete", err))
+			}
+			if err := d1.Complete(d1Org(ch), database, branch, flags.migrationID, d1NotifyAPI(client, flags.noNotify)); err != nil {
+				return writeD1(ch, d1.ErrorResponse("complete", err))
+			}
+			return writeD1(ch, resp)
+		},
+	}
+
+	cmd.Flags().StringVar(&flags.migrationID, "migration-id", "", "Migration ID")
+	cmd.Flags().BoolVar(&flags.force, "force", false, "Skip confirmation prompt")
+	cmd.Flags().BoolVar(&flags.noNotify, "no-notify", false, "Skip Slack notifications for this completion")
+	cmd.MarkFlagRequired("migration-id")
+	return cmd
+}

internal/cmd/importcmd/d1_complete_test.go

@@ -0,0 +1,89 @@
+package importcmd
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	ps "github.com/planetscale/planetscale-go/planetscale"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/config"
+	"github.com/planetscale/cli/internal/import/d1"
+	"github.com/planetscale/cli/internal/printer"
+)
+
+func TestD1CompleteCmd(t *testing.T) {
+	t.Setenv("PSCALE_TEST_MODE", "1")
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusAccepted)
+	}))
+	defer srv.Close()
+
+	client, err := ps.NewClient(
+		ps.WithBaseURL(srv.URL),
+		ps.WithAccessToken("token"),
+	)
+	if err != nil {
+		t.Fatalf("NewClient: %v", err)
+	}
+
+	const migrationID = "completecmd123"
+	fixture := d1FixturePath(t)
+	if err := d1.SavePlan(&d1.PlanResult{
+		MigrationID: migrationID,
+		Org:         "acme",
+		Database:    "mydb",
+		Branch:      "main",
+		InputPath:   fixture,
+	}); err != nil {
+		t.Fatalf("SavePlan: %v", err)
+	}
+	if err := d1.SetMigrationPhase("acme", "mydb", "main", migrationID, d1.PhaseVerified); err != nil {
+		t.Fatalf("SetMigrationPhase: %v", err)
+	}
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: "acme"},
+		Client: func() (*ps.Client, error) {
+			return client, nil
+		},
+	}
+
+	cmd := d1CompleteCmd(ch)
+	cmd.SetArgs([]string{"mydb", "--migration-id", migrationID, "--force"})
+	cmd.SetOut(io.Discard)
+	cmd.SetErr(io.Discard)
+	if err := cmd.Execute(); err != nil {
+		t.Fatalf("execute: %v", err)
+	}
+
+	assertJSONField(t, &buf, "command", "complete")
+	assertJSONField(t, &buf, "status", "ok")
+	assertJSONField(t, &buf, "migration_id", migrationID)
+	if !strings.Contains(buf.String(), "reminder") {
+		t.Fatalf("expected reminder in complete JSON output:\n%s", buf.String())
+	}
+	if !strings.Contains(buf.String(), "next_steps") {
+		t.Fatalf("expected next_steps in complete JSON output:\n%s", buf.String())
+	}
+}
+
+func TestD1CompleteCmdRequiresMigrationID(t *testing.T) {
+	ch, _ := newD1TestHelper(t)
+
+	cmd := d1CompleteCmd(ch)
+	if err := executeD1Cmd(t, cmd, "mydb"); err == nil {
+		t.Fatal("expected error when --migration-id is missing")
+	}
+}