ClickHouse/src/Functions/generateULID.cpp

#include "config.h"

#if USE_ULID

#include <Columns/ColumnFixedString.h>
#include <DataTypes/DataTypeFixedString.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
#include <Functions/IFunction.h>
#include <Interpreters/Context.h>

#include <ulid.h>


namespace DB
{

namespace ErrorCodes
{
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
}

class FunctionGenerateULID : public IFunction
{
public:
    static constexpr size_t ULID_LENGTH = 26;

    static constexpr auto name = "generateULID";

    static FunctionPtr create(ContextPtr /*context*/)
    {
        return std::make_shared<FunctionGenerateULID>();
    }

    String getName() const override { return name; }

    size_t getNumberOfArguments() const override { return 0; }

    bool isVariadic() const override { return true; }
    bool isDeterministic() const override { return false; }
    bool isDeterministicInScopeOfQuery() const override { return false; }
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }

    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
    {
        if (arguments.size() > 1)
            throw Exception(
                ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
                "Number of arguments for function {} doesn't match: passed {}, should be 0 or 1.",
                getName(), arguments.size());

        return std::make_shared<DataTypeFixedString>(ULID_LENGTH);
    }

    bool useDefaultImplementationForConstants() const override { return true; }

    ColumnPtr executeImpl(const ColumnsWithTypeAndName & /*arguments*/, const DataTypePtr &, size_t input_rows_count) const override
    {
        auto col_res = ColumnFixedString::create(ULID_LENGTH);
        auto & vec_res = col_res->getChars();

        vec_res.resize(input_rows_count * ULID_LENGTH);

        ulid_generator generator;
        ulid_generator_init(&generator, 0);

        for (size_t offset = 0, size = vec_res.size(); offset < size; offset += ULID_LENGTH)
            ulid_generate(&generator, reinterpret_cast<char *>(&vec_res[offset]));

        return col_res;
    }
};


REGISTER_FUNCTION(GenerateULID)
{
    factory.registerFunction<FunctionGenerateULID>();
}

}

#endif