diff --git a/ydb/library/workload/tpc_base/tpc_base.cpp b/ydb/library/workload/tpc_base/tpc_base.cpp index 2e2d7272a877..a1a558280dc4 100644 --- a/ydb/library/workload/tpc_base/tpc_base.cpp +++ b/ydb/library/workload/tpc_base/tpc_base.cpp @@ -5,6 +5,7 @@ #include #include #include +#include namespace NYdbWorkload { @@ -17,32 +18,55 @@ void TTpcBaseWorkloadGenerator::PatchQuery(TString& query) const { TString header; switch (Params.GetFloatMode()) { case TTpcBaseWorkloadParams::EFloatMode::FLOAT: - header = NResource::Find("consts.yql"); + header = FilterHeader(NResource::Find("consts.yql"), query); break; case TTpcBaseWorkloadParams::EFloatMode::DECIMAL: - header = NResource::Find("consts_decimal.yql"); + case TTpcBaseWorkloadParams::EFloatMode::DECIMAL_YDB: + header = FilterHeader(NResource::Find("consts_decimal.yql"), query); break; - case TTpcBaseWorkloadParams::EFloatMode::DECIMAL_YDB: { - header = NResource::Find("consts_decimal.yql"); - header.to_lower(); - const TStringBuf dec("decimal("); - auto p = header.find(dec); - while (p != TString::npos) { - p += dec.length(); - const auto q = header.find(')', p); - TVector decParams; - StringSplitter(header.cbegin() + p, q - p).SplitBySet(", ").SkipEmpty().Limit(2).ParseInto(&decParams); - TStringBuilder newDecParams; - newDecParams - << Max(decParams[0], NKikimr::NScheme::DECIMAL_PRECISION) - << "," << Max(decParams[1], NKikimr::NScheme::DECIMAL_SCALE); - header.replace(p, q - p, newDecParams); - p = header.find(dec, q); + } + PatchHeader(header); + SubstGlobal(query, "{% include 'header.sql.jinja' %}", header); +} + +TString TTpcBaseWorkloadGenerator::FilterHeader(TStringBuf header, const TString& query) const { + TStringBuilder result; + for(TStringBuf line; header.ReadLine(line);) { + const auto pos = line.find('='); + if (pos == line.npos) { + continue; + } + const auto name = StripString(line.SubString(0, pos)); + for(auto posInQ = query.find(name); posInQ != query.npos; posInQ = query.find(name, posInQ)) { + posInQ += name.length(); + if (posInQ >= query.length() || !IsAsciiAlnum(query[posInQ]) && query[posInQ] != '_') { + result << line << Endl; + break; } } - break; } - SubstGlobal(query, "{% include 'header.sql.jinja' %}", header); + return result; +} + +void TTpcBaseWorkloadGenerator::PatchHeader(TString& header) const { + if (Params.GetFloatMode() != TTpcBaseWorkloadParams::EFloatMode::DECIMAL_YDB) { + return; + } + header.to_lower(); + const TStringBuf dec("decimal("); + auto p = header.find(dec); + while (p != TString::npos) { + p += dec.length(); + const auto q = header.find(')', p); + TVector decParams; + StringSplitter(header.cbegin() + p, q - p).SplitBySet(", ").SkipEmpty().Limit(2).ParseInto(&decParams); + TStringBuilder newDecParams; + newDecParams + << Max(decParams[0], NKikimr::NScheme::DECIMAL_PRECISION) + << "," << Max(decParams[1], NKikimr::NScheme::DECIMAL_SCALE); + header.replace(p, q - p, newDecParams); + p = header.find(dec, q); + } } void TTpcBaseWorkloadParams::ConfigureOpts(NLastGetopt::TOpts& opts, const ECommandType commandType, int workloadType) { diff --git a/ydb/library/workload/tpc_base/tpc_base.h b/ydb/library/workload/tpc_base/tpc_base.h index 5216f0e90465..3bc1bdd10c26 100644 --- a/ydb/library/workload/tpc_base/tpc_base.h +++ b/ydb/library/workload/tpc_base/tpc_base.h @@ -25,6 +25,8 @@ class TTpcBaseWorkloadGenerator: public TWorkloadGeneratorBase { private: const TTpcBaseWorkloadParams& Params; + TString FilterHeader(TStringBuf header, const TString& query) const; + void PatchHeader(TString& header) const; }; } // namespace NYdbWorkload