From 3615798cab8a3d28537066a089a341ba26a90a6a Mon Sep 17 00:00:00 2001
From: Akkuman <akkumans@qq.com>
Date: Thu, 26 Feb 2026 13:57:16 +0800
Subject: [PATCH] fix: fixed `packet for query is too large` error in restore

Signed-off-by: Akkuman <akkumans@qq.com>
---
 cmd/dump.go                |  5 ++++-
 cmd/restore.go             |  9 +++++++++
 pkg/database/connection.go | 14 +++++++++-----
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/cmd/dump.go b/cmd/dump.go
index 83a3084..ec5da90 100644
--- a/cmd/dump.go
+++ b/cmd/dump.go
@@ -24,7 +24,7 @@ const (
 	defaultCompression      = "gzip"
 	defaultBegin            = "+0"
 	defaultFrequency        = 1440
-	defaultMaxAllowedPacket = 4194304
+	defaultMaxAllowedPacket = 4 << 20 // 4 MiB (4194304)
 	defaultFilenamePattern  = core.DefaultFilenamePattern
 )
 
@@ -136,6 +136,9 @@ func dumpCmd(passedExecs execs, cmdConfig *cmdConfiguration) (*cobra.Command, er
 			if !v.IsSet("max-allowed-packet") && dumpConfig != nil && dumpConfig.MaxAllowedPacket != nil && *dumpConfig.MaxAllowedPacket != 0 {
 				maxAllowedPacket = *dumpConfig.MaxAllowedPacket
 			}
+			if maxAllowedPacket != 0 && maxAllowedPacket != defaultMaxAllowedPacket {
+				cmdConfig.dbconn.MaxAllowedPacket = maxAllowedPacket
+			}
 
 			// compression algorithm: check config, then CLI/env var overrides
 			var (
diff --git a/cmd/restore.go b/cmd/restore.go
index 127247f..e10a58a 100644
--- a/cmd/restore.go
+++ b/cmd/restore.go
@@ -75,6 +75,12 @@ func restoreCmd(passedExecs execs, cmdConfig *cmdConfiguration) (*cobra.Command,
 				}
 			}
 
+			// max-allowed-packet size
+			maxAllowedPacket := v.GetInt("max-allowed-packet")
+			if maxAllowedPacket != 0 && maxAllowedPacket != defaultMaxAllowedPacket {
+				cmdConfig.dbconn.MaxAllowedPacket = maxAllowedPacket
+			}
+
 			// target URL can reference one from the config file, or an absolute one
 			// if it's not in the config file, it's an absolute one
 			// if it is in the config file, it's a reference to one of the targets in the config file
@@ -160,5 +166,8 @@ func restoreCmd(passedExecs execs, cmdConfig *cmdConfiguration) (*cobra.Command,
 	// post-restore scripts
 	flags.String("post-restore-scripts", "", "Directory wherein any file ending in `.sh` will be run post-restore.")
 
+	// max-allowed-packet size
+	flags.Int("max-allowed-packet", 0, "Maximum size of the buffer for client/server communication, similar to mysql's max_allowed_packet. 0 means to use the default size.")
+
 	return cmd, nil
 }
diff --git a/pkg/database/connection.go b/pkg/database/connection.go
index 14bd1e1..20602a8 100644
--- a/pkg/database/connection.go
+++ b/pkg/database/connection.go
@@ -9,11 +9,12 @@ import (
 )
 
 type Connection struct {
-	User            string
-	Pass            string
-	Host            string
-	Port            int
-	MultiStatements bool
+	User             string
+	Pass             string
+	Host             string
+	Port             int
+	MultiStatements  bool
+	MaxAllowedPacket int
 
 	// holds a connection to the database
 	sql *sql.DB
@@ -36,6 +37,9 @@ func (c *Connection) MySQL() (*sql.DB, error) {
 		config.ParseTime = true
 		config.TLSConfig = "preferred"
 		config.MultiStatements = c.MultiStatements
+		if c.MaxAllowedPacket != 0 {
+			config.MaxAllowedPacket = c.MaxAllowedPacket
+		}
 		dsn := config.FormatDSN()
 		handle, err := sql.Open("mysql", dsn)
 		if err != nil {