171 lines
3.7 KiB
Go
171 lines
3.7 KiB
Go
package database
|
|
|
|
import (
|
|
"encoding/hex"
|
|
"fmt"
|
|
"io"
|
|
"regexp"
|
|
"strings"
|
|
)
|
|
|
|
type DumpColumn struct {
|
|
Name string
|
|
Type string
|
|
}
|
|
|
|
type DumpRows interface {
|
|
Columns() ([]DumpColumn, error)
|
|
Next() bool
|
|
Values() ([][]byte, error)
|
|
Err() error
|
|
Close() error
|
|
}
|
|
|
|
type DumpSource interface {
|
|
Tables() ([]string, error)
|
|
CreateStatement(table string) (string, error)
|
|
Rows(table string) (DumpRows, error)
|
|
Close() error
|
|
}
|
|
|
|
var dumpPreamble = []string{
|
|
"SET NAMES utf8mb4;",
|
|
"SET sql_mode = 'NO_AUTO_VALUE_ON_ZERO';",
|
|
"SET foreign_key_checks = 0;",
|
|
}
|
|
|
|
var dumpPostamble = []string{
|
|
"SET foreign_key_checks = 1;",
|
|
}
|
|
|
|
var plainNumber = regexp.MustCompile(`^-?[0-9]+(\.[0-9]+)?([eE][+-]?[0-9]+)?$`)
|
|
|
|
var textEscapes = strings.NewReplacer(
|
|
`\`, `\\`,
|
|
`'`, `\'`,
|
|
"\x00", `\0`,
|
|
"\n", `\n`,
|
|
"\r", `\r`,
|
|
"\x1a", `\Z`,
|
|
)
|
|
|
|
func WriteDump(output io.Writer, source DumpSource) error {
|
|
if err := writeStatements(output, dumpPreamble); err != nil {
|
|
return err
|
|
}
|
|
tables, err := source.Tables()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, table := range tables {
|
|
if err := writeTable(output, source, table); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return writeStatements(output, dumpPostamble)
|
|
}
|
|
|
|
func writeStatements(output io.Writer, statements []string) error {
|
|
for _, statement := range statements {
|
|
if _, err := fmt.Fprintln(output, statement); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func writeTable(output io.Writer, source DumpSource, table string) error {
|
|
create, err := source.CreateStatement(table)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err := fmt.Fprintf(output, "%s;\n", create); err != nil {
|
|
return err
|
|
}
|
|
rows, err := source.Rows(table)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer rows.Close()
|
|
return writeRows(output, table, rows)
|
|
}
|
|
|
|
func writeRows(output io.Writer, table string, rows DumpRows) error {
|
|
columns, err := rows.Columns()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
insertPrefix := insertStatementPrefix(table, columns)
|
|
for rows.Next() {
|
|
values, err := rows.Values()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err := fmt.Fprintf(output, "%s%s);\n", insertPrefix, joinLiterals(columns, values)); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return rows.Err()
|
|
}
|
|
|
|
func insertStatementPrefix(table string, columns []DumpColumn) string {
|
|
names := make([]string, len(columns))
|
|
for index, column := range columns {
|
|
names[index] = quoteIdentifier(column.Name)
|
|
}
|
|
return fmt.Sprintf("INSERT INTO %s (%s) VALUES (", quoteIdentifier(table), strings.Join(names, ","))
|
|
}
|
|
|
|
func joinLiterals(columns []DumpColumn, values [][]byte) string {
|
|
literals := make([]string, len(columns))
|
|
for index, column := range columns {
|
|
literals[index] = sqlLiteral(column.Type, values[index])
|
|
}
|
|
return strings.Join(literals, ",")
|
|
}
|
|
|
|
func quoteIdentifier(name string) string {
|
|
return "`" + strings.ReplaceAll(name, "`", "``") + "`"
|
|
}
|
|
|
|
func sqlLiteral(databaseType string, value []byte) string {
|
|
if value == nil {
|
|
return "NULL"
|
|
}
|
|
switch {
|
|
case isNumericType(databaseType) && plainNumber.Match(value):
|
|
return string(value)
|
|
case isBinaryType(databaseType):
|
|
return hexLiteral(value)
|
|
default:
|
|
return "'" + textEscapes.Replace(string(value)) + "'"
|
|
}
|
|
}
|
|
|
|
func hexLiteral(value []byte) string {
|
|
if len(value) == 0 {
|
|
return "''"
|
|
}
|
|
return "0x" + hex.EncodeToString(value)
|
|
}
|
|
|
|
func isNumericType(databaseType string) bool {
|
|
upper := strings.ToUpper(databaseType)
|
|
for _, suffix := range []string{"INT", "DECIMAL", "NUMERIC", "FLOAT", "DOUBLE", "YEAR"} {
|
|
if strings.HasSuffix(upper, suffix) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func isBinaryType(databaseType string) bool {
|
|
upper := strings.ToUpper(databaseType)
|
|
for _, suffix := range []string{"BLOB", "BINARY", "BIT", "GEOMETRY"} {
|
|
if strings.HasSuffix(upper, suffix) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|