Tutorial: Add a QRCODE() function to TiDB

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/standard

This 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.