From 888fb291a0dc6e7094d26a8c2c77422ab88e41d8 Mon Sep 17 00:00:00 2001
From: shenzheng4 <shenzheng4@huawei.com>
Date: Wed, 24 May 2023 20:43:13 +0800
Subject: [PATCH] check max_size when create database in tablespace, and use
 unlimited tablespace when segment=on

---
 src/common/backend/utils/adt/dbsize.cpp           |  5 +++++
 src/gausskernel/optimizer/commands/dbcommands.cpp | 12 ++++++++++++
 src/gausskernel/optimizer/commands/tablecmds.cpp  | 12 ++++++++++++
 src/include/commands/dbcommands.h                 |  1 +
 src/include/commands/tablespace.h                 |  4 ++--
 5 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/common/backend/utils/adt/dbsize.cpp b/src/common/backend/utils/adt/dbsize.cpp
index 48defd20e..b3c3f7e41 100644
--- a/src/common/backend/utils/adt/dbsize.cpp
+++ b/src/common/backend/utils/adt/dbsize.cpp
@@ -253,6 +253,11 @@ static int64 calculate_database_size(Oid dbOid)
     return totalsize;
 }
 
+int64 pg_cal_database_size_oid(Oid dbOid)
+{
+    return (int64)(calculate_database_size(dbOid));
+}
+
 /*
  * @Description: calculate the compress ratio of the column-partition relation.
  * @in onerel: the relation.
diff --git a/src/gausskernel/optimizer/commands/dbcommands.cpp b/src/gausskernel/optimizer/commands/dbcommands.cpp
index 348327c66..df588ce4d 100644
--- a/src/gausskernel/optimizer/commands/dbcommands.cpp
+++ b/src/gausskernel/optimizer/commands/dbcommands.cpp
@@ -50,6 +50,7 @@
 #include "commands/comment.h"
 #include "commands/dbcommands.h"
 #include "commands/tablecmds.h"
+#include "commands/tablespace.h"
 #include "replication/slot.h"
 #include "commands/sec_rls_cmds.h"
 #include "commands/seclabel.h"
@@ -437,6 +438,17 @@ Oid createdb(const CreatedbStmt* stmt)
             ereport(ERROR,
                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE), errmsg("pg_global cannot be used as default tablespace")));
 
+        Oid defDbtemplate = get_database_oid(dbtemplate, false);
+        uint64 currSize = (uint64)pg_cal_database_size_oid(defDbtemplate);
+        uint64 tablespaceMaxSize = 0;
+        bool isLimit = TableSpaceUsageManager::IsLimited(dst_deftablespace, &tablespaceMaxSize);
+        if (isLimit && (tablespaceMaxSize < currSize)) {
+            ereport(ERROR,
+                (errcode(ERRCODE_INSUFFICIENT_RESOURCES), 
+                    errmsg("default database size %lu is bigger than max_size %lu set in tablespace \"%s\".", 
+                            currSize, tablespaceMaxSize, tablespacename)));
+        }
+
         /*
          * If we are trying to change the default tablespace of the template,
          * we require that the template not have any files in the new default
diff --git a/src/gausskernel/optimizer/commands/tablecmds.cpp b/src/gausskernel/optimizer/commands/tablecmds.cpp
index 46f6505f1..eecc141a7 100755
--- a/src/gausskernel/optimizer/commands/tablecmds.cpp
+++ b/src/gausskernel/optimizer/commands/tablecmds.cpp
@@ -2902,6 +2902,18 @@ ObjectAddress DefineRelation(CreateStmt* stmt, char relkind, Oid ownerId, Object
                 errmsg("The table %s do not support segment storage", stmt->relation->relname)));
     }
 
+    if (storage_type == SEGMENT_PAGE) {
+        Oid tbspcId = (tablespaceId == InvalidOid) ? u_sess->proc_cxt.MyDatabaseTableSpace : tablespaceId;
+        uint64 tablespaceMaxSize = 0;
+        bool isLimit = TableSpaceUsageManager::IsLimited(tbspcId, &tablespaceMaxSize);
+        if (isLimit) {
+            ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), errmodule(MOD_SEGMENT_PAGE),
+                errmsg("The table %s do not support segment-page storage", stmt->relation->relname),
+                errdetail("Segment-page storage doest not support limited tablespace \"%s\"", get_tablespace_name(tbspcId)),
+                errhint("use default or unlimited user defined tablespace before using segment-page storage.")));
+        }
+    }
+
     if (ENABLE_DMS && !u_sess->attr.attr_common.IsInplaceUpgrade) {
         if ((relkind == RELKIND_RELATION && storage_type != SEGMENT_PAGE) ||
             relkind == RELKIND_MATVIEW ||
diff --git a/src/include/commands/dbcommands.h b/src/include/commands/dbcommands.h
index ea1a70986..1a3fcf021 100644
--- a/src/include/commands/dbcommands.h
+++ b/src/include/commands/dbcommands.h
@@ -59,6 +59,7 @@ extern ObjectAddress AlterDatabaseOwner(const char* dbname, Oid newOwnerId);
 extern void AlterDatabasePermissionCheck(Oid dboid, const char* dbname);
 extern Oid get_database_oid(const char* dbname, bool missingok);
 extern char* get_database_name(Oid dbid);
+extern int64 pg_cal_database_size_oid(Oid dbOid);
 extern char* get_and_check_db_name(Oid dbid, bool is_ereport = false);
 extern bool have_createdb_privilege(void);
 
diff --git a/src/include/commands/tablespace.h b/src/include/commands/tablespace.h
index fbcf2c328..96213e579 100644
--- a/src/include/commands/tablespace.h
+++ b/src/include/commands/tablespace.h
@@ -55,9 +55,9 @@ public:
     static int ShmemSize(void);
     static void Init(void);
     static void IsExceedMaxsize(Oid tableSpaceOid, uint64 requestSize, bool segment);
-
-private:
     static bool IsLimited(Oid tableSpaceOid, uint64* maxSize);
+    
+private:
     static inline int GetBucketIndex(Oid tableSpaceOid);
     static inline void ResetUsageSlot(TableSpaceUsageSlot* info);
     static inline void ResetBucket(TableSpaceUsageBucket* bucket);