diff --git a/src/yb/util/strongly_typed_bool.h b/src/yb/util/strongly_typed_bool.h index 9713ec9b1b1..853a894780b 100644 --- a/src/yb/util/strongly_typed_bool.h +++ b/src/yb/util/strongly_typed_bool.h @@ -13,7 +13,7 @@ #pragma once -#include +#include #include @@ -31,8 +31,8 @@ template struct StronglyTypedBoolProxy { bool value; - operator bool() const { return value; } - StronglyTypedBoolProxy operator!() const { return StronglyTypedBoolProxy{!value}; } + constexpr operator bool() const { return value; } + constexpr auto operator!() const { return StronglyTypedBoolProxy{!value}; } }; template @@ -44,17 +44,17 @@ class StronglyTypedBool { StronglyTypedBool(false), StronglyTypedBool(true) }; - StronglyTypedBool(StronglyTypedBoolProxy proxy) : value_(proxy.value) {} // NOLINT + constexpr StronglyTypedBool(StronglyTypedBoolProxy proxy) : value_(proxy.value) {} // NOLINT // This is public so that we can construct a strongly-typed boolean value out of a regular one if // needed. In that case we'll have to spell out the class name, which will enforce readability. - explicit constexpr StronglyTypedBool(bool value) : value_(value) {} + constexpr explicit StronglyTypedBool(bool value) : value_(value) {} // These operators return regular bools so that it is easy to use strongly-typed bools in logical // expressions. - operator bool() const { return value_; } - StronglyTypedBool operator!() const { return StronglyTypedBool(!value_); } - bool get() const { return value_; } + constexpr operator bool() const { return value_; } + constexpr auto operator!() const { return StronglyTypedBool(!value_); } + constexpr bool get() const { return value_; } private: bool value_; diff --git a/src/yb/yql/pggate/pg_ddl.cc b/src/yb/yql/pggate/pg_ddl.cc index d27cca42a0e..fd1cbe6fadc 100644 --- a/src/yb/yql/pggate/pg_ddl.cc +++ b/src/yb/yql/pggate/pg_ddl.cc @@ -35,8 +35,7 @@ DEFINE_test_flag(int32, user_ddl_operation_timeout_sec, 0, DECLARE_int32(max_num_tablets_for_table); DECLARE_int32(yb_client_admin_operation_timeout_sec); -namespace yb { -namespace pggate { +namespace yb::pggate { using namespace std::literals; // NOLINT @@ -69,19 +68,21 @@ CoarseTimePoint CreateDatabaseDeadline() { // PgCreateDatabase //-------------------------------------------------------------------------------------------------- -PgCreateDatabase::PgCreateDatabase(PgSession::ScopedRefPtr pg_session, +PgCreateDatabase::PgCreateDatabase(const PgSession::ScopedRefPtr& pg_session, const char *database_name, - const PgOid database_oid, - const PgOid source_database_oid, - const PgOid next_oid, + PgOid database_oid, + PgOid source_database_oid, + PgOid next_oid, YbCloneInfo *yb_clone_info, - const bool colocated) - : PgDdl(std::move(pg_session)) { + bool colocated, + bool use_transaction) + : BaseType(pg_session) { req_.set_database_name(database_name); req_.set_database_oid(database_oid); req_.set_source_database_oid(source_database_oid); req_.set_next_oid(next_oid); req_.set_colocated(colocated); + req_.set_use_transaction(use_transaction); if (yb_clone_info) { req_.set_source_database_name(yb_clone_info->src_db_name); req_.set_clone_time(yb_clone_info->clone_time); @@ -90,39 +91,28 @@ PgCreateDatabase::PgCreateDatabase(PgSession::ScopedRefPtr pg_session, } } -PgCreateDatabase::~PgCreateDatabase() { -} - Status PgCreateDatabase::Exec() { return pg_session_->pg_client().CreateDatabase(&req_, CreateDatabaseDeadline()); } -PgDropDatabase::PgDropDatabase(PgSession::ScopedRefPtr pg_session, - const char *database_name, - PgOid database_oid) - : PgDdl(pg_session), +PgDropDatabase::PgDropDatabase( + const PgSession::ScopedRefPtr& pg_session, const char* database_name, PgOid database_oid) + : BaseType(pg_session), database_name_(database_name), database_oid_(database_oid) { } -PgDropDatabase::~PgDropDatabase() { -} - Status PgDropDatabase::Exec() { return pg_session_->DropDatabase(database_name_, database_oid_); } -PgAlterDatabase::PgAlterDatabase(PgSession::ScopedRefPtr pg_session, - const char *database_name, - PgOid database_oid) - : PgDdl(pg_session) { +PgAlterDatabase::PgAlterDatabase( + const PgSession::ScopedRefPtr& pg_session, const char* database_name, PgOid database_oid) + : BaseType(pg_session) { req_.set_database_name(database_name); req_.set_database_oid(database_oid); } -PgAlterDatabase::~PgAlterDatabase() { -} - Status PgAlterDatabase::Exec() { return pg_session_->pg_client().AlterDatabase(&req_, DdlDeadline()); } @@ -135,34 +125,25 @@ void PgAlterDatabase::RenameDatabase(const char *newname) { // PgCreateTablegroup / PgDropTablegroup //-------------------------------------------------------------------------------------------------- -PgCreateTablegroup::PgCreateTablegroup(PgSession::ScopedRefPtr pg_session, - const char *database_name, - const PgOid database_oid, - const PgOid tablegroup_oid, - const PgOid tablespace_oid) - : PgDdl(pg_session) { +PgCreateTablegroup::PgCreateTablegroup( + const PgSession::ScopedRefPtr& pg_session, const char* database_name, const PgOid database_oid, + const PgOid tablegroup_oid, const PgOid tablespace_oid) + : BaseType(pg_session) { req_.set_database_name(database_name); PgObjectId(database_oid, tablegroup_oid).ToPB(req_.mutable_tablegroup_id()); PgObjectId(database_oid, tablespace_oid).ToPB(req_.mutable_tablespace_id()); } -PgCreateTablegroup::~PgCreateTablegroup() { -} - Status PgCreateTablegroup::Exec() { return pg_session_->pg_client().CreateTablegroup(&req_, DdlDeadline()); } -PgDropTablegroup::PgDropTablegroup(PgSession::ScopedRefPtr pg_session, - const PgOid database_oid, - const PgOid tablegroup_oid) - : PgDdl(pg_session) { +PgDropTablegroup::PgDropTablegroup( + const PgSession::ScopedRefPtr& pg_session, PgOid database_oid, PgOid tablegroup_oid) + : BaseType(pg_session) { PgObjectId(database_oid, tablegroup_oid).ToPB(req_.mutable_tablegroup_id()); } -PgDropTablegroup::~PgDropTablegroup() { -} - Status PgDropTablegroup::Exec() { return pg_session_->pg_client().DropTablegroup(&req_, DdlDeadline()); } @@ -171,23 +152,25 @@ Status PgDropTablegroup::Exec() { // PgCreateTable //-------------------------------------------------------------------------------------------------- -PgCreateTable::PgCreateTable(PgSession::ScopedRefPtr pg_session, - const char *database_name, - const char *schema_name, - const char *table_name, - const PgObjectId& table_id, - bool is_shared_table, - bool is_sys_catalog_table, - bool if_not_exist, - PgYbrowidMode ybrowid_mode, - bool is_colocated_via_database, - const PgObjectId& tablegroup_oid, - const ColocationId colocation_id, - const PgObjectId& tablespace_oid, - bool is_matview, - const PgObjectId& pg_table_oid, - const PgObjectId& old_relfilenode_oid, - bool is_truncate) +PgCreateTableBase::PgCreateTableBase( + const PgSession::ScopedRefPtr& pg_session, + const char* database_name, + const char* schema_name, + const char* table_name, + const PgObjectId& table_id, + bool is_shared_table, + bool is_sys_catalog_table, + bool if_not_exist, + PgYbrowidMode ybrowid_mode, + bool is_colocated_via_database, + const PgObjectId& tablegroup_oid, + ColocationId colocation_id, + const PgObjectId& tablespace_oid, + bool is_matview, + const PgObjectId& pg_table_oid, + const PgObjectId& old_relfilenode_oid, + bool is_truncate, + bool use_transaction) : PgDdl(pg_session) { table_id.ToPB(req_.mutable_table_id()); req_.set_database_name(database_name); @@ -207,6 +190,7 @@ PgCreateTable::PgCreateTable(PgSession::ScopedRefPtr pg_session, pg_table_oid.ToPB(req_.mutable_pg_table_oid()); old_relfilenode_oid.ToPB(req_.mutable_old_relfilenode_oid()); req_.set_is_truncate(is_truncate); + req_.set_use_transaction(use_transaction); // Add internal primary key column to a Postgres table without a user-specified primary key. switch (ybrowid_mode) { @@ -221,13 +205,9 @@ PgCreateTable::PgCreateTable(PgSession::ScopedRefPtr pg_session, } } -Status PgCreateTable::AddColumnImpl(const char *attr_name, - int attr_num, - int attr_ybtype, - int pg_type_oid, - bool is_hash, - bool is_range, - SortingType sorting_type) { +Status PgCreateTableBase::AddColumnImpl( + const char* attr_name, int attr_num, int attr_ybtype, int pg_type_oid, bool is_hash, + bool is_range, SortingType sorting_type) { auto& column = *req_.mutable_create_columns()->Add(); column.set_attr_name(attr_name); column.set_attr_num(attr_num); @@ -239,7 +219,7 @@ Status PgCreateTable::AddColumnImpl(const char *attr_name, return Status::OK(); } -Status PgCreateTable::SetNumTablets(int32_t num_tablets) { +Status PgCreateTableBase::SetNumTablets(int32_t num_tablets) { if (num_tablets < 0) { return STATUS(InvalidArgument, "num_tablets cannot be less than zero"); } @@ -251,7 +231,7 @@ Status PgCreateTable::SetNumTablets(int32_t num_tablets) { return Status::OK(); } -Status PgCreateTable::SetVectorOptions(YbPgVectorIdxOptions *options) { +Status PgCreateTableBase::SetVectorOptions(YbPgVectorIdxOptions* options) { auto options_pb = req_.mutable_vector_idx_options(); options_pb->set_dist_type(static_cast(options->dist_type)); options_pb->set_idx_type(static_cast(options->idx_type)); @@ -266,7 +246,7 @@ Status PgCreateTable::SetVectorOptions(YbPgVectorIdxOptions *options) { return Status::OK(); } -Status PgCreateTable::AddSplitBoundary(PgExpr **exprs, int expr_count) { +Status PgCreateTableBase::AddSplitBoundary(PgExpr** exprs, int expr_count) { auto* values = req_.mutable_split_bounds()->Add()->mutable_values(); for (int i = 0; i < expr_count; ++i) { auto temp_value = VERIFY_RESULT(exprs[i]->Eval()); @@ -278,40 +258,80 @@ Status PgCreateTable::AddSplitBoundary(PgExpr **exprs, int expr_count) { return Status::OK(); } -Status PgCreateTable::Exec() { +Status PgCreateTableBase::Exec() { RETURN_NOT_OK(pg_session_->pg_client().CreateTable(&req_, DdlDeadline())); - auto base_table_id = PgObjectId::FromPB(req_.base_table_id()); + + const auto base_table_id = PgObjectId::FromPB(req_.base_table_id()); if (base_table_id.IsValid()) { pg_session_->InvalidateTableCache(base_table_id, InvalidateOnPgClient::kFalse); } return Status::OK(); } -void PgCreateTable::SetupIndex( - const PgObjectId& base_table_id, bool is_unique_index, bool skip_index_backfill) { +PgCreateTable::PgCreateTable( + const PgSession::ScopedRefPtr& pg_session, + const char* database_name, + const char* schema_name, + const char* table_name, + const PgObjectId& table_id, + bool is_shared_table, + bool is_sys_catalog_table, + bool if_not_exist, + PgYbrowidMode ybrowid_mode, + bool is_colocated_via_database, + const PgObjectId& tablegroup_oid, + ColocationId colocation_id, + const PgObjectId& tablespace_oid, + bool is_matview, + const PgObjectId& pg_table_oid, + const PgObjectId& old_relfilenode_oid, + bool is_truncate, + bool use_transaction) + : BaseType( + pg_session, database_name, schema_name, table_name, table_id, is_shared_table, + is_sys_catalog_table, if_not_exist, ybrowid_mode, is_colocated_via_database, + tablegroup_oid, colocation_id, tablespace_oid, is_matview, pg_table_oid, + old_relfilenode_oid, is_truncate, use_transaction) {} + +PgCreateIndex::PgCreateIndex( + const PgSession::ScopedRefPtr& pg_session, + const char* database_name, + const char* schema_name, + const char* table_name, + const PgObjectId& table_id, + bool is_shared_table, + bool is_sys_catalog_table, + bool if_not_exist, + PgYbrowidMode ybrowid_mode, + bool is_colocated_via_database, + const PgObjectId& tablegroup_oid, + ColocationId colocation_id, + const PgObjectId& tablespace_oid, + bool is_matview, + const PgObjectId& pg_table_oid, + const PgObjectId& old_relfilenode_oid, + bool is_truncate, + bool use_transaction, + const PgObjectId& base_table_id, + bool is_unique_index, + bool skip_index_backfill) + : BaseType( + pg_session, database_name, schema_name, table_name, table_id, is_shared_table, + is_sys_catalog_table, if_not_exist, ybrowid_mode, is_colocated_via_database, + tablegroup_oid, colocation_id, tablespace_oid, is_matview, pg_table_oid, + old_relfilenode_oid, is_truncate, use_transaction) { base_table_id.ToPB(req_.mutable_base_table_id()); req_.set_is_unique_index(is_unique_index); req_.set_skip_index_backfill(skip_index_backfill); } -StmtOp PgCreateTable::stmt_op() const { - return PgObjectId::FromPB(req_.base_table_id()).IsValid() - ? StmtOp::STMT_CREATE_INDEX : StmtOp::STMT_CREATE_TABLE; -} - //-------------------------------------------------------------------------------------------------- // PgDropTable //-------------------------------------------------------------------------------------------------- -PgDropTable::PgDropTable(PgSession::ScopedRefPtr pg_session, - const PgObjectId& table_id, - bool if_exist) - : PgDdl(pg_session), - table_id_(table_id), - if_exist_(if_exist) { -} - -PgDropTable::~PgDropTable() { +PgDropTable::PgDropTable( + const PgSession::ScopedRefPtr& pg_session, const PgObjectId& table_id, bool if_exist) + : BaseType(pg_session), table_id_(table_id), if_exist_(if_exist) { } Status PgDropTable::Exec() { @@ -327,15 +347,12 @@ Status PgDropTable::Exec() { // PgTruncateTable //-------------------------------------------------------------------------------------------------- -PgTruncateTable::PgTruncateTable(PgSession::ScopedRefPtr pg_session, - const PgObjectId& table_id) - : PgDdl(pg_session) { +PgTruncateTable::PgTruncateTable( + const PgSession::ScopedRefPtr& pg_session, const PgObjectId& table_id) + : BaseType(pg_session) { table_id.ToPB(req_.mutable_table_id()); } -PgTruncateTable::~PgTruncateTable() { -} - Status PgTruncateTable::Exec() { return pg_session_->pg_client().TruncateTable(&req_, DdlDeadline()); } @@ -344,24 +361,21 @@ Status PgTruncateTable::Exec() { // PgDropIndex //-------------------------------------------------------------------------------------------------- -PgDropIndex::PgDropIndex(PgSession::ScopedRefPtr pg_session, - const PgObjectId& index_id, - bool if_exist, - bool ddl_rollback_enabled) - : PgDropTable(pg_session, index_id, if_exist), ddl_rollback_enabled_(ddl_rollback_enabled) { -} - -PgDropIndex::~PgDropIndex() { +PgDropIndex::PgDropIndex( + const PgSession::ScopedRefPtr& pg_session, const PgObjectId& index_id, bool if_exist, + bool ddl_rollback_enabled) + : BaseType(pg_session), + index_id_(index_id), if_exist_(if_exist), ddl_rollback_enabled_(ddl_rollback_enabled) { } Status PgDropIndex::Exec() { client::YBTableName indexed_table_name; - Status s = pg_session_->DropIndex(table_id_, &indexed_table_name); + auto s = pg_session_->DropIndex(index_id_, &indexed_table_name); if (s.ok() || (s.IsNotFound() && if_exist_)) { RSTATUS_DCHECK(!indexed_table_name.empty(), Uninitialized, "indexed_table_name uninitialized"); PgObjectId indexed_table_id(indexed_table_name.table_id()); - pg_session_->InvalidateTableCache(table_id_, InvalidateOnPgClient::kFalse); + pg_session_->InvalidateTableCache(index_id_, InvalidateOnPgClient::kFalse); pg_session_->InvalidateTableCache(indexed_table_id, ddl_rollback_enabled_ ? InvalidateOnPgClient::kTrue : InvalidateOnPgClient::kFalse); return Status::OK(); @@ -373,10 +387,11 @@ Status PgDropIndex::Exec() { // PgAlterTable //-------------------------------------------------------------------------------------------------- -PgAlterTable::PgAlterTable(PgSession::ScopedRefPtr pg_session, - const PgObjectId& table_id) - : PgDdl(pg_session) { +PgAlterTable::PgAlterTable( + const PgSession::ScopedRefPtr& pg_session, const PgObjectId& table_id, bool use_transaction) + : BaseType(pg_session) { table_id.ToPB(req_.mutable_table_id()); + req_.set_use_transaction(use_transaction); } Status PgAlterTable::AddColumn(const char *name, @@ -458,35 +473,21 @@ void PgAlterTable::InvalidateTableCacheEntry() { PgObjectId::FromPB(req_.table_id()), InvalidateOnPgClient::kTrue); } -PgAlterTable::~PgAlterTable() { -} - //-------------------------------------------------------------------------------------------------- // PgDropSequence //-------------------------------------------------------------------------------------------------- -PgDropSequence::PgDropSequence(PgSession::ScopedRefPtr pg_session, - PgOid database_oid, - PgOid sequence_oid) - : PgDdl(std::move(pg_session)), - database_oid_(database_oid), - sequence_oid_(sequence_oid) { -} - -PgDropSequence::~PgDropSequence() { +PgDropSequence::PgDropSequence( + const PgSession::ScopedRefPtr& pg_session, PgOid database_oid, PgOid sequence_oid) + : BaseType(pg_session), database_oid_(database_oid), sequence_oid_(sequence_oid) { } Status PgDropSequence::Exec() { return pg_session_->pg_client().DeleteSequenceTuple(database_oid_, sequence_oid_); } -PgDropDBSequences::PgDropDBSequences(PgSession::ScopedRefPtr pg_session, - PgOid database_oid) - : PgDdl(std::move(pg_session)), - database_oid_(database_oid) { -} - -PgDropDBSequences::~PgDropDBSequences() { +PgDropDBSequences::PgDropDBSequences(const PgSession::ScopedRefPtr& pg_session, PgOid database_oid) + : BaseType(pg_session), database_oid_(database_oid) { } Status PgDropDBSequences::Exec() { @@ -496,12 +497,10 @@ Status PgDropDBSequences::Exec() { // PgCreateReplicationSlot //-------------------------------------------------------------------------------------------------- -PgCreateReplicationSlot::PgCreateReplicationSlot(PgSession::ScopedRefPtr pg_session, - const char *slot_name, - const char *plugin_name, - PgOid database_oid, - YBCPgReplicationSlotSnapshotAction snapshot_action) - : PgDdl(pg_session) { +PgCreateReplicationSlot::PgCreateReplicationSlot( + const PgSession::ScopedRefPtr& pg_session, const char* slot_name, const char* plugin_name, + PgOid database_oid, YBCPgReplicationSlotSnapshotAction snapshot_action) + : BaseType(pg_session) { req_.set_database_oid(database_oid); req_.set_replication_slot_name(slot_name); req_.set_output_plugin_name(plugin_name); @@ -524,15 +523,12 @@ Result PgCreateReplicationSlot::Exec return pg_session_->pg_client().CreateReplicationSlot(&req_, DdlDeadline()); } -PgCreateReplicationSlot::~PgCreateReplicationSlot() { -} - // PgDropReplicationSlot //-------------------------------------------------------------------------------------------------- -PgDropReplicationSlot::PgDropReplicationSlot(PgSession::ScopedRefPtr pg_session, - const char *slot_name) - : PgDdl(pg_session) { +PgDropReplicationSlot::PgDropReplicationSlot( + const PgSession::ScopedRefPtr& pg_session, const char* slot_name) + : BaseType(pg_session) { req_.set_replication_slot_name(slot_name); } @@ -540,8 +536,4 @@ Status PgDropReplicationSlot::Exec() { return pg_session_->pg_client().DropReplicationSlot(&req_, DdlDeadline()); } -PgDropReplicationSlot::~PgDropReplicationSlot() { -} - -} // namespace pggate -} // namespace yb +} // namespace yb::pggate diff --git a/src/yb/yql/pggate/pg_ddl.h b/src/yb/yql/pggate/pg_ddl.h index 432615dde7c..b1989ecff77 100644 --- a/src/yb/yql/pggate/pg_ddl.h +++ b/src/yb/yql/pggate/pg_ddl.h @@ -21,51 +21,35 @@ #include "yb/yql/pggate/pg_statement.h" -namespace yb { -namespace pggate { +namespace yb::pggate { class PgDdl : public PgStatement { - public: - explicit PgDdl(PgSession::ScopedRefPtr pg_session) : PgStatement(pg_session) { - } + protected: + explicit PgDdl(const PgSession::ScopedRefPtr& pg_session) : PgStatement(pg_session) {} }; -//-------------------------------------------------------------------------------------------------- -// CREATE DATABASE -//-------------------------------------------------------------------------------------------------- - -class PgCreateDatabase : public PgDdl { +class PgCreateDatabase final : public PgStatementLeafBase { public: - PgCreateDatabase(PgSession::ScopedRefPtr pg_session, - const char *database_name, + PgCreateDatabase(const PgSession::ScopedRefPtr& pg_session, + const char* database_name, PgOid database_oid, PgOid source_database_oid, PgOid next_oid, - YbCloneInfo *yb_clone_info, - const bool colocated); - virtual ~PgCreateDatabase(); - - void UseTransaction() { - req_.set_use_transaction(true); - } + YbCloneInfo* yb_clone_info, + bool colocated, + bool use_transaction); - StmtOp stmt_op() const override { return StmtOp::STMT_CREATE_DATABASE; } - - // Execute. Status Exec(); private: tserver::PgCreateDatabaseRequestPB req_; }; -class PgDropDatabase : public PgDdl { +class PgDropDatabase final : public PgStatementLeafBase { public: - PgDropDatabase(PgSession::ScopedRefPtr pg_session, const char *database_name, PgOid database_oid); - virtual ~PgDropDatabase(); - - StmtOp stmt_op() const override { return StmtOp::STMT_DROP_DATABASE; } + PgDropDatabase( + const PgSession::ScopedRefPtr& pg_session, const char* database_name, PgOid database_oid); - // Execute. Status Exec(); private: @@ -73,186 +57,183 @@ class PgDropDatabase : public PgDdl { const PgOid database_oid_; }; -class PgAlterDatabase : public PgDdl { +class PgAlterDatabase final : public PgStatementLeafBase { public: - PgAlterDatabase(PgSession::ScopedRefPtr pg_session, - const char *database_name, - PgOid database_oid); - virtual ~PgAlterDatabase(); - - StmtOp stmt_op() const override { return StmtOp::STMT_ALTER_DATABASE; } + PgAlterDatabase( + const PgSession::ScopedRefPtr& pg_session, const char* database_name, PgOid database_oid); void RenameDatabase(const char *newname); - // Execute. Status Exec(); private: tserver::PgAlterDatabaseRequestPB req_; }; -//-------------------------------------------------------------------------------------------------- -// CREATE / DROP TABLEGROUP -//-------------------------------------------------------------------------------------------------- - -class PgCreateTablegroup : public PgDdl { +class PgCreateTablegroup final : public PgStatementLeafBase { public: - // Assumes we are within a DDL transaction. - PgCreateTablegroup(PgSession::ScopedRefPtr pg_session, - const char *database_name, - const PgOid database_oid, - const PgOid tablegroup_oid, - const PgOid tablespace_oid); - virtual ~PgCreateTablegroup(); - - StmtOp stmt_op() const override { return StmtOp::STMT_CREATE_TABLEGROUP; } + PgCreateTablegroup( + const PgSession::ScopedRefPtr& pg_session, const char* database_name, PgOid database_oid, + PgOid tablegroup_oid, PgOid tablespace_oid); - // Execute. Status Exec(); private: tserver::PgCreateTablegroupRequestPB req_; }; -class PgDropTablegroup : public PgDdl { +class PgDropTablegroup final : public PgStatementLeafBase { public: - PgDropTablegroup(PgSession::ScopedRefPtr pg_session, - PgOid database_oid, - PgOid tablegroup_oid); - virtual ~PgDropTablegroup(); - - StmtOp stmt_op() const override { return StmtOp::STMT_DROP_TABLEGROUP; } + PgDropTablegroup( + const PgSession::ScopedRefPtr& pg_session, PgOid database_oid, PgOid tablegroup_oid); - // Execute. Status Exec(); private: tserver::PgDropTablegroupRequestPB req_; }; -//-------------------------------------------------------------------------------------------------- -// CREATE TABLE -//-------------------------------------------------------------------------------------------------- - -class PgCreateTable : public PgDdl { +class PgCreateTableBase : public PgDdl { public: - PgCreateTable(PgSession::ScopedRefPtr pg_session, - const char *database_name, - const char *schema_name, - const char *table_name, - const PgObjectId& table_id, - bool is_shared_table, - bool is_sys_catalog_table, - bool if_not_exist, - PgYbrowidMode ybrowid_mode, - bool is_colocated_via_database, - const PgObjectId& tablegroup_oid, - const ColocationId colocation_id, - const PgObjectId& tablespace_oid, - bool is_matview, - const PgObjectId& pg_table_oid, - const PgObjectId& old_relfilenode_oid, - bool is_truncate); - - void SetupIndex( - const PgObjectId& base_table_id, bool is_unique_index, bool skip_index_backfill); - - StmtOp stmt_op() const override; - - Status AddColumn(const char *attr_name, - int attr_num, - int attr_ybtype, - bool is_hash, - bool is_range, - SortingType sorting_type = SortingType::kNotSpecified) { - return AddColumnImpl(attr_name, attr_num, attr_ybtype, 20 /*INT8OID*/, - is_hash, is_range, sorting_type); + Status AddColumn( + const char* attr_name, int attr_num, int attr_ybtype, bool is_hash, bool is_range, + SortingType sorting_type = SortingType::kNotSpecified) { + return AddColumnImpl( + attr_name, attr_num, attr_ybtype, 20 /*INT8OID*/, is_hash, is_range, sorting_type); } - Status AddColumn(const char *attr_name, - int attr_num, - const YBCPgTypeEntity *attr_type, - bool is_hash, - bool is_range, - SortingType sorting_type = SortingType::kNotSpecified) { - return AddColumnImpl(attr_name, attr_num, attr_type->yb_type, attr_type->type_oid, - is_hash, is_range, sorting_type); + Status AddColumn( + const char* attr_name, int attr_num, const YBCPgTypeEntity* attr_type, bool is_hash, + bool is_range, SortingType sorting_type = SortingType::kNotSpecified) { + return AddColumnImpl( + attr_name, attr_num, attr_type->yb_type, attr_type->type_oid, is_hash, is_range, + sorting_type); } - // Specify the number of tablets explicitly. Status SetNumTablets(int32_t num_tablets); - Status SetVectorOptions(YbPgVectorIdxOptions *options); + Status SetVectorOptions(YbPgVectorIdxOptions* options); - Status AddSplitBoundary(PgExpr **exprs, int expr_count); - - void UseTransaction() { - req_.set_use_transaction(true); - } + Status AddSplitBoundary(PgExpr** exprs, int expr_count); - // Execute. - virtual Status Exec(); + Status Exec(); protected: - virtual Status AddColumnImpl( - const char *attr_name, int attr_num, int attr_ybtype, int pg_type_oid, bool is_hash, - bool is_range, SortingType sorting_type = SortingType::kNotSpecified); + PgCreateTableBase(const PgSession::ScopedRefPtr& pg_session, + const char* database_name, + const char* schema_name, + const char* table_name, + const PgObjectId& table_id, + bool is_shared_table, + bool is_sys_catalog_table, + bool if_not_exist, + PgYbrowidMode ybrowid_mode, + bool is_colocated_via_database, + const PgObjectId& tablegroup_oid, + ColocationId colocation_id, + const PgObjectId& tablespace_oid, + bool is_matview, + const PgObjectId& pg_table_oid, + const PgObjectId& old_relfilenode_oid, + bool is_truncate, + bool use_transaction); - private: tserver::PgCreateTableRequestPB req_; + + private: + Status AddColumnImpl( + const char* attr_name, int attr_num, int attr_ybtype, int pg_type_oid, bool is_hash, + bool is_range, SortingType sorting_type = SortingType::kNotSpecified); }; -class PgDropTable : public PgDdl { +class PgCreateTable final : public PgStatementLeafBase { public: - PgDropTable(PgSession::ScopedRefPtr pg_session, const PgObjectId& table_id, bool if_exist); - virtual ~PgDropTable(); + PgCreateTable( + const PgSession::ScopedRefPtr& pg_session, + const char* database_name, + const char* schema_name, + const char* table_name, + const PgObjectId& table_id, + bool is_shared_table, + bool is_sys_catalog_table, + bool if_not_exist, + PgYbrowidMode ybrowid_mode, + bool is_colocated_via_database, + const PgObjectId& tablegroup_oid, + ColocationId colocation_id, + const PgObjectId& tablespace_oid, + bool is_matview, + const PgObjectId& pg_table_oid, + const PgObjectId& old_relfilenode_oid, + bool is_truncate, + bool use_transaction); +}; - StmtOp stmt_op() const override { return StmtOp::STMT_DROP_TABLE; } +class PgCreateIndex final : public PgStatementLeafBase { + public: + PgCreateIndex( + const PgSession::ScopedRefPtr& pg_session, + const char* database_name, + const char* schema_name, + const char* table_name, + const PgObjectId& table_id, + bool is_shared_table, + bool is_sys_catalog_table, + bool if_not_exist, + PgYbrowidMode ybrowid_mode, + bool is_colocated_via_database, + const PgObjectId& tablegroup_oid, + ColocationId colocation_id, + const PgObjectId& tablespace_oid, + bool is_matview, + const PgObjectId& pg_table_oid, + const PgObjectId& old_relfilenode_oid, + bool is_truncate, + bool use_transaction, + const PgObjectId& base_table_id, + bool is_unique_index, + bool skip_index_backfill); +}; + +class PgDropTable final : public PgStatementLeafBase { + public: + PgDropTable(const PgSession::ScopedRefPtr& pg_session, const PgObjectId& table_id, bool if_exist); - // Execute. Status Exec(); protected: const PgObjectId table_id_; - bool if_exist_; + const bool if_exist_; }; -class PgTruncateTable : public PgDdl { +class PgTruncateTable final : public PgStatementLeafBase { public: - PgTruncateTable(PgSession::ScopedRefPtr pg_session, const PgObjectId& table_id); - virtual ~PgTruncateTable(); + PgTruncateTable(const PgSession::ScopedRefPtr& pg_session, const PgObjectId& table_id); - StmtOp stmt_op() const override { return StmtOp::STMT_TRUNCATE_TABLE; } - - // Execute. Status Exec(); private: tserver::PgTruncateTableRequestPB req_; }; -class PgDropIndex : public PgDropTable { +class PgDropIndex final : public PgStatementLeafBase { public: - PgDropIndex(PgSession::ScopedRefPtr pg_session, const PgObjectId& index_id, bool if_exist, - bool ddl_rollback_enabled); - virtual ~PgDropIndex(); - - StmtOp stmt_op() const override { return StmtOp::STMT_DROP_INDEX; } + PgDropIndex( + const PgSession::ScopedRefPtr& pg_session, const PgObjectId& index_id, bool if_exist, + bool ddl_rollback_enabled); - // Execute. Status Exec(); + private: - bool ddl_rollback_enabled_; + const PgObjectId index_id_; + const bool if_exist_; + const bool ddl_rollback_enabled_; }; -//-------------------------------------------------------------------------------------------------- -// ALTER TABLE -//-------------------------------------------------------------------------------------------------- - -class PgAlterTable : public PgDdl { +class PgAlterTable final : public PgStatementLeafBase { public: - PgAlterTable(PgSession::ScopedRefPtr pg_session, - const PgObjectId& table_id); + PgAlterTable( + const PgSession::ScopedRefPtr& pg_session, const PgObjectId& table_id, bool use_transaction); Status AddColumn(const char *name, const YBCPgTypeEntity *attr_type, @@ -277,92 +258,53 @@ class PgAlterTable : public PgDdl { void InvalidateTableCacheEntry(); - virtual ~PgAlterTable(); - - StmtOp stmt_op() const override { return StmtOp::STMT_ALTER_TABLE; } - - void UseTransaction() { - req_.set_use_transaction(true); - } - private: tserver::PgAlterTableRequestPB req_; }; -//-------------------------------------------------------------------------------------------------- -// DROP SEQUENCE -//-------------------------------------------------------------------------------------------------- - -class PgDropSequence : public PgDdl { +class PgDropSequence final : public PgStatementLeafBase { public: - PgDropSequence(PgSession::ScopedRefPtr pg_session, - PgOid database_oid, - PgOid sequence_oid); - virtual ~PgDropSequence(); + PgDropSequence(const PgSession::ScopedRefPtr& pg_session, PgOid database_oid, PgOid sequence_oid); - StmtOp stmt_op() const override { return StmtOp::STMT_DROP_SEQUENCE; } - - // Execute. Status Exec(); private: - PgOid database_oid_; - PgOid sequence_oid_; + const PgOid database_oid_; + const PgOid sequence_oid_; }; -class PgDropDBSequences : public PgDdl { +class PgDropDBSequences final : public PgStatementLeafBase { public: - PgDropDBSequences(PgSession::ScopedRefPtr pg_session, - PgOid database_oid); - virtual ~PgDropDBSequences(); + PgDropDBSequences(const PgSession::ScopedRefPtr& pg_session, PgOid database_oid); - StmtOp stmt_op() const override { return StmtOp::STMT_DROP_DB_SEQUENCES; } - - // Execute. Status Exec(); private: - PgOid database_oid_; + const PgOid database_oid_; }; -// CREATE REPLICATION SLOT -//-------------------------------------------------------------------------------------------------- - -class PgCreateReplicationSlot : public PgDdl { +class PgCreateReplicationSlot final : public PgStatementLeafBase< + PgDdl, StmtOp::kCreateReplicationSlot> { public: - PgCreateReplicationSlot(PgSession::ScopedRefPtr pg_session, - const char *slot_name, - const char *plugin_name, - PgOid database_oid, - YBCPgReplicationSlotSnapshotAction snapshot_action); + PgCreateReplicationSlot( + const PgSession::ScopedRefPtr& pg_session, const char* slot_name, const char* plugin_name, + PgOid database_oid, YBCPgReplicationSlotSnapshotAction snapshot_action); Result Exec(); - virtual ~PgCreateReplicationSlot(); - - StmtOp stmt_op() const override { return StmtOp::STMT_CREATE_REPLICATION_SLOT; } - private: tserver::PgCreateReplicationSlotRequestPB req_; }; -// DROP REPLICATION SLOT -//-------------------------------------------------------------------------------------------------- - -class PgDropReplicationSlot : public PgDdl { +class PgDropReplicationSlot final : public PgStatementLeafBase< + PgDdl, StmtOp::kDropReplicationSlot> { public: - PgDropReplicationSlot(PgSession::ScopedRefPtr pg_session, - const char *slot_name); + PgDropReplicationSlot(const PgSession::ScopedRefPtr& pg_session, const char *slot_name); Status Exec(); - virtual ~PgDropReplicationSlot(); - - StmtOp stmt_op() const override { return StmtOp::STMT_DROP_REPLICATION_SLOT; } - private: tserver::PgDropReplicationSlotRequestPB req_; }; -} // namespace pggate -} // namespace yb +} // namespace yb::pggate diff --git a/src/yb/yql/pggate/pg_delete.h b/src/yb/yql/pggate/pg_delete.h index 5bc8fb9ca8b..a33cb695a37 100644 --- a/src/yb/yql/pggate/pg_delete.h +++ b/src/yb/yql/pggate/pg_delete.h @@ -22,10 +22,8 @@ namespace yb::pggate { -class PgDelete : public PgDmlWrite { +class PgDelete final : public PgStatementLeafBase { public: - StmtOp stmt_op() const override { return StmtOp::STMT_DELETE; } - void SetIsPersistNeeded(bool is_persist_needed) { write_req_->set_is_delete_persist_needed(is_persist_needed); } @@ -41,7 +39,7 @@ class PgDelete : public PgDmlWrite { private: PgDelete( const PgSession::ScopedRefPtr& pg_session, YBCPgTransactionSetting transaction_setting) - : PgDmlWrite(pg_session, transaction_setting) {} + : BaseType(pg_session, transaction_setting) {} PgsqlWriteRequestPB::PgsqlStmtType stmt_type() const override { return PgsqlWriteRequestPB::PGSQL_DELETE; diff --git a/src/yb/yql/pggate/pg_dml_read.h b/src/yb/yql/pggate/pg_dml_read.h index 16a5d7b52b6..d25c136b820 100644 --- a/src/yb/yql/pggate/pg_dml_read.h +++ b/src/yb/yql/pggate/pg_dml_read.h @@ -57,8 +57,6 @@ namespace yb::pggate { class PgDmlRead : public PgDml { public: - [[nodiscard]] StmtOp stmt_op() const override { return StmtOp::STMT_SELECT; } - // Allocate binds. virtual void PrepareBinds(); diff --git a/src/yb/yql/pggate/pg_insert.h b/src/yb/yql/pggate/pg_insert.h index 8e498ca9acb..e5142118591 100644 --- a/src/yb/yql/pggate/pg_insert.h +++ b/src/yb/yql/pggate/pg_insert.h @@ -22,10 +22,8 @@ namespace yb::pggate { -class PgInsert : public PgDmlWrite { +class PgInsert final : public PgStatementLeafBase { public: - StmtOp stmt_op() const override { return StmtOp::STMT_INSERT; } - void SetUpsertMode() { write_req_->set_stmt_type(PgsqlWriteRequestPB::PGSQL_UPSERT); } @@ -46,7 +44,7 @@ class PgInsert : public PgDmlWrite { PgInsert( const PgSession::ScopedRefPtr& pg_session, YBCPgTransactionSetting transaction_setting, bool packed) - : PgDmlWrite(pg_session, transaction_setting, packed) {} + : BaseType(pg_session, transaction_setting, packed) {} PgsqlWriteRequestPB::PgsqlStmtType stmt_type() const override { return PgsqlWriteRequestPB::PGSQL_INSERT; diff --git a/src/yb/yql/pggate/pg_sample.cc b/src/yb/yql/pggate/pg_sample.cc index 52ef940b3a2..6456080947d 100644 --- a/src/yb/yql/pggate/pg_sample.cc +++ b/src/yb/yql/pggate/pg_sample.cc @@ -131,7 +131,7 @@ class PgSamplePicker : public PgSelectIndex { }; PgSample::PgSample(const PgSession::ScopedRefPtr& pg_session) - : PgDmlRead(pg_session) {} + : BaseType(pg_session) {} Status PgSample::Prepare( const PgObjectId& table_id, bool is_region_local, int targrows, diff --git a/src/yb/yql/pggate/pg_sample.h b/src/yb/yql/pggate/pg_sample.h index 8be909b5015..259911b0cb5 100644 --- a/src/yb/yql/pggate/pg_sample.h +++ b/src/yb/yql/pggate/pg_sample.h @@ -31,10 +31,8 @@ class PgSamplePicker; //-------------------------------------------------------------------------------------------------- // SAMPLE collect table statistics and take random rows sample //-------------------------------------------------------------------------------------------------- -class PgSample : public PgDmlRead { +class PgSample final : public PgStatementLeafBase { public: - StmtOp stmt_op() const override { return StmtOp::STMT_SAMPLE; } - // Make PgSamplePicker to process next block of rows in the table. // The has_more parameter is set to true if table has and needs more blocks. // PgSampler is not ready to be executed until this function returns false diff --git a/src/yb/yql/pggate/pg_select.cc b/src/yb/yql/pggate/pg_select.cc index b3fd9af2611..4a9ae53eae3 100644 --- a/src/yb/yql/pggate/pg_select.cc +++ b/src/yb/yql/pggate/pg_select.cc @@ -24,7 +24,7 @@ namespace yb::pggate { PgSelect::PgSelect(const PgSession::ScopedRefPtr& pg_session) - : PgDmlRead(pg_session) {} + : BaseType(pg_session) {} Status PgSelect::Prepare( const PgObjectId& table_id, bool is_region_local, diff --git a/src/yb/yql/pggate/pg_select.h b/src/yb/yql/pggate/pg_select.h index d4ce3f88bad..39edf9b66c6 100644 --- a/src/yb/yql/pggate/pg_select.h +++ b/src/yb/yql/pggate/pg_select.h @@ -23,7 +23,7 @@ namespace yb::pggate { -class PgSelect : public PgDmlRead { +class PgSelect : public PgStatementLeafBase { public: struct IndexQueryInfo { IndexQueryInfo(PgObjectId id_, bool is_embedded_) diff --git a/src/yb/yql/pggate/pg_statement.h b/src/yb/yql/pggate/pg_statement.h index 378f18519e6..8c47aac9750 100644 --- a/src/yb/yql/pggate/pg_statement.h +++ b/src/yb/yql/pggate/pg_statement.h @@ -16,6 +16,7 @@ #include +#include #include #include #include @@ -26,6 +27,7 @@ #include "yb/gutil/ref_counted.h" #include "yb/util/bytes_formatter.h" +#include "yb/util/enums.h" #include "yb/yql/pggate/pg_expr.h" #include "yb/yql/pggate/pg_memctx.h" @@ -34,36 +36,29 @@ namespace yb::pggate { -// Statement types. -// - Might use it for error reporting or debugging or if different operations share the same -// CAPI calls. -// - TODO(neil) Remove StmtOp if we don't need it. -enum class StmtOp { - STMT_NOOP = 0, - STMT_CREATE_DATABASE, - STMT_DROP_DATABASE, - STMT_CREATE_SCHEMA, - STMT_DROP_SCHEMA, - STMT_CREATE_TABLE, - STMT_DROP_TABLE, - STMT_TRUNCATE_TABLE, - STMT_CREATE_INDEX, - STMT_DROP_INDEX, - STMT_ALTER_TABLE, - STMT_INSERT, - STMT_UPDATE, - STMT_DELETE, - STMT_TRUNCATE, - STMT_SELECT, - STMT_ALTER_DATABASE, - STMT_CREATE_TABLEGROUP, - STMT_DROP_TABLEGROUP, - STMT_SAMPLE, - STMT_DROP_SEQUENCE, - STMT_DROP_DB_SEQUENCES, - STMT_CREATE_REPLICATION_SLOT, - STMT_DROP_REPLICATION_SLOT, -}; +YB_DEFINE_TYPED_ENUM(StmtOp, uint8_t, + (kCreateDatabase) + (kDropDatabase) + (kCreateTable) + (kDropTable) + (kTruncateTable) + (kCreateIndex) + (kDropIndex) + (kAlterTable) + (kInsert) + (kUpdate) + (kDelete) + (kTruncate) + (kSelect) + (kAlterDatabase) + (kCreateTablegroup) + (kDropTablegroup) + (kSample) + (kDropSequence) + (kDropDbSequences) + (kCreateReplicationSlot) + (kDropReplicationSlot) +); class PgStatement : public PgMemctx::Registrable { public: @@ -73,16 +68,12 @@ class PgStatement : public PgMemctx::Registrable { virtual ~PgStatement() = default; - virtual StmtOp stmt_op() const = 0; + [[nodiscard]] virtual StmtOp stmt_op() const = 0; const std::shared_ptr& arena_ptr() const { return arena_; } ThreadSafeArena& arena() const { return *arena_; } - static bool IsValidStmt(const PgStatement* stmt, StmtOp op) { - return stmt && stmt->stmt_op() == op; - } - protected: // YBSession that this statement belongs to. PgSession::ScopedRefPtr pg_session_; @@ -91,4 +82,22 @@ class PgStatement : public PgMemctx::Registrable { std::shared_ptr arena_; }; +template +concept PgStatementSubclass = std::derived_from && !std::same_as; + +template +class PgStatementLeafBase : public Base { + public: + using BaseParam = Base; + + static constexpr StmtOp kStmtOp = T; + [[nodiscard]] StmtOp stmt_op() const override final { return kStmtOp; } + + protected: + using BaseType = PgStatementLeafBase; + + template + explicit PgStatementLeafBase(Args&&... args) : Base(std::forward(args)...) {} +}; + } // namespace yb::pggate diff --git a/src/yb/yql/pggate/pg_truncate_colocated.h b/src/yb/yql/pggate/pg_truncate_colocated.h index 72ec79a823c..f0399d106bf 100644 --- a/src/yb/yql/pggate/pg_truncate_colocated.h +++ b/src/yb/yql/pggate/pg_truncate_colocated.h @@ -14,7 +14,7 @@ #pragma once -#include +#include #include "yb/util/result.h" @@ -22,10 +22,8 @@ namespace yb::pggate { -class PgTruncateColocated : public PgDmlWrite { +class PgTruncateColocated final : public PgStatementLeafBase { public: - StmtOp stmt_op() const override { return StmtOp::STMT_TRUNCATE; } - static Result> Make( const PgSession::ScopedRefPtr& pg_session, const PgObjectId& table_id, bool is_region_local, YBCPgTransactionSetting transaction_setting) { @@ -38,7 +36,7 @@ class PgTruncateColocated : public PgDmlWrite { private: PgTruncateColocated( const PgSession::ScopedRefPtr& pg_session, YBCPgTransactionSetting transaction_setting) - : PgDmlWrite(pg_session, transaction_setting) {} + : BaseType(pg_session, transaction_setting) {} PgsqlWriteRequestPB::PgsqlStmtType stmt_type() const override { return PgsqlWriteRequestPB::PGSQL_TRUNCATE_COLOCATED; diff --git a/src/yb/yql/pggate/pg_update.h b/src/yb/yql/pggate/pg_update.h index 8b6c9695b97..320d472bd8d 100644 --- a/src/yb/yql/pggate/pg_update.h +++ b/src/yb/yql/pggate/pg_update.h @@ -22,10 +22,8 @@ namespace yb::pggate { -class PgUpdate : public PgDmlWrite { +class PgUpdate final : public PgStatementLeafBase { public: - StmtOp stmt_op() const override { return StmtOp::STMT_UPDATE; } - static Result> Make( const PgSession::ScopedRefPtr& pg_session, const PgObjectId& table_id, bool is_region_local, YBCPgTransactionSetting transaction_setting) { @@ -36,7 +34,7 @@ class PgUpdate : public PgDmlWrite { private: PgUpdate(const PgSession::ScopedRefPtr& pg_session, YBCPgTransactionSetting transaction_setting) - : PgDmlWrite(pg_session, transaction_setting) {} + : BaseType(pg_session, transaction_setting) {} PgsqlWriteRequestPB::PgsqlStmtType stmt_type() const override { return PgsqlWriteRequestPB::PGSQL_UPDATE; diff --git a/src/yb/yql/pggate/pggate.cc b/src/yb/yql/pggate/pggate.cc index c3ad95c13e9..2f7a7a01d72 100644 --- a/src/yb/yql/pggate/pggate.cc +++ b/src/yb/yql/pggate/pggate.cc @@ -15,11 +15,12 @@ #include "yb/yql/pggate/pggate.h" #include -#include -#include -#include +#include +#include +#include +#include +#include #include -#include #include @@ -91,11 +92,11 @@ DEFINE_RUNTIME_PREVIEW_bool(ysql_pack_inserted_value, false, namespace yb::pggate { namespace { -Status AddColumn(PgCreateTable* pg_stmt, const char *attr_name, int attr_num, - const YBCPgTypeEntity *attr_type, bool is_hash, bool is_range, - bool is_desc, bool is_nulls_first) { - using SortingType = SortingType; - SortingType sorting_type = SortingType::kNotSpecified; +Status AddColumn( + PgCreateTableBase& create_table, const char* attr_name, int attr_num, + const YBCPgTypeEntity* attr_type, bool is_hash, bool is_range, bool is_desc, + bool is_nulls_first) { + auto sorting_type = SortingType::kNotSpecified; if (!is_hash && is_range) { if (is_desc) { @@ -105,7 +106,7 @@ Status AddColumn(PgCreateTable* pg_stmt, const char *attr_name, int attr_num, } } - return pg_stmt->AddColumn(attr_name, attr_num, attr_type, is_hash, is_range, sorting_type); + return create_table.AddColumn(attr_name, attr_num, attr_type, is_hash, is_range, sorting_type); } Result BuildMessenger( @@ -181,6 +182,204 @@ Result> MakeSelectStatement( pg_session, table_id, is_region_local, MakeIndexQueryInfo(index_id, prepare_params))); } +namespace get_statement_as::internal { + +template +concept SingleStmtOpSubclass = + std::derived_from>; + +template +concept GroupStmtOpSubclass = PgStatementSubclass && !SingleStmtOpSubclass; + +template +struct StatementTag { + constexpr StatementTag() : stmt_op(T::kStmtOp) {} + + template + requires(GroupStmtOpSubclass && std::derived_from) + constexpr StatementTag(const StatementTag& t) : stmt_op(t.stmt_op) {} + + const StmtOp stmt_op; +}; + +template +using StatementTags = std::initializer_list>; + +template +struct TagsResolver; + +template <> +struct TagsResolver { + static constexpr StatementTags kTags = { + StatementTag(), + StatementTag() + }; +}; + +template <> +struct TagsResolver { + static constexpr StatementTags kTags = { + // PgCreateTableBase + StatementTag(), + StatementTag(), + + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag() + }; +}; + +template <> +struct TagsResolver { + static constexpr StatementTags kTags = { + StatementTag(), + StatementTag() + }; +}; + +template <> +struct TagsResolver { + static constexpr StatementTags kTags = { + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag() + }; +}; + +template <> +struct TagsResolver { + static constexpr StatementTags kTags = { + // PgDmlRead + StatementTag(), + StatementTag(), + + // PgDmlWrite + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag() + }; +}; + +template +constexpr bool AreTagsCoveringSubtags( + Iterator i, Iterator end, StatementTags subtags, Args&&... args) { + if constexpr (!sizeof...(Args)) { + for (auto st : subtags) { + if (i == end || (*i++).stmt_op != st.stmt_op) { + return false; + } + } + return true; + } else { + return AreTagsCoveringSubtags(i, end, subtags) && + AreTagsCoveringSubtags(i + subtags.size(), end, std::forward(args)...); + } +} + +template +constexpr bool AreTagsCoveringSubtags(StatementTags tags, Args&&... args) { + return AreTagsCoveringSubtags(tags.begin(), tags.end(), std::forward(args)...); +} + +using AllStmtOpsBoolArray = std::array(nullptr))>; + +template +constexpr std::pair UniqTagsCount( + AllStmtOpsBoolArray& found_items, StatementTags tags, Args&&... args) { + size_t count = 0; + bool has_duplicates = false; + if constexpr (sizeof...(Args)) { + std::tie(count, has_duplicates) = UniqTagsCount(found_items, std::forward(args)...); + } + for (auto t : tags) { + const auto idx = to_underlying(t.stmt_op); + if (idx < found_items.size() && !found_items.data()[idx]) { + ++count; + found_items.data()[idx] = true; + } else { + has_duplicates = true; + } + } + return {count, has_duplicates}; +} + +template +constexpr bool DoesTagsSpanAllStmtOps(Args&&... args) { + AllStmtOpsBoolArray found_items = {}; + size_t uniq_count = 0; + bool has_duplicates = false; + std::tie(uniq_count, has_duplicates) = UniqTagsCount(found_items, std::forward(args)...); + return !has_duplicates && uniq_count == NumEnumElements(static_cast(nullptr)); +} + +static_assert( + AreTagsCoveringSubtags( + TagsResolver::kTags, + TagsResolver::kTags, TagsResolver::kTags) && + AreTagsCoveringSubtags(TagsResolver::kTags, TagsResolver::kTags) && + DoesTagsSpanAllStmtOps(TagsResolver::kTags, TagsResolver::kTags)); + +template +Result GetStatementAs(PgStatement* handle) { + RSTATUS_DCHECK( + handle && handle->stmt_op() == T::kStmtOp, InvalidArgument, "Invalid statement handle"); + return down_cast(*handle); +} + +template +bool IsStmtOpAllowed(StmtOp stmt_op, StatementTags tags) { + for (auto t : tags) { + if (t.stmt_op == stmt_op) { + return true; + } + } + return false; +} + +template +Result GetStatementAs(PgStatement* handle, StatementTags tags) { + RSTATUS_DCHECK( + handle && IsStmtOpAllowed(handle->stmt_op(), tags), + InvalidArgument, "Invalid statement handle"); + + return down_cast(*handle); +} + +template +auto GetStatementAs(PgStatement* handle) { + return GetStatementAs(handle, TagsResolver::kTags); +} + +} // namespace get_statement_as::internal + +using get_statement_as::internal::GetStatementAs; +using get_statement_as::internal::StatementTag; + +Result GetArena(PgStatement* handle) { + RSTATUS_DCHECK(handle, InvalidArgument, "Invalid statement handle"); + return handle->arena(); +} + +template +requires(std::derived_from) +Status ExecDdlWithSyscatalogChanges(PgStatement* handle, PgSession& session) { + auto& stmt = VERIFY_RESULT_REF(GetStatementAs(handle)); + session.SetDdlHasSyscatalogChanges(); + return stmt.Exec(); +} + } // namespace //-------------------------------------------------------------------------------------------------- @@ -608,66 +807,42 @@ Status PgApiImpl::IsDatabaseColocated(const PgOid database_oid, bool *colocated, } Status PgApiImpl::NewCreateDatabase( - const char* database_name, const PgOid database_oid, const PgOid source_database_oid, - const PgOid next_oid, const bool colocated, YbCloneInfo *yb_clone_info, PgStatement** handle) { - auto stmt = std::make_unique( - pg_session_, database_name, database_oid, source_database_oid, next_oid, - yb_clone_info, colocated); - if (pg_txn_manager_->IsDdlMode()) { - stmt->UseTransaction(); - } - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); + const char* database_name, PgOid database_oid, PgOid source_database_oid, PgOid next_oid, + bool colocated, YbCloneInfo* yb_clone_info, PgStatement **handle) { + return AddToCurrentPgMemctx( + std::make_unique( + pg_session_, database_name, database_oid, source_database_oid, next_oid, yb_clone_info, + colocated, pg_txn_manager_->IsDdlMode()), + handle); } -Status PgApiImpl::ExecCreateDatabase(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_DATABASE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecCreateDatabase(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } -Status PgApiImpl::NewDropDatabase(const char *database_name, - PgOid database_oid, - PgStatement **handle) { - auto stmt = std::make_unique(pg_session_, database_name, database_oid); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewDropDatabase( + const char* database_name, PgOid database_oid, PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, database_name, database_oid), handle); } -Status PgApiImpl::ExecDropDatabase(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_DROP_DATABASE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecDropDatabase(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } -Status PgApiImpl::NewAlterDatabase(const char *database_name, - PgOid database_oid, - PgStatement **handle) { - auto stmt = std::make_unique(pg_session_, database_name, database_oid); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewAlterDatabase( + const char* database_name, PgOid database_oid, PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, database_name, database_oid), handle); } -Status PgApiImpl::AlterDatabaseRenameDatabase(PgStatement *handle, const char *newname) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_DATABASE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - down_cast(handle)->RenameDatabase(newname); +Status PgApiImpl::AlterDatabaseRenameDatabase(PgStatement* handle, const char* newname) { + VERIFY_RESULT_REF(GetStatementAs(handle)).RenameDatabase(newname); return Status::OK(); } -Status PgApiImpl::ExecAlterDatabase(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_DATABASE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecAlterDatabase(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } Status PgApiImpl::ReserveOids(const PgOid database_oid, @@ -705,57 +880,42 @@ void PgApiImpl::InvalidateTableCache(const PgObjectId& table_id) { //-------------------------------------------------------------------------------------------------- -Status PgApiImpl::NewCreateTablegroup(const char *database_name, - const PgOid database_oid, - const PgOid tablegroup_oid, - const PgOid tablespace_oid, - PgStatement **handle) { +Status PgApiImpl::NewCreateTablegroup( + const char* database_name, PgOid database_oid, PgOid tablegroup_oid, PgOid tablespace_oid, + PgStatement** handle) { SCHECK(pg_txn_manager_->IsDdlMode(), IllegalState, "Tablegroup is being created outside of DDL mode"); - auto stmt = std::make_unique(pg_session_, database_name, - database_oid, tablegroup_oid, tablespace_oid); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); + return AddToCurrentPgMemctx( + std::make_unique( + pg_session_, database_name, database_oid, tablegroup_oid, tablespace_oid), + handle); } -Status PgApiImpl::ExecCreateTablegroup(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_TABLEGROUP)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - pg_session_->SetDdlHasSyscatalogChanges(); - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecCreateTablegroup(PgStatement* handle) { + return ExecDdlWithSyscatalogChanges(handle, *pg_session_); } -Status PgApiImpl::NewDropTablegroup(const PgOid database_oid, - const PgOid tablegroup_oid, - PgStatement **handle) { +Status PgApiImpl::NewDropTablegroup( + PgOid database_oid, PgOid tablegroup_oid, PgStatement** handle) { SCHECK(pg_txn_manager_->IsDdlMode(), IllegalState, "Tablegroup is being dropped outside of DDL mode"); - auto stmt = std::make_unique(pg_session_, database_oid, tablegroup_oid); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, database_oid, tablegroup_oid), handle); } -Status PgApiImpl::ExecDropTablegroup(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_DROP_TABLEGROUP)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - pg_session_->SetDdlHasSyscatalogChanges(); - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecDropTablegroup(PgStatement* handle) { + return ExecDdlWithSyscatalogChanges(handle, *pg_session_); } //-------------------------------------------------------------------------------------------------- -Status PgApiImpl::NewCreateTable(const char *database_name, - const char *schema_name, - const char *table_name, +Status PgApiImpl::NewCreateTable(const char* database_name, + const char* schema_name, + const char* table_name, const PgObjectId& table_id, bool is_shared_table, bool is_sys_catalog_table, @@ -769,227 +929,120 @@ Status PgApiImpl::NewCreateTable(const char *database_name, const PgObjectId& pg_table_oid, const PgObjectId& old_relfilenode_oid, bool is_truncate, - PgStatement **handle) { - auto stmt = std::make_unique( - pg_session_, database_name, schema_name, table_name, table_id, is_shared_table, - is_sys_catalog_table, if_not_exist, ybrowid_mode, is_colocated_via_database, - tablegroup_oid, colocation_id, tablespace_oid, is_matview, pg_table_oid, - old_relfilenode_oid, is_truncate); - if (pg_txn_manager_->IsDdlMode()) { - stmt->UseTransaction(); - } - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); + PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique( + pg_session_, database_name, schema_name, table_name, table_id, is_shared_table, + is_sys_catalog_table, if_not_exist, ybrowid_mode, is_colocated_via_database, + tablegroup_oid, colocation_id, tablespace_oid, is_matview, pg_table_oid, + old_relfilenode_oid, is_truncate, pg_txn_manager_->IsDdlMode()), + handle); } -Status PgApiImpl::CreateTableAddColumn(PgStatement *handle, const char *attr_name, int attr_num, - const YBCPgTypeEntity *attr_type, - bool is_hash, bool is_range, - bool is_desc, bool is_nulls_first) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return AddColumn(down_cast(handle), attr_name, attr_num, attr_type, - is_hash, is_range, is_desc, is_nulls_first); +Status PgApiImpl::CreateTableAddColumn( + PgStatement* handle, const char* attr_name, int attr_num, const YBCPgTypeEntity* attr_type, + bool is_hash, bool is_range, bool is_desc, bool is_nulls_first) { + return AddColumn( + VERIFY_RESULT_REF(GetStatementAs(handle)), + attr_name, attr_num, attr_type, is_hash, is_range, is_desc, is_nulls_first); } -Status PgApiImpl::CreateTableSetNumTablets(PgStatement *handle, int32_t num_tablets) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return down_cast(handle)->SetNumTablets(num_tablets); +Status PgApiImpl::CreateTableSetNumTablets(PgStatement* handle, int32_t num_tablets) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).SetNumTablets(num_tablets); } -Status PgApiImpl::AddSplitBoundary(PgStatement *handle, PgExpr **exprs, int expr_count) { - // Partitioning a TABLE or an INDEX. - if (PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_TABLE) || - PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_INDEX)) { - return down_cast(handle)->AddSplitBoundary(exprs, expr_count); - } - - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); +Status PgApiImpl::AddSplitBoundary(PgStatement* handle, PgExpr** exprs, int expr_count) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).AddSplitBoundary( + exprs, expr_count); } -Status PgApiImpl::ExecCreateTable(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - pg_session_->SetDdlHasSyscatalogChanges(); - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecCreateTable(PgStatement* handle) { + return ExecDdlWithSyscatalogChanges(handle, *pg_session_); } -Status PgApiImpl::NewAlterTable(const PgObjectId& table_id, - PgStatement **handle) { - auto stmt = std::make_unique(pg_session_, table_id); - if (pg_txn_manager_->IsDdlMode()) { - stmt->UseTransaction(); - } - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewAlterTable(const PgObjectId& table_id, PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, table_id, pg_txn_manager_->IsDdlMode()), handle); } -Status PgApiImpl::AlterTableAddColumn(PgStatement *handle, const char *name, - int order, - const YBCPgTypeEntity *attr_type, - YBCPgExpr missing_value) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->AddColumn(name, attr_type, order, missing_value); +Status PgApiImpl::AlterTableAddColumn( + PgStatement* handle, const char* name, int order, const YBCPgTypeEntity* attr_type, + YBCPgExpr missing_value) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).AddColumn( + name, attr_type, order, missing_value); } -Status PgApiImpl::AlterTableRenameColumn(PgStatement *handle, const char *oldname, - const char *newname) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->RenameColumn(oldname, newname); +Status PgApiImpl::AlterTableRenameColumn( + PgStatement* handle, const char* oldname, const char* newname) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).RenameColumn(oldname, newname); } -Status PgApiImpl::AlterTableDropColumn(PgStatement *handle, const char *name) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->DropColumn(name); +Status PgApiImpl::AlterTableDropColumn(PgStatement* handle, const char* name) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).DropColumn(name); } -Status PgApiImpl::AlterTableSetReplicaIdentity(PgStatement *handle, const char identity_type) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->SetReplicaIdentity(identity_type); +Status PgApiImpl::AlterTableSetReplicaIdentity(PgStatement* handle, char identity_type) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).SetReplicaIdentity(identity_type); } -Status PgApiImpl::AlterTableRenameTable(PgStatement *handle, const char *db_name, - const char *newname) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->RenameTable(db_name, newname); +Status PgApiImpl::AlterTableRenameTable( + PgStatement* handle, const char* db_name, const char* newname) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).RenameTable(db_name, newname); } -Status PgApiImpl::AlterTableIncrementSchemaVersion(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->IncrementSchemaVersion(); +Status PgApiImpl::AlterTableIncrementSchemaVersion(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).IncrementSchemaVersion(); } -Status PgApiImpl::AlterTableSetTableId(PgStatement *handle, const PgObjectId &table_id) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->SetTableId(table_id); +Status PgApiImpl::AlterTableSetTableId(PgStatement* handle, const PgObjectId& table_id) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).SetTableId(table_id); } -Status PgApiImpl::AlterTableSetSchema(PgStatement *handle, const char *schema_name) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->SetSchema(schema_name); +Status PgApiImpl::AlterTableSetSchema(PgStatement* handle, const char* schema_name) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).SetSchema(schema_name); } -Status PgApiImpl::ExecAlterTable(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - pg_session_->SetDdlHasSyscatalogChanges(); - PgAlterTable *pg_stmt = down_cast(handle); - return pg_stmt->Exec(); +Status PgApiImpl::ExecAlterTable(PgStatement* handle) { + return ExecDdlWithSyscatalogChanges(handle, *pg_session_); } -Status PgApiImpl::AlterTableInvalidateTableCacheEntry(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_ALTER_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - PgAlterTable *pg_stmt = down_cast(handle); - pg_stmt->InvalidateTableCacheEntry(); +Status PgApiImpl::AlterTableInvalidateTableCacheEntry(PgStatement* handle) { + VERIFY_RESULT_REF(GetStatementAs(handle)).InvalidateTableCacheEntry(); return Status::OK(); } -Status PgApiImpl::NewDropTable(const PgObjectId& table_id, - bool if_exist, - PgStatement **handle) { +Status PgApiImpl::NewDropTable(const PgObjectId& table_id, bool if_exist, PgStatement** handle) { SCHECK(pg_txn_manager_->IsDdlMode(), IllegalState, "Table is being dropped outside of DDL mode"); - auto stmt = std::make_unique(pg_session_, table_id, if_exist); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, table_id, if_exist), handle); } -Status PgApiImpl::NewTruncateTable(const PgObjectId& table_id, - PgStatement **handle) { - auto stmt = std::make_unique(pg_session_, table_id); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewTruncateTable(const PgObjectId& table_id, PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, table_id), handle); } -Status PgApiImpl::ExecTruncateTable(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_TRUNCATE_TABLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecTruncateTable(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } -Status PgApiImpl::NewDropSequence(const YBCPgOid database_oid, - const YBCPgOid sequence_oid, - PgStatement **handle) { - auto stmt = std::make_unique(pg_session_, database_oid, - sequence_oid); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewDropSequence(PgOid database_oid, PgOid sequence_oid, PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, database_oid, sequence_oid), handle); } -Status PgApiImpl::ExecDropSequence(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_DROP_SEQUENCE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - PgDropSequence *pg_stmt = down_cast(handle); - return pg_stmt->Exec(); +Status PgApiImpl::ExecDropSequence(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } -Status PgApiImpl::NewDropDBSequences(const YBCPgOid database_oid, - PgStatement **handle) { - auto stmt = std::make_unique(pg_session_, database_oid); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewDropDBSequences(PgOid database_oid, PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, database_oid), handle); } -Status PgApiImpl::GetTableDesc(const PgObjectId& table_id, - PgTableDesc **handle) { +Status PgApiImpl::GetTableDesc(const PgObjectId& table_id, PgTableDesc **handle) { // First read from memory context. size_t hash_id = hash_value(table_id); RETURN_NOT_OK(GetTabledescFromCurrentPgMemctx(hash_id, handle)); @@ -1010,66 +1063,44 @@ Result PgApiImpl::GetDatabaseClones() { return pg_session_->pg_client().ListDatabaseClones(); } -Result PgApiImpl::GetColumnInfo(YBCPgTableDesc table_desc, - int16_t attr_number) { +Result PgApiImpl::GetColumnInfo(YBCPgTableDesc table_desc, int16_t attr_number) { return table_desc->GetColumnInfo(attr_number); } -Status PgApiImpl::DmlModifiesRow(PgStatement *handle, bool *modifies_row) { - if (!handle) { - return STATUS(InvalidArgument, "Invalid statement handle"); - } +Result PgApiImpl::DmlModifiesRow(PgStatement* handle) { + auto& dml = VERIFY_RESULT_REF(GetStatementAs(handle)); - *modifies_row = false; - - switch (handle->stmt_op()) { - case StmtOp::STMT_UPDATE: - case StmtOp::STMT_DELETE: - *modifies_row = true; - break; + switch (dml.stmt_op()) { + case PgUpdate::kStmtOp: + case PgDelete::kStmtOp: + return true; default: break; } - return Status::OK(); + return false; } -Status PgApiImpl::SetIsSysCatalogVersionChange(PgStatement *handle) { - if (!handle) { - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - switch (handle->stmt_op()) { - case StmtOp::STMT_UPDATE: - case StmtOp::STMT_DELETE: - case StmtOp::STMT_INSERT: - down_cast(handle)->SetIsSystemCatalogChange(); - return Status::OK(); - default: - break; - } - - return STATUS(InvalidArgument, "Invalid statement handle"); +Status PgApiImpl::SetIsSysCatalogVersionChange(PgStatement* handle) { + VERIFY_RESULT_REF(GetStatementAs( + handle, + { + StatementTag(), + StatementTag(), + StatementTag() + })).SetIsSystemCatalogChange(); + return Status::OK(); } Status PgApiImpl::SetCatalogCacheVersion( - PgStatement *handle, uint64_t version, std::optional db_oid) { - if (!handle) { - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - switch (handle->stmt_op()) { - case StmtOp::STMT_SELECT: - case StmtOp::STMT_INSERT: - case StmtOp::STMT_UPDATE: - case StmtOp::STMT_DELETE: - down_cast(handle)->SetCatalogCacheVersion(db_oid, version); - return Status::OK(); - default: - break; - } - - return STATUS(InvalidArgument, "Invalid statement handle"); + PgStatement* handle, uint64_t version, std::optional db_oid) { + VERIFY_RESULT_REF(GetStatementAs(handle, { + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag() + })).SetCatalogCacheVersion(db_oid, version); + return Status::OK(); } Result PgApiImpl::GetTableDiskSize(const PgObjectId& table_oid) { @@ -1078,15 +1109,15 @@ Result PgApiImpl::GetTableDiskSize(const PgObjectId& tabl //-------------------------------------------------------------------------------------------------- -Status PgApiImpl::NewCreateIndex(const char *database_name, - const char *schema_name, - const char *index_name, +Status PgApiImpl::NewCreateIndex(const char* database_name, + const char* schema_name, + const char* index_name, const PgObjectId& index_id, const PgObjectId& base_table_id, bool is_shared_index, bool is_sys_catalog_index, bool is_unique_index, - const bool skip_index_backfill, + bool skip_index_backfill, bool if_not_exist, bool is_colocated_via_database, const PgObjectId& tablegroup_oid, @@ -1094,106 +1125,83 @@ Status PgApiImpl::NewCreateIndex(const char *database_name, const PgObjectId& tablespace_oid, const PgObjectId& pg_table_id, const PgObjectId& old_relfilenode_id, - PgStatement **handle) { - auto stmt = std::make_unique( - pg_session_, database_name, schema_name, index_name, index_id, is_shared_index, - is_sys_catalog_index, if_not_exist, PG_YBROWID_MODE_NONE, - is_colocated_via_database, tablegroup_oid, colocation_id, - tablespace_oid, false /* is_matview */, pg_table_id, old_relfilenode_id, - false /* is_truncate */); - stmt->SetupIndex(base_table_id, is_unique_index, skip_index_backfill); - if (pg_txn_manager_->IsDdlMode()) { - stmt->UseTransaction(); - } - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); + PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique( + pg_session_, database_name, schema_name, index_name, index_id, is_shared_index, + is_sys_catalog_index, if_not_exist, PG_YBROWID_MODE_NONE, is_colocated_via_database, + tablegroup_oid, colocation_id, tablespace_oid, false /* is_matview */, pg_table_id, + old_relfilenode_id, false /* is_truncate */, pg_txn_manager_->IsDdlMode(), base_table_id, + is_unique_index, skip_index_backfill), + handle); } -Status PgApiImpl::CreateIndexAddColumn(PgStatement *handle, const char *attr_name, int attr_num, - const YBCPgTypeEntity *attr_type, - bool is_hash, bool is_range, - bool is_desc, bool is_nulls_first) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_INDEX)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - - return AddColumn(down_cast(handle), attr_name, attr_num, attr_type, - is_hash, is_range, is_desc, is_nulls_first); +Status PgApiImpl::CreateIndexAddColumn( + PgStatement* handle, const char* attr_name, int attr_num, const YBCPgTypeEntity* attr_type, + bool is_hash, bool is_range, bool is_desc, bool is_nulls_first) { + return AddColumn( + VERIFY_RESULT_REF(GetStatementAs(handle)), + attr_name, attr_num, attr_type, is_hash, is_range, is_desc, is_nulls_first); } -Status PgApiImpl::CreateIndexSetNumTablets(PgStatement *handle, int32_t num_tablets) { - SCHECK(PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_INDEX), - InvalidArgument, - "Invalid statement handle"); - return down_cast(handle)->SetNumTablets(num_tablets); +Status PgApiImpl::CreateIndexSetNumTablets(PgStatement* handle, int32_t num_tablets) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).SetNumTablets(num_tablets); } -Status PgApiImpl::CreateIndexSetVectorOptions(PgStatement *handle, YbPgVectorIdxOptions *options) { - SCHECK(PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_INDEX), - InvalidArgument, - "Invalid statement handle"); - return down_cast(handle)->SetVectorOptions(options); +Status PgApiImpl::CreateIndexSetVectorOptions(PgStatement* handle, YbPgVectorIdxOptions* options) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).SetVectorOptions(options); } -Status PgApiImpl::ExecCreateIndex(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_INDEX)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - pg_session_->SetDdlHasSyscatalogChanges(); - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecCreateIndex(PgStatement* handle) { + return ExecDdlWithSyscatalogChanges(handle, *pg_session_); } -Status PgApiImpl::NewDropIndex(const PgObjectId& index_id, - bool if_exist, - bool ddl_rollback_enabled, - PgStatement **handle) { - SCHECK(pg_txn_manager_->IsDdlMode(), - IllegalState, - "Index is being dropped outside of DDL mode"); - auto stmt = std::make_unique(pg_session_, index_id, if_exist, ddl_rollback_enabled); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewDropIndex( + const PgObjectId& index_id, bool if_exist, bool ddl_rollback_enabled, PgStatement** handle) { + SCHECK(pg_txn_manager_->IsDdlMode(), IllegalState, "Index is being dropped outside of DDL mode"); + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, index_id, if_exist, ddl_rollback_enabled), handle); } -Status PgApiImpl::ExecPostponedDdlStmt(PgStatement *handle) { - if (!handle) { - return STATUS(InvalidArgument, "Invalid statement handle"); - } +Status PgApiImpl::ExecPostponedDdlStmt(PgStatement* handle) { + auto& ddl = VERIFY_RESULT_REF(GetStatementAs( + handle, + { + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag(), + StatementTag() + })); + + const auto stmt_op = ddl.stmt_op(); - switch (handle->stmt_op()) { - case StmtOp::STMT_DROP_TABLE: - return down_cast(handle)->Exec(); - case StmtOp::STMT_DROP_INDEX: - return down_cast(handle)->Exec(); - case StmtOp::STMT_DROP_TABLEGROUP: - return down_cast(handle)->Exec(); - case StmtOp::STMT_DROP_SEQUENCE: - return down_cast(handle)->Exec(); - case StmtOp::STMT_DROP_DB_SEQUENCES: - return down_cast(handle)->Exec(); +#define YB_EXECUTE_DDL_STMT_CASE(name) \ + case name::kStmtOp: return down_cast(ddl).Exec() + + switch (stmt_op) { + YB_EXECUTE_DDL_STMT_CASE(PgDropTable); + YB_EXECUTE_DDL_STMT_CASE(PgDropIndex); + YB_EXECUTE_DDL_STMT_CASE(PgDropTablegroup); + YB_EXECUTE_DDL_STMT_CASE(PgDropSequence); + YB_EXECUTE_DDL_STMT_CASE(PgDropDBSequences); default: break; } - return STATUS(InvalidArgument, "Invalid statement handle"); + +#undef YB_EXECUTE_PG_STMT + + RSTATUS_DCHECK(false, IllegalState, "Unexpected stmt_op $0", to_underlying(stmt_op)); + return Status::OK(); } -Status PgApiImpl::ExecDropTable(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_DROP_TABLE)) { - return STATUS(InvalidArgument, "Invalid statement handle"); - } - pg_session_->SetDdlHasSyscatalogChanges(); - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecDropTable(PgStatement* handle) { + return ExecDdlWithSyscatalogChanges(handle, *pg_session_); } -Status PgApiImpl::ExecDropIndex(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_DROP_INDEX)) { - return STATUS(InvalidArgument, "Invalid statement handle"); - } - pg_session_->SetDdlHasSyscatalogChanges(); - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecDropIndex(PgStatement* handle) { + return ExecDdlWithSyscatalogChanges(handle, *pg_session_); } Result PgApiImpl::WaitForBackendsCatalogVersion(PgOid dboid, uint64_t version) { @@ -1220,95 +1228,86 @@ Status PgApiImpl::BackfillIndex(const PgObjectId& table_id) { // DML Statement Support. //-------------------------------------------------------------------------------------------------- -// Binding ----------------------------------------------------------------------------------------- - -Status PgApiImpl::DmlAppendTarget(PgStatement *handle, PgExpr *target) { - return down_cast(handle)->AppendTarget(target); +Status PgApiImpl::DmlAppendTarget(PgStatement* handle, PgExpr* target) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).AppendTarget(target); } -Status PgApiImpl::DmlAppendQual(PgStatement *handle, PgExpr *qual, bool is_primary) { - return down_cast(handle)->AppendQual(qual, is_primary); +Status PgApiImpl::DmlAppendQual(PgStatement* handle, PgExpr* qual, bool is_primary) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).AppendQual(qual, is_primary); } -Status PgApiImpl::DmlAppendColumnRef(PgStatement *handle, PgColumnRef *colref, bool is_primary) { - return down_cast(handle)->AppendColumnRef(colref, is_primary); +Status PgApiImpl::DmlAppendColumnRef(PgStatement* handle, PgColumnRef* colref, bool is_primary) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).AppendColumnRef(colref, is_primary); } -Status PgApiImpl::DmlBindColumn(PgStatement *handle, int attr_num, PgExpr *attr_value) { - return down_cast(handle)->BindColumn(attr_num, attr_value); +Status PgApiImpl::DmlBindColumn(PgStatement* handle, int attr_num, PgExpr* attr_value) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).BindColumn(attr_num, attr_value); } Status PgApiImpl::DmlBindRow( - YBCPgStatement handle, uint64_t ybctid, YBCBindColumn* columns, int count) { - return down_cast(handle)->BindRow(ybctid, columns, count); + PgStatement* handle, uint64_t ybctid, YBCBindColumn* columns, int count) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).BindRow(ybctid, columns, count); } -Status PgApiImpl::DmlBindColumnCondBetween(PgStatement *handle, int attr_num, - PgExpr *attr_value, - bool start_inclusive, - PgExpr *attr_value_end, - bool end_inclusive) { - return down_cast(handle)->BindColumnCondBetween(attr_num, - attr_value, - start_inclusive, - attr_value_end, - end_inclusive); +Status PgApiImpl::DmlBindColumnCondBetween( + PgStatement* handle, int attr_num, PgExpr* attr_value, bool start_inclusive, + PgExpr* attr_value_end, bool end_inclusive) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).BindColumnCondBetween( + attr_num, attr_value, start_inclusive, attr_value_end, end_inclusive); } -Status PgApiImpl::DmlBindColumnCondIsNotNull(PgStatement *handle, int attr_num) { - return down_cast(handle)->BindColumnCondIsNotNull(attr_num); +Status PgApiImpl::DmlBindColumnCondIsNotNull(PgStatement* handle, int attr_num) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).BindColumnCondIsNotNull(attr_num); } -Status PgApiImpl::DmlBindColumnCondIn(PgStatement *handle, YBCPgExpr lhs, int n_attr_values, - PgExpr **attr_values) { - return down_cast(handle)->BindColumnCondIn(lhs, n_attr_values, attr_values); +Status PgApiImpl::DmlBindColumnCondIn( + PgStatement* handle, YBCPgExpr lhs, int n_attr_values, PgExpr** attr_values) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).BindColumnCondIn( + lhs, n_attr_values, attr_values); } -Status PgApiImpl::DmlAddRowUpperBound(YBCPgStatement handle, - int n_col_values, PgExpr **col_values, bool is_inclusive) { - return down_cast(handle)->AddRowUpperBound(handle, - n_col_values, - col_values, - is_inclusive); +Status PgApiImpl::DmlAddRowUpperBound( + PgStatement* handle, int n_col_values, PgExpr** col_values, bool is_inclusive) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).AddRowUpperBound( + handle, n_col_values, col_values, is_inclusive); } -Status PgApiImpl::DmlAddRowLowerBound(YBCPgStatement handle, - int n_col_values, PgExpr **col_values, bool is_inclusive) { - return down_cast(handle)->AddRowLowerBound(handle, - n_col_values, - col_values, - is_inclusive); +Status PgApiImpl::DmlAddRowLowerBound( + PgStatement* handle, int n_col_values, PgExpr** col_values, bool is_inclusive) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).AddRowLowerBound( + handle, n_col_values, col_values, is_inclusive); } -void PgApiImpl::DmlBindHashCode( - PgStatement* handle, const std::optional& start, const std::optional& end) { - down_cast(handle)->BindHashCode(start, end); +Status PgApiImpl::DmlBindHashCode( + PgStatement* handle, const std::optional& start, const std::optional& end) { + VERIFY_RESULT_REF(GetStatementAs(handle)).BindHashCode(start, end); + return Status::OK(); } -Status PgApiImpl::DmlBindRange(YBCPgStatement handle, - Slice lower_bound, - bool lower_bound_inclusive, - Slice upper_bound, - bool upper_bound_inclusive) { - return down_cast(handle)->BindRange( +Status PgApiImpl::DmlBindRange( + PgStatement* handle, Slice lower_bound, bool lower_bound_inclusive, Slice upper_bound, + bool upper_bound_inclusive) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).BindRange( lower_bound, lower_bound_inclusive, upper_bound, upper_bound_inclusive); } -Status PgApiImpl::DmlBindTable(PgStatement *handle) { - return down_cast(handle)->BindTable(); +Status PgApiImpl::DmlBindTable(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).BindTable(); } -Result PgApiImpl::DmlGetColumnInfo(YBCPgStatement handle, int attr_num) { - return down_cast(handle)->GetColumnInfo(attr_num); +Result PgApiImpl::DmlGetColumnInfo(PgStatement* handle, int attr_num) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).GetColumnInfo(attr_num); } -Status PgApiImpl::DmlAssignColumn(PgStatement *handle, int attr_num, PgExpr *attr_value) { - return down_cast(handle)->AssignColumn(attr_num, attr_value); +Status PgApiImpl::DmlAssignColumn(PgStatement* handle, int attr_num, PgExpr* attr_value) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).AssignColumn(attr_num, attr_value); } -Status PgApiImpl::DmlFetch(PgStatement *handle, int32_t natts, uint64_t *values, bool *isnulls, - PgSysColumns *syscols, bool *has_data) { - return down_cast(handle)->Fetch(natts, values, isnulls, syscols, has_data); +Status PgApiImpl::DmlFetch( + PgStatement* handle, int32_t natts, uint64_t* values, bool* isnulls, PgSysColumns* syscols, + bool* has_data) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Fetch( + natts, values, isnulls, syscols, has_data); } Result PgApiImpl::BuildTupleId(const YBCPgYBTupleIdDescriptor& descr) { @@ -1332,23 +1331,12 @@ Status PgApiImpl::FlushBufferedOperations() { } Status PgApiImpl::DmlExecWriteOp(PgStatement *handle, int32_t *rows_affected_count) { - switch (handle->stmt_op()) { - case StmtOp::STMT_INSERT: - case StmtOp::STMT_UPDATE: - case StmtOp::STMT_DELETE: - case StmtOp::STMT_TRUNCATE: - { - auto dml_write = down_cast(handle); - RETURN_NOT_OK(dml_write->Exec(ForceNonBufferable(rows_affected_count != nullptr))); - if (rows_affected_count) { - *rows_affected_count = dml_write->GetRowsAffectedCount(); - } - return Status::OK(); - } - default: - break; + auto& dml_write = VERIFY_RESULT_REF(GetStatementAs(handle)); + RETURN_NOT_OK(dml_write.Exec(ForceNonBufferable{rows_affected_count != nullptr})); + if (rows_affected_count) { + *rows_affected_count = dml_write.GetRowsAffectedCount(); } - return STATUS(InvalidArgument, "Invalid statement handle"); + return Status::OK(); } // Insert ------------------------------------------------------------------------------------------ @@ -1369,10 +1357,9 @@ Result PgApiImpl::NewInsertBlock( return result; } -Status PgApiImpl::NewInsert(const PgObjectId& table_id, - bool is_region_local, - PgStatement **handle, - YBCPgTransactionSetting transaction_setting) { +Status PgApiImpl::NewInsert( + const PgObjectId& table_id, bool is_region_local, PgStatement **handle, + YBCPgTransactionSetting transaction_setting) { *handle = nullptr; return AddToCurrentPgMemctx( VERIFY_RESULT(PgInsert::Make( @@ -1380,80 +1367,52 @@ Status PgApiImpl::NewInsert(const PgObjectId& table_id, handle); } -Status PgApiImpl::ExecInsert(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_INSERT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecInsert(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } -Status PgApiImpl::InsertStmtSetUpsertMode(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_INSERT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - down_cast(handle)->SetUpsertMode(); - +Status PgApiImpl::InsertStmtSetUpsertMode(PgStatement* handle) { + VERIFY_RESULT_REF(GetStatementAs(handle)).SetUpsertMode(); return Status::OK(); } -Status PgApiImpl::InsertStmtSetWriteTime(PgStatement *handle, const HybridTime write_time) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_INSERT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - RETURN_NOT_OK(down_cast(handle)->SetWriteTime(write_time)); - return Status::OK(); +Status PgApiImpl::InsertStmtSetWriteTime(PgStatement* handle, HybridTime write_time) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).SetWriteTime(write_time); } -Status PgApiImpl::InsertStmtSetIsBackfill(PgStatement *handle, const bool is_backfill) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_INSERT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - down_cast(handle)->SetIsBackfill(is_backfill); +Status PgApiImpl::InsertStmtSetIsBackfill(PgStatement* handle, bool is_backfill) { + VERIFY_RESULT_REF(GetStatementAs(handle)).SetIsBackfill(is_backfill); return Status::OK(); } // Update ------------------------------------------------------------------------------------------ -Status PgApiImpl::NewUpdate(const PgObjectId& table_id, - bool is_region_local, - PgStatement **handle, - YBCPgTransactionSetting transaction_setting) { +Status PgApiImpl::NewUpdate( + const PgObjectId& table_id, bool is_region_local, PgStatement** handle, + YBCPgTransactionSetting transaction_setting) { *handle = nullptr; return AddToCurrentPgMemctx( VERIFY_RESULT(PgUpdate::Make(pg_session_, table_id, is_region_local, transaction_setting)), handle); } -Status PgApiImpl::ExecUpdate(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_UPDATE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecUpdate(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } // Delete ------------------------------------------------------------------------------------------ -Status PgApiImpl::NewDelete(const PgObjectId& table_id, - bool is_region_local, - PgStatement **handle, - YBCPgTransactionSetting transaction_setting) { +Status PgApiImpl::NewDelete( + const PgObjectId& table_id, bool is_region_local, PgStatement** handle, + YBCPgTransactionSetting transaction_setting) { *handle = nullptr; return AddToCurrentPgMemctx( VERIFY_RESULT(PgDelete::Make(pg_session_, table_id, is_region_local, transaction_setting)), handle); } -Status PgApiImpl::ExecDelete(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_DELETE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecDelete(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } Status PgApiImpl::NewSample( @@ -1467,43 +1426,27 @@ Status PgApiImpl::NewSample( } Result PgApiImpl::SampleNextBlock(PgStatement* handle) { - RSTATUS_DCHECK( - PgStatement::IsValidStmt(handle, StmtOp::STMT_SAMPLE), - InvalidArgument, "Invalid statement handle"); - return down_cast(handle)->SampleNextBlock(); + return VERIFY_RESULT_REF(GetStatementAs(handle)).SampleNextBlock(); } -Status PgApiImpl::ExecSample(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_SAMPLE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - RETURN_NOT_OK(down_cast(handle)->Exec(nullptr)); - return Status::OK(); +Status PgApiImpl::ExecSample(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(nullptr); } Result PgApiImpl::GetEstimatedRowCount(PgStatement* handle) { - RSTATUS_DCHECK( - PgStatement::IsValidStmt(handle, StmtOp::STMT_SAMPLE), - InvalidArgument, "Invalid statement handle"); - return down_cast(handle)->GetEstimatedRowCount(); + return VERIFY_RESULT_REF(GetStatementAs(handle)).GetEstimatedRowCount(); } -Status PgApiImpl::DeleteStmtSetIsPersistNeeded(PgStatement *handle, const bool is_persist_needed) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_DELETE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - down_cast(handle)->SetIsPersistNeeded(is_persist_needed); +Status PgApiImpl::DeleteStmtSetIsPersistNeeded(PgStatement* handle, bool is_persist_needed) { + VERIFY_RESULT_REF(GetStatementAs(handle)).SetIsPersistNeeded(is_persist_needed); return Status::OK(); } // Colocated Truncate ------------------------------------------------------------------------------ -Status PgApiImpl::NewTruncateColocated(const PgObjectId& table_id, - bool is_region_local, - PgStatement **handle, - YBCPgTransactionSetting transaction_setting) { +Status PgApiImpl::NewTruncateColocated( + const PgObjectId& table_id, bool is_region_local, PgStatement** handle, + YBCPgTransactionSetting transaction_setting) { *handle = nullptr; return AddToCurrentPgMemctx( VERIFY_RESULT(PgTruncateColocated::Make( @@ -1511,12 +1454,8 @@ Status PgApiImpl::NewTruncateColocated(const PgObjectId& table_id, handle); } -Status PgApiImpl::ExecTruncateColocated(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_TRUNCATE)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - return down_cast(handle)->Exec(); +Status PgApiImpl::ExecTruncateColocated(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } // Select ------------------------------------------------------------------------------------------ @@ -1531,30 +1470,19 @@ Status PgApiImpl::NewSelect( handle); } -Status PgApiImpl::SetForwardScan(PgStatement *handle, bool is_forward_scan) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_SELECT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - down_cast(handle)->SetForwardScan(is_forward_scan); +Status PgApiImpl::SetForwardScan(PgStatement* handle, bool is_forward_scan) { + VERIFY_RESULT_REF(GetStatementAs(handle)).SetForwardScan(is_forward_scan); return Status::OK(); } -Status PgApiImpl::SetDistinctPrefixLength(PgStatement *handle, int distinct_prefix_length) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_SELECT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - down_cast(handle)->SetDistinctPrefixLength(distinct_prefix_length); +Status PgApiImpl::SetDistinctPrefixLength(PgStatement* handle, int distinct_prefix_length) { + VERIFY_RESULT_REF(GetStatementAs(handle)).SetDistinctPrefixLength( + distinct_prefix_length); return Status::OK(); } -Status PgApiImpl::SetHashBounds(PgStatement *handle, uint16_t low_bound, uint16_t high_bound) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_SELECT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - down_cast(handle)->SetHashBounds(low_bound, high_bound); +Status PgApiImpl::SetHashBounds(PgStatement* handle, uint16_t low_bound, uint16_t high_bound) { + VERIFY_RESULT_REF(GetStatementAs(handle)).SetHashBounds(low_bound, high_bound); return Status::OK(); } @@ -1565,25 +1493,21 @@ Slice PgApiImpl::GetYbctidAsSlice(uint64_t ybctid) { return Slice(value, bytes); } -Status PgApiImpl::RetrieveYbctids(PgStatement *handle, const YBCPgExecParameters *exec_params, - int natts, SliceVector *ybctids, size_t *count, - bool *exceeded_work_mem) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_SELECT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - auto& dml_read = *down_cast(handle); +Status PgApiImpl::RetrieveYbctids( + PgStatement* handle, const YBCPgExecParameters* exec_params, int natts, SliceVector* ybctids, + size_t* count, bool* exceeded_work_mem) { + auto& select = VERIFY_RESULT_REF(GetStatementAs(handle)); auto vec = new std::vector(); // If there's a secondary index request (secondary index scans or primary key scans on // non-colocated tables) return the ybctids from the secondary index request. - if (dml_read.has_secondary_index_with_doc_op()) { - RETURN_NOT_OK(dml_read.RetrieveYbctidsFromSecondaryIndex(exec_params, vec, - exceeded_work_mem)); + if (select.has_secondary_index_with_doc_op()) { + RETURN_NOT_OK(select.RetrieveYbctidsFromSecondaryIndex( + exec_params, vec, exceeded_work_mem)); } else { // Otherwise (colocated case): fetch the results one at a time. - RETURN_NOT_OK(dml_read.Exec(exec_params)); + RETURN_NOT_OK(select.Exec(exec_params)); uint64_t *values = new uint64_t[natts]; bool *nulls = new bool[natts]; @@ -1592,7 +1516,7 @@ Status PgApiImpl::RetrieveYbctids(PgStatement *handle, const YBCPgExecParameters size_t consumed_bytes = 0; do { - RETURN_NOT_OK(dml_read.Fetch(natts, values, nulls, &syscols, &has_data)); + RETURN_NOT_OK(select.Fetch(natts, values, nulls, &syscols, &has_data)); if (has_data && syscols.ybctid != NULL) { Slice s = GetYbctidAsSlice((uintptr_t) syscols.ybctid); @@ -1622,47 +1546,37 @@ Status PgApiImpl::RetrieveYbctids(PgStatement *handle, const YBCPgExecParameters return Status::OK(); } -Status PgApiImpl::FetchRequestedYbctids(PgStatement *handle, const PgExecParameters *exec_params, - ConstSliceVector ybctids) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_SELECT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - auto& dml_read = *down_cast(handle); - - RETURN_NOT_OK(dml_read.SetRequestedYbctids((const std::vector *) ybctids)); - - return dml_read.Exec(exec_params); +Status PgApiImpl::FetchRequestedYbctids( + PgStatement* handle, const PgExecParameters* exec_params, ConstSliceVector ybctids) { + auto& select = VERIFY_RESULT_REF(GetStatementAs(handle)); + RETURN_NOT_OK(select.SetRequestedYbctids((const std::vector *) ybctids)); + return select.Exec(exec_params); } -Status PgApiImpl::DmlANNBindVector(PgStatement *handle, int vec_att_no, PgExpr *vector) { - return down_cast(handle)->ANNBindVector(vec_att_no, vector); +Status PgApiImpl::DmlANNBindVector(PgStatement* handle, int vec_att_no, PgExpr* vector) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).ANNBindVector(vec_att_no, vector); } -Status PgApiImpl::DmlANNSetPrefetchSize(PgStatement *handle, int prefetch_size) { - return down_cast(handle)->ANNSetPrefetchSize(prefetch_size); +Status PgApiImpl::DmlANNSetPrefetchSize(PgStatement* handle, int prefetch_size) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).ANNSetPrefetchSize(prefetch_size); } -Status PgApiImpl::ExecSelect(PgStatement *handle, const PgExecParameters *exec_params) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_SELECT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - auto& dml_read = *down_cast(handle); - if (pg_sys_table_prefetcher_ && dml_read.IsReadFromYsqlCatalog() && dml_read.read_req()) { - // In case of sys tables prefething is enabled all reads from sys table must use cached data. +Status PgApiImpl::ExecSelect(PgStatement* handle, const PgExecParameters* exec_params) { + auto& select = VERIFY_RESULT_REF(GetStatementAs(handle)); + if (pg_sys_table_prefetcher_ && select.IsReadFromYsqlCatalog() && select.read_req()) { + // In case of sys tables prefetching is enabled all reads from sys table must use cached data. auto data = pg_sys_table_prefetcher_->GetData( - *dml_read.read_req(), dml_read.IsIndexOrderedScan()); + *select.read_req(), select.IsIndexOrderedScan()); if (!data) { // LOG(DFATAL) is used instead of SCHECK to let user on release build proceed by reading // data from a master in a non efficient way (by using separate RPC). LOG(DFATAL) << "Data was not prefetched for request " - << dml_read.read_req()->ShortDebugString(); + << select.read_req()->ShortDebugString(); } else { - dml_read.UpgradeDocOp(MakeDocReadOpWithData(pg_session_, std::move(data))); + select.UpgradeDocOp(MakeDocReadOpWithData(pg_session_, std::move(data))); } } - return dml_read.Exec(exec_params); + return select.Exec(exec_params); } @@ -1670,12 +1584,10 @@ Status PgApiImpl::ExecSelect(PgStatement *handle, const PgExecParameters *exec_p // Functions. //-------------------------------------------------------------------------------------------------- -Status PgApiImpl::NewSRF( - PgFunction **handle, PgFunctionDataProcessor processor) { +Status PgApiImpl::NewSRF(PgFunction** handle, PgFunctionDataProcessor processor) { *handle = nullptr; - auto func = std::make_unique(std::move(processor), pg_session_); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(func), handle)); - return Status::OK(); + return AddToCurrentPgMemctx( + std::make_unique(std::move(processor), pg_session_), handle); } Status PgApiImpl::AddFunctionParam( @@ -1725,55 +1637,41 @@ Status PgApiImpl::NewGetLockStatusDataSRF(PgFunction **handle) { // Column references ------------------------------------------------------------------------------- Status PgApiImpl::NewColumnRef( - PgStatement *stmt, int attr_num, const PgTypeEntity *type_entity, bool collate_is_valid_non_c, - const PgTypeAttrs *type_attrs, PgExpr **expr_handle) { - if (!stmt) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } + PgStatement* stmt, int attr_num, const PgTypeEntity* type_entity, bool collate_is_valid_non_c, + const PgTypeAttrs* type_attrs, PgExpr** expr_handle) { *expr_handle = PgColumnRef::Create( - &stmt->arena(), attr_num, type_entity, collate_is_valid_non_c, type_attrs); - + &VERIFY_RESULT_REF(GetArena(stmt)), attr_num, type_entity, collate_is_valid_non_c, type_attrs); return Status::OK(); } // Constant ---------------------------------------------------------------------------------------- Status PgApiImpl::NewConstant( - YBCPgStatement stmt, const YBCPgTypeEntity *type_entity, bool collate_is_valid_non_c, - const char *collation_sortkey, uint64_t datum, bool is_null, YBCPgExpr *expr_handle) { - if (!stmt) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - *expr_handle = stmt->arena().NewObject( - &stmt->arena(), type_entity, collate_is_valid_non_c, collation_sortkey, datum, is_null); + PgStatement* stmt, const YBCPgTypeEntity* type_entity, bool collate_is_valid_non_c, + const char* collation_sortkey, uint64_t datum, bool is_null, YBCPgExpr* expr_handle) { + auto& arena = VERIFY_RESULT_REF(GetArena(stmt)); + *expr_handle = arena.NewObject( + &arena, type_entity, collate_is_valid_non_c, collation_sortkey, datum, is_null); return Status::OK(); } Status PgApiImpl::NewConstantVirtual( - YBCPgStatement stmt, const YBCPgTypeEntity *type_entity, - YBCPgDatumKind datum_kind, YBCPgExpr *expr_handle) { - if (!stmt) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } + PgStatement* stmt, const YBCPgTypeEntity* type_entity, YBCPgDatumKind datum_kind, + YBCPgExpr* expr_handle) { + auto& arena = VERIFY_RESULT_REF(GetArena(stmt)); - *expr_handle = stmt->arena().NewObject( - &stmt->arena(), type_entity, false /* collate_is_valid_non_c */, datum_kind); + *expr_handle = arena.NewObject( + &arena, type_entity, false /* collate_is_valid_non_c */, datum_kind); return Status::OK(); } Status PgApiImpl::NewConstantOp( - YBCPgStatement stmt, const YBCPgTypeEntity *type_entity, bool collate_is_valid_non_c, - const char *collation_sortkey, uint64_t datum, bool is_null, YBCPgExpr *expr_handle, + PgStatement* stmt, const YBCPgTypeEntity* type_entity, bool collate_is_valid_non_c, + const char* collation_sortkey, uint64_t datum, bool is_null, YBCPgExpr* expr_handle, bool is_gt) { - if (!stmt) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - *expr_handle = stmt->arena().NewObject( - &stmt->arena(), type_entity, collate_is_valid_non_c, collation_sortkey, datum, is_null, + auto& arena = VERIFY_RESULT_REF(GetArena(stmt)); + *expr_handle = arena.NewObject( + &arena, type_entity, collate_is_valid_non_c, collation_sortkey, datum, is_null, is_gt ? PgExpr::Opcode::PG_EXPR_GT : PgExpr::Opcode::PG_EXPR_LT); return Status::OK(); @@ -1802,16 +1700,12 @@ Status PgApiImpl::UpdateConstant(PgExpr *expr, const void *value, int64_t bytes, // Text constant ----------------------------------------------------------------------------------- Status PgApiImpl::NewOperator( - PgStatement *stmt, const char *opname, const YBCPgTypeEntity *type_entity, - bool collate_is_valid_non_c, PgExpr **op_handle) { - if (!stmt) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } + PgStatement* stmt, const char* opname, const YBCPgTypeEntity* type_entity, + bool collate_is_valid_non_c, PgExpr** op_handle) { RETURN_NOT_OK(PgExpr::CheckOperatorName(opname)); - // Create operator. - *op_handle = PgOperator::Create(&stmt->arena(), opname, type_entity, collate_is_valid_non_c); + *op_handle = PgOperator::Create( + &VERIFY_RESULT_REF(GetArena(stmt)), opname, type_entity, collate_is_valid_non_c); return Status::OK(); } @@ -1920,16 +1814,12 @@ pid_t PgApiImpl::GetLocalTServerPid() const { // Tuple Expression ----------------------------------------------------------------------------- Status PgApiImpl::NewTupleExpr( - YBCPgStatement stmt, const YBCPgTypeEntity *tuple_type_entity, - const YBCPgTypeAttrs *type_attrs, int num_elems, - const YBCPgExpr *elems, YBCPgExpr *expr_handle) { - if (!stmt) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } + PgStatement* stmt, const YBCPgTypeEntity* tuple_type_entity, const YBCPgTypeAttrs* type_attrs, + int num_elems, const YBCPgExpr* elems, YBCPgExpr* expr_handle) { + auto& arena = VERIFY_RESULT_REF(GetArena(stmt)); - *expr_handle = stmt->arena().NewObject( - &stmt->arena(), tuple_type_entity, type_attrs, num_elems, elems); + *expr_handle = arena.NewObject( + &arena, tuple_type_entity, type_attrs, num_elems, elems); return Status::OK(); } @@ -2256,25 +2146,18 @@ void PgApiImpl::RestoreSessionState(const YBCPgSessionState& session_data) { //-------------------------------------------------------------------------------------------------- -Status PgApiImpl::NewCreateReplicationSlot(const char *slot_name, - const char *plugin_name, - const PgOid database_oid, - YBCPgReplicationSlotSnapshotAction snapshot_action, - PgStatement **handle) { - auto stmt = std::make_unique( - pg_session_, slot_name, plugin_name, database_oid, snapshot_action); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewCreateReplicationSlot( + const char* slot_name, const char* plugin_name, PgOid database_oid, + YBCPgReplicationSlotSnapshotAction snapshot_action, PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique( + pg_session_, slot_name, plugin_name, database_oid, snapshot_action), + handle); } Result PgApiImpl::ExecCreateReplicationSlot( - PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_CREATE_REPLICATION_SLOT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - PgCreateReplicationSlot *pg_stmt = down_cast(handle); - return pg_stmt->Exec(); + PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } Result PgApiImpl::ListReplicationSlots() { @@ -2310,20 +2193,13 @@ Result PgApiImpl::UpdateAndPersistLSN( return pg_session_->pg_client().UpdateAndPersistLSN(stream_id, restart_lsn, confirmed_flush); } -Status PgApiImpl::NewDropReplicationSlot(const char *slot_name, - PgStatement **handle) { - auto stmt = std::make_unique(pg_session_, slot_name); - RETURN_NOT_OK(AddToCurrentPgMemctx(std::move(stmt), handle)); - return Status::OK(); +Status PgApiImpl::NewDropReplicationSlot(const char* slot_name, PgStatement** handle) { + return AddToCurrentPgMemctx( + std::make_unique(pg_session_, slot_name), handle); } -Status PgApiImpl::ExecDropReplicationSlot(PgStatement *handle) { - if (!PgStatement::IsValidStmt(handle, StmtOp::STMT_DROP_REPLICATION_SLOT)) { - // Invalid handle. - return STATUS(InvalidArgument, "Invalid statement handle"); - } - PgDropReplicationSlot *pg_stmt = down_cast(handle); - return pg_stmt->Exec(); +Status PgApiImpl::ExecDropReplicationSlot(PgStatement* handle) { + return VERIFY_RESULT_REF(GetStatementAs(handle)).Exec(); } Result PgApiImpl::YCQLStatementStats() { diff --git a/src/yb/yql/pggate/pggate.h b/src/yb/yql/pggate/pggate.h index e3159e21346..50b29b8bc9a 100644 --- a/src/yb/yql/pggate/pggate.h +++ b/src/yb/yql/pggate/pggate.h @@ -368,7 +368,7 @@ class PgApiImpl { Result GetColumnInfo(YBCPgTableDesc table_desc, int16_t attr_number); - Status DmlModifiesRow(PgStatement *handle, bool *modifies_row); + Result DmlModifiesRow(PgStatement* handle); Status SetIsSysCatalogVersionChange(PgStatement *handle); @@ -473,7 +473,7 @@ class PgApiImpl { Status DmlBindColumnCondIsNotNull(PgStatement *handle, int attr_num); Status DmlBindRow(YBCPgStatement handle, uint64_t ybctid, YBCBindColumn* columns, int count); - void DmlBindHashCode( + Status DmlBindHashCode( PgStatement* handle, const std::optional& start, const std::optional& end); Status DmlBindRange(YBCPgStatement handle, diff --git a/src/yb/yql/pggate/ybc_pggate.cc b/src/yb/yql/pggate/ybc_pggate.cc index 017ea8e0d9d..c19bae186cd 100644 --- a/src/yb/yql/pggate/ybc_pggate.cc +++ b/src/yb/yql/pggate/ybc_pggate.cc @@ -1112,7 +1112,7 @@ YBCStatus YBCPgSetDBCatalogCacheVersion( } YBCStatus YBCPgDmlModifiesRow(YBCPgStatement handle, bool *modifies_row) { - return ToYBCStatus(pgapi->DmlModifiesRow(handle, modifies_row)); + return ExtractValueFromResult(pgapi->DmlModifiesRow(handle), modifies_row); } YBCStatus YBCPgSetIsSysCatalogVersionChange(YBCPgStatement handle) { @@ -1346,8 +1346,7 @@ YBCStatus YBCPgDmlBindHashCodes( const auto start = MakeBound(start_type, start_value); const auto end = MakeBound(end_type, end_value); DCHECK(start || end); - pgapi->DmlBindHashCode(handle, start, end); - return YBCStatusOK(); + return ToYBCStatus(pgapi->DmlBindHashCode(handle, start, end)); } YBCStatus YBCPgDmlBindRange(YBCPgStatement handle,