fix(tdengine): fix table data preview

crates/dbx-core/src/sql_dialect/table_select.rs

@@ -33,7 +33,11 @@ pub fn build_table_data_select_sql(options: TableDataSelectSqlOptions) -> String
     let select_columns = if options.include_row_id && database_type == Some(DatabaseType::Oracle) {
         format!("ROWIDTOCHAR(t.ROWID) AS \"{DBX_ROWID_COLUMN}\", t.*")
     } else {
-        build_select_columns(database_type, &options.columns)
+        build_select_columns(
+            database_type,
+            &options.columns,
+            tdengine_should_include_tbname(database_type, options.table_type.as_deref()),
+        )
     };
     let rownum_select_columns = quoted_table_columns_or_star(database_type, &options.columns);
     let page_select_columns = if options.include_row_id && database_type == Some(DatabaseType::Oracle) {
@@ -209,16 +213,40 @@ pub(super) fn is_tdengine_tbname(database_type: Option<DatabaseType>, name: &str
     database_type == Some(DatabaseType::Tdengine) && name.eq_ignore_ascii_case(DBX_TDENGINE_TBNAME_COLUMN)
 }
 
-pub(super) fn build_select_columns(database_type: Option<DatabaseType>, columns: &[String]) -> String {
+fn tdengine_should_include_tbname(database_type: Option<DatabaseType>, table_type: Option<&str>) -> bool {
+    if database_type != Some(DatabaseType::Tdengine) {
+        return false;
+    }
+    matches!(
+        table_type.map(|value| value.trim().to_ascii_uppercase()),
+        Some(value) if value == "STABLE" || value == "SUPER TABLE" || value == "SUPERTABLE"
+    )
+}
+
+pub(super) fn build_select_columns(
+    database_type: Option<DatabaseType>,
+    columns: &[String],
+    include_tdengine_tbname: bool,
+) -> String {
     if columns.is_empty() {
         return "*".to_string();
     }
     if database_type == Some(DatabaseType::Tdengine) {
         let mut tdengine_columns = Vec::new();
-        if !columns.iter().any(|column| column.eq_ignore_ascii_case(DBX_TDENGINE_TBNAME_COLUMN)) {
+        if include_tdengine_tbname
+            && !columns.iter().any(|column| column.eq_ignore_ascii_case(DBX_TDENGINE_TBNAME_COLUMN))
+        {
             tdengine_columns.push(DBX_TDENGINE_TBNAME_COLUMN.to_string());
         }
-        tdengine_columns.extend(columns.iter().cloned());
+        tdengine_columns.extend(
+            columns
+                .iter()
+                .filter(|column| include_tdengine_tbname || !column.eq_ignore_ascii_case(DBX_TDENGINE_TBNAME_COLUMN))
+                .cloned(),
+        );
+        if tdengine_columns.is_empty() {
+            return "*".to_string();
+        }
         return tdengine_columns
             .iter()
             .map(|column| {

crates/dbx-core/src/sql_dialect/tests.rs

@@ -224,6 +224,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::Mysql),
             schema: None,
             table_name: "users".to_string(),
+            table_type: None,
             primary_keys: vec!["id".to_string()],
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -240,6 +241,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::Goldendb),
             schema: None,
             table_name: "sys_dic".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -256,6 +258,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::Postgres),
             schema: Some("public".to_string()),
             table_name: "orders".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -272,6 +275,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::Xugu),
             schema: Some("DBX_TEST".to_string()),
             table_name: "PRODUCTS".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -288,6 +292,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::StarRocks),
             schema: None,
             table_name: "sales_report".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: vec!["customer_name".to_string(), "amount".to_string()],
             fallback_order_columns: Vec::new(),
@@ -304,6 +309,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::Db2),
             schema: Some("DB2INST1".to_string()),
             table_name: "ORDERS".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -336,6 +342,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::OceanbaseOracle),
             schema: Some("DBXTEST".to_string()),
             table_name: "ORDERS".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -352,6 +359,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::OceanbaseOracle),
             schema: Some("DBXTEST".to_string()),
             table_name: "ORDERS".to_string(),
+            table_type: None,
             primary_keys: vec!["ID".to_string()],
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -368,6 +376,7 @@ fn builds_table_data_where_and_schema_queries() {
                 database_type: Some(DatabaseType::Db2),
                 schema: Some("DB2INST1".to_string()),
                 table_name: "ORDERS".to_string(),
+                table_type: None,
                 primary_keys: vec!["ID".to_string()],
                 columns: vec!["ID".to_string(), "AMOUNT".to_string()],
                 fallback_order_columns: Vec::new(),
@@ -384,6 +393,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::Iris),
             schema: Some("Ens".to_string()),
             table_name: "AlarmResponse".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -400,6 +410,7 @@ fn builds_table_data_where_and_schema_queries() {
             database_type: Some(DatabaseType::Iotdb),
             schema: Some("root.test".to_string()),
             table_name: "device2".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -420,6 +431,7 @@ fn builds_informix_table_data_with_skip_first_pagination() {
             database_type: Some(DatabaseType::Informix),
             schema: Some("ignored".to_string()),
             table_name: "users".to_string(),
+            table_type: None,
             primary_keys: vec!["id".to_string()],
             columns: vec!["id".to_string(), "name".to_string()],
             fallback_order_columns: Vec::new(),
@@ -452,6 +464,7 @@ fn explicit_table_data_order_is_preserved() {
             database_type: Some(DatabaseType::Postgres),
             schema: Some("public".to_string()),
             table_name: "country_gdp".to_string(),
+            table_type: None,
             primary_keys: vec!["year".to_string()],
             columns: vec!["iso3".to_string(), "year".to_string(), "gdp_pc".to_string()],
             fallback_order_columns: Vec::new(),
@@ -471,6 +484,7 @@ fn builds_iris_table_data_sql_with_literal_top_and_quoted_object() {
         database_type: Some(DatabaseType::Iris),
         schema: Some("Ens".to_string()),
         table_name: "AlarmResponse".to_string(),
+        table_type: None,
         primary_keys: vec!["ID".to_string()],
         columns: vec!["ID".to_string(), "Status".to_string()],
         fallback_order_columns: Vec::new(),
@@ -493,32 +507,68 @@ fn builds_iris_table_data_sql_with_literal_top_and_quoted_object() {
 #[test]
 fn builds_table_data_special_column_queries() {
     assert_eq!(
-            build_table_data_select_sql(TableDataSelectSqlOptions {
-                database_type: Some(DatabaseType::Tdengine),
-                schema: Some("test_db".to_string()),
-                table_name: "meters".to_string(),
-                primary_keys: vec!["ts".to_string()],
-                columns: vec![
-                    "ts".to_string(),
-                    "current".to_string(),
-                    "voltage".to_string(),
-                    "location".to_string(),
-                    "groupid".to_string(),
-                ],
-                fallback_order_columns: Vec::new(),
-                order_by: None,
-                limit: Some(100),
-                offset: None,
-                where_input: None,
-                include_row_id: false,
-            }),
-            "SELECT tbname, `ts` AS `ts`, `current` AS `current`, `voltage` AS `voltage`, `location` AS `location`, `groupid` AS `groupid` FROM `test_db`.`meters` LIMIT 100;"
-        );
+        build_table_data_select_sql(TableDataSelectSqlOptions {
+            database_type: Some(DatabaseType::Tdengine),
+            schema: Some("test_db".to_string()),
+            table_name: "meters".to_string(),
+            table_type: Some("STABLE".to_string()),
+            primary_keys: vec!["ts".to_string()],
+            columns: vec![
+                "ts".to_string(),
+                "current".to_string(),
+                "voltage".to_string(),
+                "location".to_string(),
+                "groupid".to_string(),
+            ],
+            fallback_order_columns: Vec::new(),
+            order_by: None,
+            limit: Some(100),
+            offset: None,
+            where_input: None,
+            include_row_id: false,
+        }),
+        "SELECT tbname, `ts` AS `ts`, `current` AS `current`, `voltage` AS `voltage`, `location` AS `location`, `groupid` AS `groupid` FROM `test_db`.`meters` LIMIT 100;"
+    );
+    assert_eq!(
+        build_table_data_select_sql(TableDataSelectSqlOptions {
+            database_type: Some(DatabaseType::Tdengine),
+            schema: Some("test_db".to_string()),
+            table_name: "d1001".to_string(),
+            table_type: Some("TABLE".to_string()),
+            primary_keys: vec!["ts".to_string()],
+            columns: vec!["ts".to_string(), "current".to_string()],
+            fallback_order_columns: Vec::new(),
+            order_by: None,
+            limit: Some(100),
+            offset: None,
+            where_input: None,
+            include_row_id: false,
+        }),
+        "SELECT `ts` AS `ts`, `current` AS `current` FROM `test_db`.`d1001` LIMIT 100;"
+    );
+    assert_eq!(
+        build_table_data_select_sql(TableDataSelectSqlOptions {
+            database_type: Some(DatabaseType::Tdengine),
+            schema: Some("test_db".to_string()),
+            table_name: "d1001".to_string(),
+            table_type: Some("TABLE".to_string()),
+            primary_keys: vec!["ts".to_string()],
+            columns: vec!["tbname".to_string(), "ts".to_string(), "current".to_string()],
+            fallback_order_columns: Vec::new(),
+            order_by: None,
+            limit: Some(100),
+            offset: None,
+            where_input: None,
+            include_row_id: false,
+        }),
+        "SELECT `ts` AS `ts`, `current` AS `current` FROM `test_db`.`d1001` LIMIT 100;"
+    );
     assert_eq!(
         build_table_data_select_sql(TableDataSelectSqlOptions {
             database_type: Some(DatabaseType::Hive),
             schema: None,
             table_name: "departments".to_string(),
+            table_type: None,
             primary_keys: Vec::new(),
             columns: vec!["id".to_string(), "name".to_string()],
             fallback_order_columns: Vec::new(),
@@ -539,6 +589,7 @@ fn builds_sqlserver_table_data_pages() {
             database_type: Some(DatabaseType::SqlServer),
             schema: Some("dbo".to_string()),
             table_name: "accounts".to_string(),
+            table_type: None,
             primary_keys: vec!["id".to_string()],
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -555,6 +606,7 @@ fn builds_sqlserver_table_data_pages() {
                 database_type: Some(DatabaseType::SqlServer),
                 schema: Some("sales".to_string()),
                 table_name: "orders".to_string(),
+                table_type: None,
                 primary_keys: vec!["order_id".to_string()],
                 columns: vec!["order_id".to_string(), "customer".to_string()],
                 fallback_order_columns: Vec::new(),
@@ -575,6 +627,7 @@ fn builds_oracle_and_neo4j_table_data_queries() {
             database_type: Some(DatabaseType::Oracle),
             schema: Some("DBXTEST".to_string()),
             table_name: "DBX_LOAD_TABLE_006".to_string(),
+            table_type: None,
             primary_keys: vec![DBX_ROWID_COLUMN.to_string()],
             columns: Vec::new(),
             fallback_order_columns: Vec::new(),
@@ -607,6 +660,7 @@ fn builds_oracle_and_neo4j_table_data_queries() {
                 database_type: Some(DatabaseType::Neo4j),
                 schema: None,
                 table_name: "Employee".to_string(),
+                table_type: None,
                 primary_keys: vec!["id".to_string()],
                 columns: vec!["id".to_string(), "first name".to_string(), "role".to_string()],
                 fallback_order_columns: Vec::new(),

crates/dbx-core/src/sql_dialect/types.rs

@@ -24,6 +24,8 @@ pub struct TableDataSelectSqlOptions {
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub schema: Option<String>,
     pub table_name: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub table_type: Option<String>,
     #[serde(default)]
     pub primary_keys: Vec<String>,
     #[serde(default)]