aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/clickhouse/src/Functions/UserDefined/UserDefinedSQLObjectsBackup.cpp
blob: 6920e8ce2c2686e4b14427b654d68dac02bfd7a2 (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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
#include <Functions/UserDefined/UserDefinedSQLObjectsBackup.h>

#include <Backups/BackupEntriesCollector.h>
#include <Backups/BackupEntryFromMemory.h>
#include <Backups/IBackup.h>
#include <Backups/IBackupCoordination.h>
#include <Backups/IRestoreCoordination.h>
#include <Backups/RestorerFromBackup.h>
#include <Functions/UserDefined/IUserDefinedSQLObjectsLoader.h>
#include <Functions/UserDefined/UserDefinedSQLObjectType.h>
#include <Interpreters/Context.h>
#include <Parsers/ParserCreateFunctionQuery.h>
#include <Parsers/parseQuery.h>
#include <Parsers/queryToString.h>
#include <Common/escapeForFileName.h>


namespace DB
{

namespace ErrorCodes
{
    extern const int CANNOT_RESTORE_TABLE;
}


void backupUserDefinedSQLObjects(
    BackupEntriesCollector & backup_entries_collector,
    const String & data_path_in_backup,
    UserDefinedSQLObjectType object_type,
    const std::vector<std::pair<String, ASTPtr>> & objects)
{
    std::vector<std::pair<String, BackupEntryPtr>> backup_entries;
    backup_entries.reserve(objects.size());
    for (const auto & [object_name, create_object_query] : objects)
        backup_entries.emplace_back(
            escapeForFileName(object_name) + ".sql", std::make_shared<BackupEntryFromMemory>(queryToString(create_object_query)));

    auto context = backup_entries_collector.getContext();
    const auto & loader = context->getUserDefinedSQLObjectsLoader();

    if (!loader.isReplicated())
    {
        fs::path data_path_in_backup_fs{data_path_in_backup};
        for (const auto & [file_name, entry] : backup_entries)
            backup_entries_collector.addBackupEntry(data_path_in_backup_fs / file_name, entry);
        return;
    }

    String replication_id = loader.getReplicationID();

    auto backup_coordination = backup_entries_collector.getBackupCoordination();
    backup_coordination->addReplicatedSQLObjectsDir(replication_id, object_type, data_path_in_backup);

    // On the stage of running post tasks, all directories will already be added to the backup coordination object.
    // They will only be returned for one of the hosts below, for the rest an empty list.
    // See also BackupCoordinationReplicatedSQLObjects class.
    backup_entries_collector.addPostTask(
        [my_backup_entries = std::move(backup_entries),
         my_replication_id = std::move(replication_id),
         object_type,
         &backup_entries_collector,
         backup_coordination]
        {
            auto dirs = backup_coordination->getReplicatedSQLObjectsDirs(my_replication_id, object_type);

            for (const auto & dir : dirs)
            {
                fs::path dir_fs{dir};
                for (const auto & [file_name, entry] : my_backup_entries)
                {
                    backup_entries_collector.addBackupEntry(dir_fs / file_name, entry);
                }
            }
        });
}


std::vector<std::pair<String, ASTPtr>>
restoreUserDefinedSQLObjects(RestorerFromBackup & restorer, const String & data_path_in_backup, UserDefinedSQLObjectType object_type)
{
    auto context = restorer.getContext();
    const auto & loader = context->getUserDefinedSQLObjectsLoader();

    if (loader.isReplicated() && !restorer.getRestoreCoordination()->acquireReplicatedSQLObjects(loader.getReplicationID(), object_type))
        return {}; /// Other replica is already restoring user-defined SQL objects.

    auto backup = restorer.getBackup();
    fs::path data_path_in_backup_fs{data_path_in_backup};

    Strings filenames = backup->listFiles(data_path_in_backup);
    if (filenames.empty())
        return {}; /// Nothing to restore.

    for (const auto & filename : filenames)
    {
        if (!filename.ends_with(".sql"))
        {
            throw Exception(
                ErrorCodes::CANNOT_RESTORE_TABLE,
                "Cannot restore user-defined SQL objects: File name {} doesn't have the extension .sql",
                String{data_path_in_backup_fs / filename});
        }
    }

    std::vector<std::pair<String, ASTPtr>> res;

    for (const auto & filename : filenames)
    {
        String escaped_object_name = filename.substr(0, filename.length() - strlen(".sql"));
        String object_name = unescapeForFileName(escaped_object_name);

        String filepath = data_path_in_backup_fs / filename;
        auto in = backup->readFile(filepath);
        String statement_def;
        readStringUntilEOF(statement_def, *in);

        ASTPtr ast;

        switch (object_type)
        {
            case UserDefinedSQLObjectType::Function:
            {
                ParserCreateFunctionQuery parser;
                ast = parseQuery(
                    parser,
                    statement_def.data(),
                    statement_def.data() + statement_def.size(),
                    "in file " + filepath + " from backup " + backup->getNameForLogging(),
                    0,
                    context->getSettingsRef().max_parser_depth);
                break;
            }
        }

        res.emplace_back(std::move(object_name), ast);
    }

    return res;
}

}