forked from zeromicro/go-zero
/
sqlconn.go
162 lines (130 loc) · 4.64 KB
/
sqlconn.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
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
// copy from core/stores/sqlx/sqlconn.go
package mocksql
import (
"context"
"database/sql"
"github.com/zeromicro/go-zero/core/stores/sqlx"
)
type (
// MockConn defines a mock connection instance for mysql
MockConn struct {
db *sql.DB
}
statement struct {
stmt *sql.Stmt
}
)
// NewMockConn creates an instance for MockConn
func NewMockConn(db *sql.DB) *MockConn {
return &MockConn{db: db}
}
// Exec executes sql and returns the result
func (conn *MockConn) Exec(query string, args ...any) (sql.Result, error) {
return exec(conn.db, query, args...)
}
// ExecCtx executes sql and returns the result
func (conn *MockConn) ExecCtx(_ context.Context, query string, args ...any) (sql.Result, error) {
return exec(conn.db, query, args...)
}
// Prepare executes sql by sql.DB
func (conn *MockConn) Prepare(query string) (sqlx.StmtSession, error) {
st, err := conn.db.Prepare(query)
return statement{stmt: st}, err
}
// PrepareCtx executes sql by sql.DB
func (conn *MockConn) PrepareCtx(_ context.Context, query string) (sqlx.StmtSession, error) {
return conn.Prepare(query)
}
// QueryRow executes sql and returns a query row
func (conn *MockConn) QueryRow(v any, q string, args ...any) error {
return query(conn.db, func(rows *sql.Rows) error {
return unmarshalRow(v, rows, true)
}, q, args...)
}
// QueryRowCtx executes sql and returns a query row
func (conn *MockConn) QueryRowCtx(_ context.Context, v any, query string, args ...any) error {
return conn.QueryRow(v, query, args...)
}
// QueryRowPartial executes sql and returns a partial query row
func (conn *MockConn) QueryRowPartial(v any, q string, args ...any) error {
return query(conn.db, func(rows *sql.Rows) error {
return unmarshalRow(v, rows, false)
}, q, args...)
}
// QueryRowPartialCtx executes sql and returns a partial query row
func (conn *MockConn) QueryRowPartialCtx(_ context.Context, v any, query string, args ...any) error {
return conn.QueryRowPartial(v, query, args...)
}
// QueryRows executes sql and returns query rows
func (conn *MockConn) QueryRows(v any, q string, args ...any) error {
return query(conn.db, func(rows *sql.Rows) error {
return unmarshalRows(v, rows, true)
}, q, args...)
}
// QueryRowsCtx executes sql and returns query rows
func (conn *MockConn) QueryRowsCtx(_ context.Context, v any, query string, args ...any) error {
return conn.QueryRows(v, query, args...)
}
// QueryRowsPartial executes sql and returns partial query rows
func (conn *MockConn) QueryRowsPartial(v any, q string, args ...any) error {
return query(conn.db, func(rows *sql.Rows) error {
return unmarshalRows(v, rows, false)
}, q, args...)
}
// QueryRowsPartialCtx executes sql and returns partial query rows
func (conn *MockConn) QueryRowsPartialCtx(_ context.Context, v any, query string, args ...any) error {
return conn.QueryRowsPartial(v, query, args...)
}
// RawDB returns the underlying sql.DB.
func (conn *MockConn) RawDB() (*sql.DB, error) {
return conn.db, nil
}
// Transact is the implemention of sqlx.SqlConn, nothing to do
func (conn *MockConn) Transact(func(session sqlx.Session) error) error {
return nil
}
// TransactCtx is the implemention of sqlx.SqlConn, nothing to do
func (conn *MockConn) TransactCtx(ctx context.Context, fn func(context.Context, sqlx.Session) error) error {
return nil
}
func (s statement) Close() error {
return s.stmt.Close()
}
func (s statement) Exec(args ...any) (sql.Result, error) {
return execStmt(s.stmt, args...)
}
func (s statement) ExecCtx(_ context.Context, args ...any) (sql.Result, error) {
return s.Exec(args...)
}
func (s statement) QueryRow(v any, args ...any) error {
return queryStmt(s.stmt, func(rows *sql.Rows) error {
return unmarshalRow(v, rows, true)
}, args...)
}
func (s statement) QueryRowCtx(_ context.Context, v any, args ...any) error {
return s.QueryRow(v, args...)
}
func (s statement) QueryRowPartial(v any, args ...any) error {
return queryStmt(s.stmt, func(rows *sql.Rows) error {
return unmarshalRow(v, rows, false)
}, args...)
}
func (s statement) QueryRowPartialCtx(_ context.Context, v any, args ...any) error {
return s.QueryRowPartial(v, args...)
}
func (s statement) QueryRows(v any, args ...any) error {
return queryStmt(s.stmt, func(rows *sql.Rows) error {
return unmarshalRows(v, rows, true)
}, args...)
}
func (s statement) QueryRowsCtx(_ context.Context, v any, args ...any) error {
return s.QueryRows(v, args...)
}
func (s statement) QueryRowsPartial(v any, args ...any) error {
return queryStmt(s.stmt, func(rows *sql.Rows) error {
return unmarshalRows(v, rows, false)
}, args...)
}
func (s statement) QueryRowsPartialCtx(_ context.Context, v any, args ...any) error {
return s.QueryRowsPartial(v, args...)
}