Crafting An Extension Framework for MySQL

static char rotate(char c, int n) {
  if (c >= 'a' && c <= 'z') return 'a' + (c - 'a' + n) % 26;
  if (c >= 'A' && c <= 'Z') return 'A' + (c - 'A' + n) % 26;
  return c;
}

mysql> CREATE TABLE msgs (msg ROT13);

mysql> INSERT INTO msgs VALUES ('Alice'), ('Bob'), ('Hello'), ('World');

mysql> SELECT msg FROM msgs;
+-------+
| msg   |
+-------+
| Alice |
| Bob   |
| Hello |
| World |
+-------+

mysql> CREATE TABLE msgs (msg ROT13);

mysql> INSERT INTO msgs VALUES ('Alice'), ('Bob'), ('Hello'), ('World');

mysql> SELECT msg, rot13_ciphertext(msg) FROM msgs;
+-------+-----------------------+
| msg   | rot13_ciphertext(msg) |
+-------+-----------------------+
| Alice | Nyvpr                 |
| Bob   | Obo                   |
| Hello | Uryyb                 |
| World | Jbeyq                 |
+-------+-----------------------+

mysql> SELECT msg, rot13_ciphertext(msg) FROM msgs ORDER BY msg;
+-------+-----------------------+
| msg   | rot13_ciphertext(msg) |
+-------+-----------------------+
| World | Jbeyq                 |
| Alice | Nyvpr                 |
| Bob   | Obo                   |
| Hello | Uryyb                 |
+-------+-----------------------+

static constexpr const char kRot13[] = "rot13";

constexpr auto ROT13 = make_type<kRot13>()
                           .variable_length_type()
                           .from_string<&rot13_from_string>()
                           .to_string<&rot13_to_string>()
                           .compare<&NormalStringCompare>()
                           .build();

VEF_GENERATE_ENTRY_POINTS(make_extension().type(ROT13).func(
    make_func<&rot13_ciphertext>("rot13_ciphertext")
        .returns(STRING)
        .param(ROT13)
        .build()))

void rot13_from_string(std::string_view from, CustomResult out) {
  for (size_t i = 0; i < from.size(); i++)
    out.buffer()[i] = rotate(from[i], 13);
  out.set_length(from.size());
}

void rot13_to_string(CustomArg in, StringResult out) {
  for (size_t i = 0; i < in.value().size(); i++)
    out.buffer()[i] = rotate(in.value()[i], 13);
  out.set_length(in.value().size());
}

// The server won't HEX() or CAST() a custom type, so expose the stored bytes
// for the demo.
void rot13_ciphertext(CustomArg in, StringResult out) {
  auto data = in.value();
  out.set({reinterpret_cast<const char *>(data.data()), data.size()});
}

mysql> CREATE FUNCTION add RETURNS INTEGER SONAME "add_udf.so";

mysql> SELECT add(40, 2);
+------------+
| add(40, 2) |
+------------+
|         42 |
+------------+

mysql> DROP FUNCTION add;

extern "C" bool add_init(UDF_INIT *initid, UDF_ARGS *args, char *message) {
  if (args->arg_count != 2) {
    strcpy(message, "add() requires exactly two arguments");
    return true;
  }
  args->arg_type[0] = args->arg_type[1] = INT_RESULT;
  initid->maybe_null = true;
  return false;
}

extern "C" long long add(UDF_INIT *, UDF_ARGS *args,
                         unsigned char *is_null, unsigned char *) {
  if (!args->args[0] || !args->args[1]) {
    *is_null = 1;
    return 0;
  }
  return *((long long *)args->args[0]) + *((long long *)args->args[1]);
}

void vdf_add(IntArg a, IntArg b, IntResult out) {
  if (a.is_null() || b.is_null()) return out.set_null();
  out.set(a.value() + b.value());
}

VEF_GENERATE_ENTRY_POINTS(
    make_extension()
        .func(make_func<&vdf_add>("add")
                  .returns(INT)
                  .param(INT)
                  .param(INT)
                  .build()))

void add_prerun(PrerunArgs args, PrerunResult out) {
  if (args.size() != 2) {
    out.error("add() requires exactly two arguments");
    return;
  }
  if (!args.type_at(0).is_int() || !args.type_at(1).is_int()) {
    out.error("add() requires integer arguments");
    return;
  }
}

void add(VarArgs args, IntResult out) {
  out.set(args[0].as_int() + args[1].as_int());
}

make_func<&add>("add")
  .returns(INT)
  .varargs()
  .prerun<&add_prerun>()
  .build()

extern "C" bool udf_sum_init(UDF_INIT *initid, UDF_ARGS *args, char *message) {
  if (args->arg_count != 1) {
    strcpy(message, "udf_sum() requires exactly one argument");
    return true;
  }
  args->arg_type[0] = INT_RESULT;
  initid->ptr = (char *)new long long(0);
  return false;
}

extern "C" void udf_sum_deinit(UDF_INIT *initid) {
  delete (long long *)initid->ptr;
}

extern "C" void udf_sum_clear(UDF_INIT *initid, unsigned char *,
                              unsigned char *) {
  *(long long *)initid->ptr = 0;
}

extern "C" void udf_sum_add(UDF_INIT *initid, UDF_ARGS *args,
                            unsigned char *, unsigned char *) {
  if (args->args[0]) *(long long *)initid->ptr += *(long long *)args->args[0];
}

extern "C" long long udf_sum(UDF_INIT *initid, UDF_ARGS *,
                             unsigned char *, unsigned char *) {
  return *(long long *)initid->ptr;
}

using SumState = std::optional<long long>;

void vdf_sum_clear(SumState &state) { state = std::nullopt; }

void vdf_sum_accumulate(SumState &state, IntArg val) {
  if (!val.is_null()) state = state.value_or(0) + val.value();
}

void vdf_sum_result(const SumState &state, IntResult out) {
  if (state) out.set(*state);
  else out.set_null();
}

VEF_GENERATE_ENTRY_POINTS(make_extension().func(
    make_aggregate_func<SumState, &vdf_sum_result>("vdf_sum")
        .returns(INT).param(INT)
        .clear<&vdf_sum_clear>()
        .accumulate<&vdf_sum_accumulate>()
        .build()))

mysql> CREATE TABLE msgs (r13 ROT13, r3 ROT(3));

mysql> INSERT INTO msgs VALUES ('Alice', 'Alice'), ('Bob', 'Bob'),
    ->   ('Hello', 'Hello'), ('World', 'World');

mysql> SELECT r13, rot13_ciphertext(r13), r3, rot_ciphertext(r3) FROM msgs;
+-------+-----------------------+-------+--------------------+
| r13   | rot13_ciphertext(r13) | r3    | rot_ciphertext(r3) |
+-------+-----------------------+-------+--------------------+
| Alice | Nyvpr                 | Alice | Dolfh              |
| Bob   | Obo                   | Bob   | Ere                |
| Hello | Uryyb                 | Hello | Khoor              |
| World | Jbeyq                 | World | Zruog              |
+-------+-----------------------+-------+--------------------+

static constexpr const char kRot[] = "rot";

constexpr auto ROT =
    make_type<kRot>()
        ...()
        .params<RotParams, &RotParams::parse, &RotParams::to_strings>()
        .int_to_params<&rot_int_to_params>()
        .resolve_params<&rot_resolve_params>()
        ...()
        .build();

struct RotParams {
  int n;

  static RotParams parse(const std::map<std::string, std::string> &params) {
    return {atoi(params.at("n").c_str())};
  }

  static void to_strings(const RotParams &p,
                         std::map<std::string, std::string> &out) {
    out["n"] = std::to_string(p.n);
  }
};

// rot(3) is shorthand for rot('n=3').
bool rot_int_to_params(int64_t value,
                       std::map<std::string, std::string> &params, char *) {
  params["n"] = std::to_string(value);
  return false;
}

// Unknown params mean infer them from the string (rot::from_string('...')),
// but a rotated string carries no hint of n.
void rot_from_string(MaybeParams<RotParams> &p, std::string_view from,
                     CustomResult out) {
  if (!p.is_known()) {
    out.error("rot: no rot(n) type here, so no n to rotate by");
    return;
  }
  for (size_t i = 0; i < from.size(); i++)
    out.buffer()[i] = rotate(from[i], p.value().n);
  out.set_length(from.size());
}

void rot_to_string(CustomArgWith<RotParams> in, StringResult out) {
  for (size_t i = 0; i < in.value().size(); i++)
    out.buffer()[i] = rotate(in.value()[i], 26 - in.params().n);
  out.set_length(in.value().size());
}

bool rot_resolve_params(const std::map<std::string, std::string> &params,
                        ResolvedTypeParams *result, char *error_msg) {
  auto it = params.find("n");
  int n = it == params.end() ? -1 : atoi(it->second.c_str());
  if (n < 0 || n > 26) {
    snprintf(error_msg, VEF_MAX_ERROR_LEN, "rot(n) needs 0 <= n <= 26");
    return true;
  }
  return false;
}

mysql> CREATE TABLE bad (msg ROT(27));
ERROR 1064 (42000): rot(n) needs 0 <= n <= 26 near 'ROT(27))' at line 1

mysql> CREATE TABLE both13 (general ROT(13), special ROT13);

mysql> INSERT INTO both13 VALUES ('Hello', 'Hello');

mysql> SELECT rot_ciphertext(general), rot13_ciphertext(special) FROM both13;
+-------------------------+---------------------------+
| rot_ciphertext(general) | rot13_ciphertext(special) |
+-------------------------+---------------------------+
| Uryyb                   | Uryyb                     |
+-------------------------+---------------------------+