summaryrefslogtreecommitdiffstats
path: root/yql/essentials/sql/v1/complete/check/check_complete.cpp
blob: 2ed7ccd835ee63084dfa423bbe756e82c1f9eaca (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
#include "check_complete.h"

#include <yql/essentials/sql/v1/complete/sql_complete.h>
#include <yql/essentials/sql/v1/complete/analysis/yql/yql.h>
#include <yql/essentials/sql/v1/complete/name/cluster/static/discovery.h>
#include <yql/essentials/sql/v1/complete/name/object/simple/static/schema.h>
#include <yql/essentials/sql/v1/complete/name/service/cluster/name_service.h>
#include <yql/essentials/sql/v1/complete/name/service/schema/name_service.h>
#include <yql/essentials/sql/v1/complete/name/service/static/name_service.h>
#include <yql/essentials/sql/v1/complete/name/service/union/name_service.h>

#include <yql/essentials/sql/v1/lexer/antlr4_pure/lexer.h>
#include <yql/essentials/sql/v1/lexer/antlr4_pure_ansi/lexer.h>

#include <util/charset/utf8.h>
#include <util/random/random.h>

namespace NSQLComplete {

namespace {

TLexerSupplier MakePureLexerSupplier() {
    NSQLTranslationV1::TLexers lexers;
    lexers.Antlr4Pure = NSQLTranslationV1::MakeAntlr4PureLexerFactory();
    lexers.Antlr4PureAnsi = NSQLTranslationV1::MakeAntlr4PureAnsiLexerFactory();
    return [lexers = std::move(lexers)](bool ansi) {
        return NSQLTranslationV1::MakeLexer(
            lexers, ansi, /* antlr4 = */ true,
            NSQLTranslationV1::ELexerFlavor::Pure);
    };
}

INameService::TPtr MakeClusterNameService(const TYqlContext& ctx) {
    THashSet<TString> clusterSet = ctx.Clusters();
    TVector<TString> clusterVec(begin(clusterSet), end(clusterSet));
    Sort(clusterVec);

    return MakeClusterNameService(MakeStaticClusterDiscovery(std::move(clusterVec)));
}

INameService::TPtr MakeSchemaNameService(const TYqlContext& ctx) {
    THashMap<TString, THashMap<TString, TVector<TFolderEntry>>> fs;
    for (const auto& [cluster, tables] : ctx.TablesByCluster) {
        for (TString table : tables) {
            fs[cluster]["/"].push_back(TFolderEntry{
                .Type = TFolderEntry::Table,
                .Name = std::move(table),
            });
        }
    }

    return MakeSchemaNameService(MakeSimpleSchema(MakeStaticSimpleSchema({.Folders = std::move(fs)})));
}

} // namespace

bool CheckComplete(TStringBuf query, TYqlContext ctx) {
    constexpr size_t Seed = 97651231;
    constexpr size_t Attempts = 64;
    constexpr size_t MaxAttempts = 256;
    SetRandomSeed(Seed);

    auto service = MakeUnionNameService(
        {
            MakeClusterNameService(ctx),
            MakeSchemaNameService(ctx),
        },
        MakeDefaultRanking());

    auto engine = MakeSqlCompletionEngine(
        MakePureLexerSupplier(),
        std::move(service),
        MakeYQLConfiguration());

    for (size_t i = 0, j = 0; i < Attempts && j < MaxAttempts; ++j) {
        size_t pos = RandomNumber<size_t>(query.size() + 1);
        if (pos < query.size() && IsUTF8ContinuationByte(query.at(pos))) {
            continue;
        }

        TCompletionInput input = {
            .Text = query,
            .CursorPosition = pos,
        };

        auto output = engine->Complete(input).ExtractValueSync();
        Y_DO_NOT_OPTIMIZE_AWAY(output);

        i += 1;
    }

    return true;
}

bool CheckComplete(TStringBuf query, NYql::TExprNode::TPtr root, NYql::TExprContext& ctx, NYql::TIssues& issues) try {
    return CheckComplete(query, MakeYqlAnalysis()->Analyze(root, ctx));
} catch (...) {
    issues.AddIssue(FormatCurrentException());
    return false;
}

bool CheckComplete(TStringBuf query, NYql::TAstNode& root, NYql::TIssues& issues) try {
    return CheckComplete(query, MakeYqlAnalysis()->Analyze(root, issues).GetOrElse({}));
} catch (...) {
    issues.AddIssue(FormatCurrentException());
    return false;
}

} // namespace NSQLComplete