pg_timetable/internal/pgengine/copy.go at 0a42d6bf9a76266521b3e37f48304ba18dd571e3 · cybertec-postgresql/pg_timetable · GitHub

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
package pgengine

import (
	"context"
	"errors"
	"os"
	"os/exec"
)

// CopyToFile copies data from database into local file using COPY format specified by sql
func (pge *PgEngine) CopyToFile(ctx context.Context, filename string, sql string) (int64, error) {
	dbconn, err := pge.ConfigDb.Acquire(ctx)
	if err != nil {
		return -1, err
	}
	defer dbconn.Release()
	f, err := os.Create(filename)
	defer func() { _ = f.Close() }()
	if err != nil {
		return -1, err
	}
	res, err := dbconn.Conn().PgConn().CopyTo(ctx, f, sql)
	return res.RowsAffected(), err
}

// CopyFromFile copies data from local file into database using COPY format specified by sql
func (pge *PgEngine) CopyFromFile(ctx context.Context, filename string, sql string) (int64, error) {
	dbconn, err := pge.ConfigDb.Acquire(ctx)
	if err != nil {
		return -1, err
	}
	defer dbconn.Release()
	f, err := os.Open(filename)
	defer func() { _ = f.Close() }()
	if err != nil {
		return -1, err
	}
	res, err := dbconn.Conn().PgConn().CopyFrom(ctx, f, sql)
	return res.RowsAffected(), err
}

// CopyToProgram copies data from database to the standard input of the command using COPY format specified by sql
func (pge *PgEngine) CopyToProgram(ctx context.Context, sql string, cmd string, args ...string) (int64, error) {
	dbconn, err := pge.ConfigDb.Acquire(ctx)
	if err != nil {
		return -1, err
	}
	defer dbconn.Release()
	c := exec.CommandContext(ctx, cmd, args...)
	inPipe, err := c.StdinPipe()
	if err != nil {
		return -1, err
	}
	if err := c.Start(); err != nil {
		return -1, err
	}
	res, sqlErr := dbconn.Conn().PgConn().CopyTo(ctx, inPipe, sql)
	_ = inPipe.Close()
	cmdError := c.Wait()
	return res.RowsAffected(), errors.Join(sqlErr, cmdError)
}

// CopyFromProgram copies data from the standard output of the command into database using COPY format specified by sql
func (pge *PgEngine) CopyFromProgram(ctx context.Context, sql string, cmd string, args ...string) (int64, error) {
	dbconn, err := pge.ConfigDb.Acquire(ctx)
	if err != nil {
		return -1, err
	}
	defer dbconn.Release()
	c := exec.CommandContext(ctx, cmd, args...)
	outPipe, err := c.StdoutPipe()
	if err != nil {
		return -1, err
	}
	if err := c.Start(); err != nil {
		return -1, err
	}
	res, err := dbconn.Conn().PgConn().CopyFrom(ctx, outPipe, sql)
	waitErr := c.Wait()
	return res.RowsAffected(), errors.Join(waitErr, err)
}