-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
122 lines (102 loc) · 3.11 KB
/
main.go
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
package main
import (
"context"
"fmt"
"log"
"path"
"time"
"github.com/TheAschr/pgseeder"
"github.com/TheAschr/pgseeder/examples/basic/seeders"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/vbauerster/mpb/v8"
)
const dataDir = "./data"
func main() {
ctx := context.Background()
pool, err := pgxpool.New(ctx, "postgresql://user:password@localhost:5435/example")
if err != nil {
log.Fatalf("Failed to create connection pool: %v", err)
}
if err := initDb(pool); err != nil {
log.Fatalf("Failed to initialize database: %v", err)
}
p := mpb.NewWithContext(ctx)
s := pgseeder.New(pool, pgseeder.WithProgress(p))
startTime := time.Now()
if err := s.Run(ctx, []pgseeder.Config{
seeders.NewUsers(
path.Join(dataDir, "users.gz"),
nil,
),
seeders.NewUsStates(
path.Join(dataDir, "us-states.gz"),
[]pgseeder.Config{
seeders.NewUsCounties(path.Join(dataDir, "us-counties.gz"), nil),
},
),
seeders.NewEsriLandformPolygons(path.Join(dataDir, "esri-landform-polygons.gz"), nil),
}); err != nil {
log.Fatalf("Failed to seed database: %v", err)
}
fmt.Printf("Completed in %v\n", time.Since(startTime).Round(time.Millisecond))
}
func initDb(pool *pgxpool.Pool) error {
if _, err := pool.Exec(context.Background(), `
CREATE TABLE IF NOT EXISTS "User" (
"id" INT NOT NULL,
"name" TEXT NOT NULL,
CONSTRAINT "User_pkey" PRIMARY KEY ("id")
)`); err != nil {
return fmt.Errorf("failed to create user table: %w", err)
}
if _, err := pool.Exec(context.Background(), `
CREATE TABLE IF NOT EXISTS "EsriLandformPolygon" (
"id" TEXT NOT NULL,
"featureCodeId" INTEGER NOT NULL,
"gazId" INTEGER NOT NULL,
"name" TEXT NOT NULL,
"geoJSON" JSONB NOT NULL,
CONSTRAINT "EsriLandformPolygon_pkey" PRIMARY KEY ("id")
)`); err != nil {
return fmt.Errorf("failed to create esri landform polygon table: %w", err)
}
if _, err := pool.Exec(context.Background(), `
CREATE TABLE IF NOT EXISTS "UsState" (
"id" TEXT NOT NULL,
"fipsCode" TEXT NOT NULL,
"alpha" TEXT NOT NULL,
"name" TEXT NOT NULL,
"shapeGeoJSON" JSONB NOT NULL,
"districtOfColumbiaId" TEXT,
CONSTRAINT "UsState_pkey" PRIMARY KEY ("id")
)`); err != nil {
return fmt.Errorf("failed to create us state table: %w", err)
}
if _, err := pool.Exec(context.Background(), `
CREATE TABLE IF NOT EXISTS "UsCounty" (
"id" TEXT NOT NULL,
"stcoFipsCode" TEXT NOT NULL,
"shortName" TEXT NOT NULL,
"longName" TEXT NOT NULL,
"deprecated" BOOLEAN NOT NULL DEFAULT false,
"shapeGeoJSON" JSONB NOT NULL,
"stateId" TEXT,
"districtOfColumbiaId" TEXT,
"territoryId" TEXT,
CONSTRAINT "UsCounty_pkey" PRIMARY KEY ("id")
)`); err != nil {
return fmt.Errorf("failed to create us county table: %w", err)
}
if _, err := pool.Exec(context.Background(), `
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'UsCounty_stateId_fkey') THEN
ALTER TABLE "UsCounty" ADD CONSTRAINT "UsCounty_stateId_fkey" FOREIGN KEY ("stateId") REFERENCES "UsState"("id") ON DELETE SET NULL ON UPDATE CASCADE;
END IF;
END;
$$;
`); err != nil {
return fmt.Errorf("failed to create us state reference key: %w", err)
}
return nil
}