clickhouse-go/clickhouse_rows.go at main · cloudquery/clickhouse-go · 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
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
package clickhouse

import (
	"database/sql"
	"io"

	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
)

type rows struct {
	err       error
	row       int
	block     *proto.Block
	totals    *proto.Block
	errors    chan error
	stream    chan *proto.Block
	columns   []string
	structMap *structMap
}

func (r *rows) Next() (result bool) {
	defer func() {
		if !result {
			r.Close()
		}
	}()
	if r.block == nil {
		return false
	}
next:
	if r.row >= r.block.Rows() {
		if r.stream == nil {
			return false
		}
		select {
		case err := <-r.errors:
			if err != nil {
				r.err = err
				return false
			}
		case block := <-r.stream:
			if block == nil {
				return false
			}
			if block.Packet == proto.ServerTotals {
				r.row, r.block, r.totals = 0, nil, block
				return false
			}
			r.row, r.block = 0, block
		}
		goto next
	}
	r.row++
	return r.row <= r.block.Rows()
}

func (r *rows) Scan(dest ...any) error {
	if r.block == nil || (r.row == 0 && r.row >= r.block.Rows()) { // call without next when result is empty
		return io.EOF
	}
	return scan(r.block, r.row, dest...)
}

func (r *rows) ScanStruct(dest any) error {
	values, err := r.structMap.Map("ScanStruct", r.columns, dest, true)
	if err != nil {
		return err
	}
	return r.Scan(values...)
}

func (r *rows) Totals(dest ...any) error {
	if r.totals == nil {
		return sql.ErrNoRows
	}
	return scan(r.totals, 1, dest...)
}

func (r *rows) Columns() []string {
	return r.columns
}

func (r *rows) Close() error {
	if r.errors == nil && r.stream == nil {
		return r.err
	}

	if r.errors == nil {
		for range r.stream {
		}
		return nil
	}

	if r.stream == nil {
		for err := range r.errors {
			r.err = err
		}
		return r.err
	}

	errorsClosed := false
	streamClosed := false
	for {
		select {
		case _, ok := <-r.stream:
			if !ok {
				streamClosed = true
			}
		case err, ok := <-r.errors:
			if err != nil {
				r.err = err
			}
			if !ok {
				errorsClosed = true
			}
		}

		if errorsClosed && streamClosed {
			return r.err
		}
	}
}

func (r *rows) Err() error {
	return r.err
}

type row struct {
	err  error
	rows *rows
}

func (r *row) Err() error {
	return r.err
}

func (r *row) ScanStruct(dest any) error {
	if r.err != nil {
		return r.err
	}
	values, err := r.rows.structMap.Map("ScanStruct", r.rows.columns, dest, true)
	if err != nil {
		return err
	}
	return r.Scan(values...)
}

func (r *row) Scan(dest ...any) error {
	if r.err != nil {
		return r.err
	}
	if !r.rows.Next() {
		r.rows.Close()
		if err := r.rows.Err(); err != nil {
			return err
		}
		return sql.ErrNoRows
	}
	if err := r.rows.Scan(dest...); err != nil {
		return err
	}
	return r.rows.Close()
}