~louis77/pg2sqlite

ref: bcc3f2acb5d11bc917c96519c8a810fc06058a24 pg2sqlite/main.go -rw-r--r-- 4.3 KiB
bcc3f2ac — Louis Brauer Update README 4 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
/*	pg2sqlite - Migrate tables from PostgreSQL to SQLite
	Copyright (C) 2021  Louis Brauer

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package main

import (
	_ "embed"
	"fmt"
	"github.com/gosuri/uiprogress"
	"github.com/mkideal/cli"
	"log"
	"os"
	"strings"
)

//go:embed VERSION
var Version string

const AppName = "pg2sqlite"
const Copyright = "Copyright © Louis Brauer <louis@brauer.family>"

type stringListDecoder struct {
	List []string
}

func (d *stringListDecoder) Decode(s string) error {
	d.List = strings.Split(s, ",")
	return nil
}

type argT struct {
	cli.Helper
	// Required
	PGURL     string `cli:"*pg-url" usage:"Postgres connection string (i.e. postgres://localhost:5432/mydb)"`
	SLFile    string `cli:"*sqlite-file" usage:"Path to SQLite database file (i.e. mydatabase.db)"`
	Tablename string `cli:"*t,table" usage:"Name of table to export"`
	// Change behaviour
	IgnoreColumns     stringListDecoder `cli:"ignore-columns" usage:"comma-separated list of columns to ignore" default:""`
	DropTableIfExists bool              `cli:"drop-table-if-exists" usage:"DANGER: Drop target table if it already exists" default:"false"`
	// Comfort options
	Confirm bool `cli:"confirm" usage:"Confirm prompts with Y, useful if used in script" default:""`
	Verify  bool `cli:"verify" usage:"Verify that the number of rows inserted into SQLite equals the number of rows loaded from Postgres. In case of failure, exits with status code 2" default:"false"`
}

func (argv *argT) AutoHelp() bool {
	return argv.Help
}

func run(ctx *cli.Context) error {
	argv := ctx.Argv().(*argT)

	if err := ValidatePG(argv.PGURL); err != nil {
		log.Fatal(err)
	}

	if err := ValidateSqlite(argv.SLFile, argv.Tablename, argv.DropTableIfExists); err != nil {
		log.Fatal(err)
	}

	schema, err := FetchSchema(argv.Tablename, argv.IgnoreColumns.List)
	if err != nil {
		log.Fatal(err)
	}

	PrintSchema(schema)

	createTableSQL, err := BuildCreateTableSQL(schema)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Will create SQLite table with the following statement:\n%s\n\n", createTableSQL)

	if !argv.Confirm && !AskYesNo("Does this look ok?") {
		log.Fatal("Cancelled")
	}

	if argv.DropTableIfExists {
		if err := DropTable(schema.Name); err != nil {
			log.Fatalf("Unable to drop target table: %v\n", err)
		}
	}

	if err := CreateTable(createTableSQL); err != nil {
		log.Fatal(err)
	}

	estimatedRows, err := EstimateRows(schema.Name)
	if err != nil {
		log.Fatalln(err)
	}
	fmt.Printf("Estimated row count: %d\n", estimatedRows)

	uiprogress.Start()
	bar := uiprogress.AddBar(int(estimatedRows))
	bar.AppendCompleted()
	bar.PrependElapsed()

	rowChan := make(chan []interface{})
	finished := make(chan bool)
	transferredRows := uint64(0)

	go func() {
		for row := range rowChan {
			if err := InsertRow(schema.Name, row); err != nil {
				log.Fatalln("error inserting a row:", err)
			}
			transferredRows++
			bar.Incr()
		}
		finished <- true
	}()

	go func() {
		if err := LoadData(schema, rowChan); err != nil {
			log.Println("Error while loading data", err)
			finished <- true
		}
	}()

	<-finished
	fmt.Println("Finished.")
	fmt.Println()

	if argv.Verify {
		fmt.Println("Verifying number of rows, this could take a while...")
		rowcount, err := CountRows(schema.Name)
		if err != nil {
			log.Fatalln("Unable to verify rowcount:", err)
		}
		if rowcount != transferredRows {
			log.Println("VERIFICATION FAILED")
			log.Printf("Discrepancy: counted rows: %d, rows in SQLite table: %d", transferredRows, rowcount)
			os.Exit(2)
		}
		fmt.Println("OK, row counts match")
	}

	if err := CloseSqlite(); err != nil {
		log.Println("Unable to close Sqlite database", err)
	}

	return nil
}

func main() {
	cli.SetUsageStyle(cli.DenseManualStyle)
	fmt.Printf("%s v%s\n%s\n\n", AppName, Version, Copyright)

	os.Exit(cli.Run(new(argT), run))
}