File: fastpath.go

package info (click to toggle)
golang-github-jackc-pgx 3.6.2-2
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 1,940 kB
  • sloc: sh: 101; makefile: 9; sql: 8
file content (119 lines) | stat: -rw-r--r-- 2,692 bytes parent folder | download
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
package pgx

import (
	"encoding/binary"

	"github.com/jackc/pgx/pgio"
	"github.com/jackc/pgx/pgproto3"
	"github.com/jackc/pgx/pgtype"
)

func newFastpath(cn *Conn) *fastpath {
	return &fastpath{cn: cn, fns: make(map[string]pgtype.OID)}
}

type fastpath struct {
	cn  *Conn
	fns map[string]pgtype.OID
}

func (f *fastpath) functionOID(name string) pgtype.OID {
	return f.fns[name]
}

func (f *fastpath) addFunction(name string, oid pgtype.OID) {
	f.fns[name] = oid
}

func (f *fastpath) addFunctions(rows *Rows) error {
	for rows.Next() {
		var name string
		var oid pgtype.OID
		if err := rows.Scan(&name, &oid); err != nil {
			return err
		}
		f.addFunction(name, oid)
	}
	return rows.Err()
}

type fpArg []byte

func fpIntArg(n int32) fpArg {
	res := make([]byte, 4)
	binary.BigEndian.PutUint32(res, uint32(n))
	return res
}

func fpInt64Arg(n int64) fpArg {
	res := make([]byte, 8)
	binary.BigEndian.PutUint64(res, uint64(n))
	return res
}

func (f *fastpath) Call(oid pgtype.OID, args []fpArg) (res []byte, err error) {
	if err := f.cn.ensureConnectionReadyForQuery(); err != nil {
		return nil, err
	}

	buf := f.cn.wbuf
	buf = append(buf, 'F') // function call
	sp := len(buf)
	buf = pgio.AppendInt32(buf, -1)

	buf = pgio.AppendInt32(buf, int32(oid))       // function object id
	buf = pgio.AppendInt16(buf, 1)                // # of argument format codes
	buf = pgio.AppendInt16(buf, 1)                // format code: binary
	buf = pgio.AppendInt16(buf, int16(len(args))) // # of arguments
	for _, arg := range args {
		buf = pgio.AppendInt32(buf, int32(len(arg))) // length of argument
		buf = append(buf, arg...)                    // argument value
	}
	buf = pgio.AppendInt16(buf, 1) // response format code (binary)
	pgio.SetInt32(buf[sp:], int32(len(buf[sp:])))

	if _, err := f.cn.conn.Write(buf); err != nil {
		return nil, err
	}

	f.cn.pendingReadyForQueryCount++

	for {
		msg, err := f.cn.rxMsg()
		if err != nil {
			return nil, err
		}
		switch msg := msg.(type) {
		case *pgproto3.FunctionCallResponse:
			res = make([]byte, len(msg.Result))
			copy(res, msg.Result)
		case *pgproto3.ReadyForQuery:
			f.cn.rxReadyForQuery(msg)
			// done
			return res, err
		default:
			if err := f.cn.processContextFreeMsg(msg); err != nil {
				return nil, err
			}
		}
	}
}

func (f *fastpath) CallFn(fn string, args []fpArg) ([]byte, error) {
	return f.Call(f.functionOID(fn), args)
}

func fpInt32(data []byte, err error) (int32, error) {
	if err != nil {
		return 0, err
	}
	n := int32(binary.BigEndian.Uint32(data))
	return n, nil
}

func fpInt64(data []byte, err error) (int64, error) {
	if err != nil {
		return 0, err
	}
	return int64(binary.BigEndian.Uint64(data)), nil
}