The code for this tutorial is available here.
Objectives
This tutorial demonstrates how easy it is to add a new function to TiDB that can be used in a SQL-statement. We want to add a function for creating QRCodes.
TiDB is aiming to be compatible with MySQL. However as the functionality we’re adding doesn’t exist in MySQL this isn’t a real concern.
For reference here is the architecture of a TiDB cluster:
We’re only going to add the function to TiDB (in red in the above image), which means this can’t be pushed down to TiKV or TiFlash, however for this function that’s fine.
The TiDB Development Guide has a page that describes some of this and more.
When adding functions to TiDB you should aim for things that can be merged into upstream TiDB instead of running and maintaining your own fork. Might be good to discuss your plans in a GitHub issue before actually starting to do any work.
Step 1: Adding the function to the parser
For this open parser/ast/functions.go and add this:
// List scalar function names.
const (
...
Qrcode = "qrcode"
...
)Here you need to use the lowercase function name.
Step 2: Register the function class
// funcs holds all registered builtin functions. When new function is added,
// check expression/function_traits.go to see if it should be appended to
// any set there.
var funcs = map[string]functionClass{
...
ast.Qrcode: &qrcodeFunctionClass{baseFunctionClass{ast.Qrcode, 1, 1}},
...
}Here we use the Qrcode that we defined in
parser/ast/functions.go. The
qrcodeFunctionClass is something that we’ll create in the
next step.
Step 3: Create the function class
We could put this in any of the existing
expression/builtin_*.go files, but for this demonstration
we create a new file: expression/builtin_qrcode.go.
In this file we’ll define the function class
qrcodeFunctionClass and the builtinQrcodeSig
function implementation.
The function class looks like this:
type qrcodeFunctionClass struct {
baseFunctionClass
}
func (c *qrcodeFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
if err := c.verifyArgs(args); err != nil {
return nil, err
}
bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETString, types.ETString)
if err != nil {
return nil, err
}
types.SetBinChsClnFlag(bf.tp)
sig := &builtinQrcodeSig{bf}
return sig, nil
}The getFunction() method verifies the arguments and then
helps to find the right function implementation (only one option here)
and takes care of things like setting the character set and collation
flag to binary via types.SetBinChsClnFlag(). If our
function would have multiple argument types (integer, string, etc) or
would have multiple optional arguments this function would have been
more complex.
Note that the arguments for newBaseBuiltinFuncWithTp():
- Sets the context to ctx. - Sets the function name to
c.funcName. - Sets the arguments to args. -
Sets the return type to types.ETString. - Sets the argument
types to a single argument of types.ETString.
And the function implementation looks like this:
type builtinQrcodeSig struct {
baseBuiltinFunc
}
func (b *builtinQrcodeSig) Clone() builtinFunc {
newSig := &builtinQrcodeSig{}
newSig.cloneFrom(&b.baseBuiltinFunc)
return newSig
}
func (b *builtinQrcodeSig) evalString(row chunk.Row) (string, bool, error) {
d, isNull, err := b.args[0].EvalString(b.ctx, row)
if isNull || err != nil {
return d, isNull, err
}
qrc, err := qrcode.New(d)
if err != nil {
fmt.Printf("could not generate QRCode: %v", err)
return "", isNull, err
}
var bb bytes.Buffer
qrbw := &qrBufWriter{&bb}
qrWriter := standard.NewWithWriter(qrbw)
if err = qrc.Save(qrWriter); err != nil {
fmt.Printf("could not save image: %v", err)
return "", isNull, err
}
return qrbw.String(), isNull, nil
}We also add another type that is used in the function implementation:
type qrBufWriter struct {
*bytes.Buffer
}
func (b *qrBufWriter) Close() error {
return nil
}This is because we need a writer that implements the io.WriteCloser
interface.
The evalString() method is where we actually create the
QRCode. We get the first argument and store this in d. Then
we use create a QRCode based on d which we then write to a
buffer and eventually return via qrbw.String().
The full code is available in the repository that is mentioned at the start of this article.
Note that many of the other functions serve as examples of how to implement these.
Step 4: Adding dependencies and buildign
We need to add the library for creating QRCodes to our repository.
go get github.com/yeqown/go-qrcode/v2
go get github.com/yeqown/go-qrcode/writer/standardThis modifies the go.mod and go.sum
files.
Now we can run make server to create the
bin/tidb-server binary.
Just running ./bin/tidb-server after building the binary
should give you a TiDB server that listens on port 4000 and uses the
unistore storage. You can also use
tiup playground and specify the path to the binary that you
just build.
Step 5: The result
If you use a tool like DBeaver or MySQL Workbench and call the new function you get a working QRCode image.
SELECT QRCODE('https://github.com/dveeden/tidb/tree/qrcode')
Conclusion
Adding a new function isn’t too difficult. While a
QRCODE() function in SQL is nice this might not be very
useful as it is probably better to move this functionality to the
application layer.
This function is lacking unittests, which would be required before
actually merging into the master branch of pingcap/tidb on
GitHub.
