Merge pull request #35313 from kssenii/fix-pg-db

Fix possible segfault in DatabasePostgreSQL::getCreateTableQuery
This commit is contained in:
Kruglov Pavel 2022-03-16 12:23:25 +01:00 committed by GitHub
commit ec6968c17e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 21 additions and 6 deletions

View File

@ -406,13 +406,24 @@ ASTPtr DatabasePostgreSQL::getCreateTableQueryImpl(const String & table_name, Co
ASTs storage_children = ast_storage->children;
auto storage_engine_arguments = ast_storage->engine->arguments;
/// Remove extra engine argument (`schema` and `use_table_cache`)
if (storage_engine_arguments->children.size() >= 5)
storage_engine_arguments->children.resize(4);
if (storage_engine_arguments->children.empty())
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Unexpected number of arguments: {}", storage_engine_arguments->children.size());
/// Add table_name to engine arguments
assert(storage_engine_arguments->children.size() >= 2);
storage_engine_arguments->children.insert(storage_engine_arguments->children.begin() + 2, std::make_shared<ASTLiteral>(table_id.table_name));
/// Check for named collection.
if (typeid_cast<ASTIdentifier *>(storage_engine_arguments->children[0].get()))
{
storage_engine_arguments->children.push_back(makeASTFunction("equals", std::make_shared<ASTIdentifier>("table"), std::make_shared<ASTLiteral>(table_id.table_name)));
}
else
{
/// Remove extra engine argument (`schema` and `use_table_cache`)
if (storage_engine_arguments->children.size() >= 5)
storage_engine_arguments->children.resize(4);
/// Add table_name to engine arguments.
if (storage_engine_arguments->children.size() >= 2)
storage_engine_arguments->children.insert(storage_engine_arguments->children.begin() + 2, std::make_shared<ASTLiteral>(table_id.table_name));
}
return create_table_query;
}

View File

@ -226,6 +226,10 @@ def test_predefined_connection_configuration(started_cluster):
node1.query("DROP DATABASE IF EXISTS postgres_database")
node1.query("CREATE DATABASE postgres_database ENGINE = PostgreSQL(postgres1)")
result = node1.query("select create_table_query from system.tables where database ='postgres_database'")
assert(result.strip().endswith("ENGINE = PostgreSQL(postgres1, table = \\'test_table\\')"))
node1.query("INSERT INTO postgres_database.test_table SELECT number, number from numbers(100)")
assert (node1.query(f"SELECT count() FROM postgres_database.test_table").rstrip() == '100')