Add optional default_port argument for the port() function

docs/en/sql-reference/functions/url-functions.md

@@ -117,9 +117,9 @@ Returns the part of the domain that includes top-level subdomains up to the “f
 
 For example, `cutToFirstSignificantSubdomain('https://news.yandex.com.tr/') = 'yandex.com.tr'`.
 
-### port(URL) {#port}
+### port(URL[, default_port = 0]) {#port}
 
-Returns the port or zero if there is no port in the URL (or in case of validation error).
+Returns the port or `default_port` if there is no port in the URL (or in case of validation error).
 
 ### path {#path}
 

src/Functions/URL/port.cpp

@@ -1,6 +1,10 @@
 #include <Functions/FunctionFactory.h>
-#include <Functions/FunctionStringOrArrayToT.h>
+#include <Functions/IFunctionImpl.h>
 #include <Common/StringUtils/StringUtils.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <Columns/ColumnsNumber.h>
+#include <Columns/ColumnArray.h>
+#include <Columns/ColumnConst.h>
 #include "domain.h"
 
 
@@ -9,83 +13,113 @@ namespace DB
 
 namespace ErrorCodes
 {
+    extern const int ILLEGAL_COLUMN;
     extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
 }
 
-struct ExtractPort
+struct FunctionPort : public IFunction
 {
     static constexpr auto name = "port";
-    static constexpr auto is_fixed_to_constant = true;
+    static FunctionPtr create(const Context &) { return std::make_shared<FunctionPort>(); }
 
-    static void vector(const ColumnString::Chars & data, const ColumnString::Offsets & offsets, PaddedPODArray<UInt16> & res)
+    String getName() const override { return name; }
+    bool isVariadic() const override { return true; }
+    size_t getNumberOfArguments() const override { return 0; }
+    bool useDefaultImplementationForConstants() const override { return true; }
+    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {1}; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        if (arguments.size() != 1 && arguments.size() != 2)
+            throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
+                            + std::to_string(arguments.size()) + ", should be 1 or 2",
+                            ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+
+        if (!WhichDataType(arguments[0].type).isString())
+            throw Exception("Illegal type " + arguments[0].type->getName() + " of first argument of function " + getName() + ". Must be String.",
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        if (arguments.size() == 2 && !WhichDataType(arguments[1].type).isUInt16())
+            throw Exception("Illegal type " + arguments[1].type->getName() + " of second argument of function " + getName() + ". Must be UInt16.",
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        return std::make_shared<DataTypeUInt16>();
+    }
+
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t) override
+    {
+        UInt16 default_port = 0;
+        if (arguments.size() == 2)
+        {
+            const auto * port_column = checkAndGetColumn<ColumnConst>(block.getByPosition(arguments[1]).column.get());
+            if (!port_column)
+                throw Exception("Second argument for function " + getName() + " must be constant UInt16", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+            default_port = port_column->getValue<UInt16>();
+        }
+
+        const ColumnPtr url_column = block.getByPosition(arguments[0]).column;
+        if (const ColumnString * url_strs = checkAndGetColumn<ColumnString>(url_column.get()))
+        {
+            auto col_res = ColumnVector<UInt16>::create();
+            typename ColumnVector<UInt16>::Container & vec_res = col_res->getData();
+            vec_res.resize(url_column->size());
+
+            vector(default_port, url_strs->getChars(), url_strs->getOffsets(), vec_res);
+            block.getByPosition(result).column = std::move(col_res);
+        }
+        else
+            throw Exception(
+                "Illegal column " + block.getByPosition(arguments[0]).column->getName() + " of argument of function " + getName(),
+                ErrorCodes::ILLEGAL_COLUMN);
+}
+
+private:
+    static void vector(UInt16 default_port, const ColumnString::Chars & data, const ColumnString::Offsets & offsets, PaddedPODArray<UInt16> & res)
     {
         size_t size = offsets.size();
 
         ColumnString::Offset prev_offset = 0;
         for (size_t i = 0; i < size; ++i)
         {
-            res[i] = parse(data, prev_offset, offsets[i] - prev_offset - 1);
+            res[i] = extractPort(default_port, data, prev_offset, offsets[i] - prev_offset - 1);
             prev_offset = offsets[i];
         }
-    }
+}
 
-    static void vectorFixedToConstant(const ColumnString::Chars & data, size_t n, UInt16 & res) { res = parse(data, 0, n); }
-
-    static void vectorFixedToVector(const ColumnString::Chars & data, size_t n, PaddedPODArray<UInt16> & res)
-    {
-        size_t size = data.size() / n;
-
-        for (size_t i = 0; i < size; ++i)
-        {
-            res[i] = parse(data, i * n, n);
-        }
-    }
-
-    [[noreturn]] static void array(const ColumnString::Offsets &, PaddedPODArray<UInt16> &)
-    {
-        throw Exception("Cannot apply function port to Array argument", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
-    }
-
-private:
-    static UInt16 parse(const ColumnString::Chars & buf, size_t offset, size_t size)
+    static UInt16 extractPort(UInt16 default_port, const ColumnString::Chars & buf, size_t offset, size_t size)
     {
         const char * p = reinterpret_cast<const char *>(&buf[0]) + offset;
         const char * end = p + size;
 
         StringRef host = getURLHost(p, size);
         if (!host.size)
-            return 0;
+            return default_port;
         if (host.size == size)
-            return 0;
+            return default_port;
 
         p = host.data + host.size;
         if (*p++ != ':')
-            return 0;
+            return default_port;
 
-        Int64 port = 0;
+        Int64 port = default_port;
         while (p < end)
         {
             if (*p == '/')
                 break;
             if (!isNumericASCII(*p))
-                return 0;
+                return default_port;
 
             port = (port * 10) + (*p - '0');
             if (port < 0 || port > UInt16(-1))
-                return 0;
+                return default_port;
             ++p;
         }
         return port;
     }
 };
 
-struct NamePort
-{
-    static constexpr auto name = "port";
-};
-
-using FunctionPort = FunctionStringOrArrayToT<ExtractPort, NamePort, UInt16>;
-
 void registerFunctionPort(FunctionFactory & factory)
 {
     factory.registerFunction<FunctionPort>();

tests/queries/0_stateless/01284_port.reference

@@ -10,6 +10,9 @@ hostname
 80
 80
 80
+default-port
+80
+80
 ipv6
 0
 0

tests/queries/0_stateless/01284_port.sql

@@ -11,6 +11,15 @@ select port('http://foobar.com:80/');
 select port('//foobar.com:80/');
 select port('foobar.com:80');
 
+select 'default-port';
+select port('http://127.0.0.1/', toUInt16(80));
+select port('http://foobar.com/', toUInt16(80));
+
+-- unsupported
+/* ILLEGAL_TYPE_OF_ARGUMENT */ select port(toFixedString('', 1)); -- { serverError 43; }
+/* ILLEGAL_TYPE_OF_ARGUMENT */ select port('', 1); -- { serverError 43; }
+/* NUMBER_OF_ARGUMENTS_DOESNT_MATCH */ select port('', 1, 1); -- { serverError 42; }
+
 --
 -- Known limitations of domain() (getURLHost())
 --