From f857d7fbb423b023b39884bee9b2c895796b47e6 Mon Sep 17 00:00:00 2001
From: aykutbozkurt <aykut.bozkurt1995@gmail.com>
Date: Wed, 7 Jun 2023 11:24:08 +0300
Subject: [PATCH] support single shard tables for query generator

---
 .../regress/citus_tests/query_generator/README.md  |  2 +-
 .../citus_tests/query_generator/config/config.py   |  8 ++++++--
 .../citus_tests/query_generator/config/config.yaml | 14 +++++++++++++-
 .../citus_tests/query_generator/data_gen.py        |  5 ++++-
 .../regress/citus_tests/query_generator/ddl_gen.py | 12 ++++++++++--
 .../citus_tests/query_generator/node_defs.py       |  7 ++++---
 6 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/src/test/regress/citus_tests/query_generator/README.md b/src/test/regress/citus_tests/query_generator/README.md
index b35a96c0a..3a52f245a 100644
--- a/src/test/regress/citus_tests/query_generator/README.md
+++ b/src/test/regress/citus_tests/query_generator/README.md
@@ -172,7 +172,7 @@ Tool supports following citus table types:
 targetTables:
   - Table:
     ...
-    citusType: <one of (DISTRIBUTED || REFERENCE || POSTGRES)>
+    citusType: <one of (HASH_DISTRIBUTED || SINGLE_SHARD_DISTRIBUTED || REFERENCE || POSTGRES)>
     ...
 ```
 
diff --git a/src/test/regress/citus_tests/query_generator/config/config.py b/src/test/regress/citus_tests/query_generator/config/config.py
index fec93ef19..0f8b3e010 100644
--- a/src/test/regress/citus_tests/query_generator/config/config.py
+++ b/src/test/regress/citus_tests/query_generator/config/config.py
@@ -111,8 +111,12 @@ def getMaxAllowedCountForTable(tableName):
     return filtered[0].maxAllowedUseOnQuery
 
 
-def isTableDistributed(table):
-    return table.citusType == CitusType.DISTRIBUTED
+def isTableHashDistributed(table):
+    return table.citusType == CitusType.HASH_DISTRIBUTED
+
+
+def isTableSingleShardDistributed(table):
+    return table.citusType == CitusType.SINGLE_SHARD_DISTRIBUTED
 
 
 def isTableReference(table):
diff --git a/src/test/regress/citus_tests/query_generator/config/config.yaml b/src/test/regress/citus_tests/query_generator/config/config.yaml
index 1920966ee..10f5b4804 100644
--- a/src/test/regress/citus_tests/query_generator/config/config.yaml
+++ b/src/test/regress/citus_tests/query_generator/config/config.yaml
@@ -26,7 +26,7 @@ commonColName: id
 targetTables:
   - Table:
       name: dist
-      citusType: DISTRIBUTED
+      citusType: HASH_DISTRIBUTED
       maxAllowedUseOnQuery: 10
       rowCount: 10
       nullRate: 0.1
@@ -36,6 +36,18 @@ targetTables:
           name: id
           type: int
       distinctCopyCount: 2
+  # - Table:
+  #     name: single_dist
+  #     citusType: SINGLE_SHARD_DISTRIBUTED
+  #     maxAllowedUseOnQuery: 10
+  #     rowCount: 10
+  #     nullRate: 0.1
+  #     duplicateRate: 0.1
+  #     columns:
+  #       - Column:
+  #         name: id
+  #         type: int
+  #     distinctCopyCount: 2
   - Table:
       name: ref
       citusType: REFERENCE
diff --git a/src/test/regress/citus_tests/query_generator/data_gen.py b/src/test/regress/citus_tests/query_generator/data_gen.py
index 96f7a1366..a6ff5f232 100644
--- a/src/test/regress/citus_tests/query_generator/data_gen.py
+++ b/src/test/regress/citus_tests/query_generator/data_gen.py
@@ -17,7 +17,10 @@ def getTableData():
         dataGenerationSql += "\n"
 
         # generate null rows
-        if not table.citusType == CitusType.DISTRIBUTED:
+        if table.citusType not in (
+            CitusType.HASH_DISTRIBUTED,
+            CitusType.SINGLE_SHARD_DISTRIBUTED,
+        ):
             targetNullRows = int(table.rowCount * table.nullRate)
             dataGenerationSql += _genNullData(table.name, targetNullRows)
             dataGenerationSql += "\n"
diff --git a/src/test/regress/citus_tests/query_generator/ddl_gen.py b/src/test/regress/citus_tests/query_generator/ddl_gen.py
index b2f97f694..2e7d2e999 100755
--- a/src/test/regress/citus_tests/query_generator/ddl_gen.py
+++ b/src/test/regress/citus_tests/query_generator/ddl_gen.py
@@ -1,4 +1,9 @@
-from config.config import getConfig, isTableDistributed, isTableReference
+from config.config import (
+    getConfig,
+    isTableHashDistributed,
+    isTableReference,
+    isTableSingleShardDistributed,
+)
 
 
 def getTableDDLs():
@@ -23,7 +28,7 @@ def _genTableDDL(table):
         ddl += _genColumnDDL(table.columns[-1])
     ddl += ");\n"
 
-    if isTableDistributed(table):
+    if isTableHashDistributed(table):
         ddl += (
             "SELECT create_distributed_table("
             + "'"
@@ -34,6 +39,9 @@ def _genTableDDL(table):
             + ");"
         )
         ddl += "\n"
+    if isTableSingleShardDistributed(table):
+        ddl += "SELECT create_distributed_table(" + "'" + table.name + "'" ",NULL);"
+        ddl += "\n"
     elif isTableReference(table):
         ddl += "SELECT create_reference_table(" + "'" + table.name + "'" + ");"
         ddl += "\n"
diff --git a/src/test/regress/citus_tests/query_generator/node_defs.py b/src/test/regress/citus_tests/query_generator/node_defs.py
index b0db1da63..62135d177 100755
--- a/src/test/regress/citus_tests/query_generator/node_defs.py
+++ b/src/test/regress/citus_tests/query_generator/node_defs.py
@@ -22,9 +22,10 @@ class RestrictOp(Enum):
 
 
 class CitusType(Enum):
-    DISTRIBUTED = 1
-    REFERENCE = 2
-    POSTGRES = 3
+    HASH_DISTRIBUTED = 1
+    SINGLE_SHARD_DISTRIBUTED = 2
+    REFERENCE = 3
+    POSTGRES = 4
 
 
 class Table: