Skip to content

Commit c462979

Browse files
lex007inlex007in
authored
Support ALTER TABLE for adding PG types (#11630)
1 parent 0a30799 commit c462979

File tree

4 files changed

+132
-13
lines changed

4 files changed

+132
-13
lines changed

‎ydb/core/kqp/provider/yql_kikimr_exec.cpp

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1390,8 +1390,11 @@ class TKiSinkCallableExecutionTransformer : public TAsyncCallbackTransformer<TKi
13901390
auto typeNode = columnTuple.Item(1);
13911391
auto columnType = typeNode.Ref().GetTypeAnn();
13921392
auto type = columnType->Cast<TTypeExprType>()->GetType();
1393-
auto notNull = type->GetKind() != ETypeAnnotationKind::Optional;
1394-
auto actualType = notNull ? type : type->Cast<TOptionalExprType>()->GetItemType();
1393+
auto notNull = type->GetKind() != ETypeAnnotationKind::Optional
1394+
&& type->GetKind() != ETypeAnnotationKind::Pg; // pg types always optional
1395+
auto actualType = type->GetKind() == ETypeAnnotationKind::Optional
1396+
? type->Cast<TOptionalExprType>()->GetItemType()
1397+
: type;
13951398

13961399
TString error;
13971400
if (!SetColumnType(actualType, notNull, *add_column->mutable_type(), error)) {

‎ydb/core/kqp/provider/yql_kikimr_gateway.cpp

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -179,18 +179,20 @@ EYqlIssueCode YqlStatusFromYdbStatus(ui32 ydbStatus) {
179179
bool SetColumnType(const TTypeAnnotationNode* typeNode, bool notNull, Ydb::Type& protoType, TString& error) {
180180
switch (typeNode->GetKind()) {
181181
case ETypeAnnotationKind::Pg: {
182-
const auto pgTypeNode = typeNode->Cast<TPgExprType>();
183-
const TString typeName = pgTypeNode->GetName();
184-
const auto typeDesc = NKikimr::NPg::TypeDescFromPgTypeName(typeName);
182+
const auto* pgTypeNode = typeNode->Cast<TPgExprType>();
183+
const auto& typeId = pgTypeNode->GetId();
184+
const auto typeDesc = NKikimr::NPg::TypeDescFromPgTypeId(typeId);
185185
if (typeDesc) {
186186
Y_ABORT_UNLESS(!notNull, "It is not allowed to create NOT NULL pg columns");
187187
auto* pg = protoType.mutable_pg_type();
188188
pg->set_type_name(NKikimr::NPg::PgTypeNameFromTypeDesc(typeDesc));
189-
pg->set_type_modifier(NKikimr::NPg::TypeModFromPgTypeName(typeName));
190189
pg->set_oid(NKikimr::NPg::PgTypeIdFromTypeDesc(typeDesc));
191190
pg->set_typlen(0);
192191
pg->set_typmod(0);
193192
return true;
193+
} else {
194+
error = TStringBuilder() << "Unknown pg type: " << FormatType(pgTypeNode);
195+
return false;
194196
}
195197
}
196198
case ETypeAnnotationKind::Data: {

‎ydb/core/kqp/provider/yql_kikimr_type_ann.cpp

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1348,19 +1348,24 @@ virtual TStatus HandleCreateTable(TKiCreateTable create, TExprContext& ctx) over
13481348
auto actualType = (type->GetKind() == ETypeAnnotationKind::Optional) ?
13491349
type->Cast<TOptionalExprType>()->GetItemType() : type;
13501350

1351-
if (actualType->GetKind() != ETypeAnnotationKind::Data) {
1352-
columnTypeError(typeNode.Pos(), name, "Only core YQL data types are currently supported");
1351+
1352+
if (actualType->GetKind() != ETypeAnnotationKind::Data
1353+
&& actualType->GetKind() != ETypeAnnotationKind::Pg)
1354+
{
1355+
columnTypeError(typeNode.Pos(), name, "Only YQL data types and PG types are currently supported");
13531356
return TStatus::Error;
13541357
}
13551358

1356-
auto dataType = actualType->Cast<TDataExprType>();
1357-
1358-
if (!ValidateColumnDataType(dataType, typeNode, name, ctx)) {
1359-
return IGraphTransformer::TStatus::Error;
1359+
if (actualType->GetKind() == ETypeAnnotationKind::Data) {
1360+
if (!ValidateColumnDataType(actualType->Cast<TDataExprType>(), typeNode, name, ctx)) {
1361+
return IGraphTransformer::TStatus::Error;
1362+
}
1363+
} else {
1364+
// TODO: Validate pg modifiers
13601365
}
13611366

13621367
TKikimrColumnMetadata columnMeta;
1363-
// columnMeta.Name = columnName;
1368+
columnMeta.Name = name;
13641369
if (columnTuple.Size() > 2) {
13651370
const auto& columnConstraints = columnTuple.Item(2).Cast<TCoNameValueTuple>();
13661371
for(const auto& constraint: columnConstraints.Value().Cast<TCoNameValueTupleList>()) {

‎ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2433,6 +2433,49 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
24332433
createAndCheck(999, 99);
24342434
}
24352435

2436+
Y_UNIT_TEST(CreateTableWithPgColumn) {
2437+
TKikimrRunner kikimr;
2438+
auto db = kikimr.GetTableClient();
2439+
auto session = db.CreateSession().GetValueSync().GetSession();
2440+
2441+
auto createAndCheck = [&](const TString& typeName) {
2442+
TString tableName = TStringBuilder() << "/Root/TableWithPgColumn_" << typeName;
2443+
auto createQuery = TStringBuilder() << Sprintf(R"(
2444+
CREATE TABLE `%s` (
2445+
Key Uint64,
2446+
Value %s,
2447+
PRIMARY KEY (Key)
2448+
);)", tableName.c_str(), typeName.c_str());
2449+
auto createResult = session.ExecuteSchemeQuery(createQuery).GetValueSync();
2450+
2451+
UNIT_ASSERT_VALUES_EQUAL_C(createResult.GetStatus(), EStatus::SUCCESS, createResult.GetIssues().ToString());
2452+
2453+
TDescribeTableResult describe = session.DescribeTable(tableName).GetValueSync();
2454+
UNIT_ASSERT_EQUAL_C(describe.GetStatus(), EStatus::SUCCESS, describe.GetIssues().ToString());
2455+
auto tableDesc = describe.GetTableDescription();
2456+
TVector<TTableColumn> columns = tableDesc.GetTableColumns();
2457+
UNIT_ASSERT_VALUES_EQUAL(columns.size(), 2);
2458+
TType valueType = columns[1].Type;
2459+
TTypeParser parser(valueType);
2460+
auto kind = parser.GetKind();
2461+
UNIT_ASSERT_EQUAL(kind, TTypeParser::ETypeKind::Pg);
2462+
const auto& pgType = parser.GetPg();
2463+
UNIT_ASSERT_VALUES_EQUAL(pgType.TypeName, typeName);
2464+
};
2465+
2466+
createAndCheck("pgbool");
2467+
createAndCheck("pgint2");
2468+
createAndCheck("pgint4");
2469+
createAndCheck("pgint8");
2470+
createAndCheck("pgnumeric");
2471+
createAndCheck("pgfloat4");
2472+
createAndCheck("pgfloat8");
2473+
createAndCheck("pgtext");
2474+
createAndCheck("pgjson");
2475+
createAndCheck("pgtimestamp");
2476+
createAndCheck("pgvarchar");
2477+
}
2478+
24362479
void AlterTableAddIndex(EIndexTypeSql type) {
24372480
NKikimrConfig::TFeatureFlags featureFlags;
24382481
featureFlags.SetEnableVectorIndex(true);
@@ -3136,6 +3179,72 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
31363179
checkColumn(6, 35, 10);
31373180
}
31383181

3182+
Y_UNIT_TEST(AlterTableWithPgColumn) {
3183+
TKikimrRunner kikimr;
3184+
auto db = kikimr.GetTableClient();
3185+
auto session = db.CreateSession().GetValueSync().GetSession();
3186+
TString tableName = "/Root/TableWithPgColumn";
3187+
{
3188+
auto query = TStringBuilder() << R"(
3189+
CREATE TABLE `)" << tableName << R"(` (
3190+
Key Uint64,
3191+
Value String,
3192+
PRIMARY KEY (Key)
3193+
);)";
3194+
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
3195+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
3196+
}
3197+
3198+
auto addColumn = [&] (const TString& typeName) {
3199+
TString columnName = TStringBuilder() << "Column_" << typeName;
3200+
auto query = TStringBuilder() << Sprintf(R"(
3201+
ALTER TABLE `%s`
3202+
ADD COLUMN %s %s
3203+
)", tableName.c_str(), columnName.c_str(), typeName.c_str());
3204+
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
3205+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
3206+
};
3207+
3208+
addColumn("pgbool");
3209+
addColumn("pgint2");
3210+
addColumn("pgint4");
3211+
addColumn("pgint8");
3212+
addColumn("pgnumeric");
3213+
addColumn("pgfloat4");
3214+
addColumn("pgfloat8");
3215+
addColumn("pgtext");
3216+
addColumn("pgjson");
3217+
addColumn("pgtimestamp");
3218+
addColumn("pgvarchar");
3219+
3220+
TDescribeTableResult describe = session.DescribeTable(tableName).GetValueSync();
3221+
UNIT_ASSERT_EQUAL_C(describe.GetStatus(), EStatus::SUCCESS, describe.GetIssues().ToString());
3222+
auto tableDesc = describe.GetTableDescription();
3223+
TVector<TTableColumn> columns = tableDesc.GetTableColumns();
3224+
UNIT_ASSERT_VALUES_EQUAL(columns.size(), 13);
3225+
3226+
auto checkColumn = [&] (ui64 columnIdx, const TString& typeName) {
3227+
TType valueType = columns[columnIdx].Type;
3228+
TTypeParser parser(valueType);
3229+
auto kind = parser.GetKind();
3230+
UNIT_ASSERT_VALUES_EQUAL(kind, TTypeParser::ETypeKind::Pg);
3231+
const auto& pgType = parser.GetPg();
3232+
UNIT_ASSERT_VALUES_EQUAL(pgType.TypeName, typeName);
3233+
};
3234+
3235+
checkColumn(2, "pgbool");
3236+
checkColumn(3, "pgint2");
3237+
checkColumn(4, "pgint4");
3238+
checkColumn(5, "pgint8");
3239+
checkColumn(6, "pgnumeric");
3240+
checkColumn(7, "pgfloat4");
3241+
checkColumn(8, "pgfloat8");
3242+
checkColumn(9, "pgtext");
3243+
checkColumn(10, "pgjson");
3244+
checkColumn(11, "pgtimestamp");
3245+
checkColumn(12, "pgvarchar");
3246+
}
3247+
31393248
Y_UNIT_TEST(CreateUserWithPassword) {
31403249
TKikimrRunner kikimr;
31413250
auto db = kikimr.GetTableClient();

0 commit comments

Comments
 (0)