From 84bedee3066681e1eb3958780d1b5d1a02634d40 Mon Sep 17 00:00:00 2001 From: "Kevin.Shin" Date: Mon, 18 Mar 2024 16:58:19 +0800 Subject: [PATCH] fix(datasource): Fix the bug of doris missing comments (#1308) Co-authored-by: shenk-b --- dbgpt/datasource/rdbms/conn_doris.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/dbgpt/datasource/rdbms/conn_doris.py b/dbgpt/datasource/rdbms/conn_doris.py index ddf28d397..47aaea2a5 100644 --- a/dbgpt/datasource/rdbms/conn_doris.py +++ b/dbgpt/datasource/rdbms/conn_doris.py @@ -1,4 +1,4 @@ -from typing import Any, Iterable, List, Optional, Tuple +from typing import Any, Dict, Iterable, List, Optional, Tuple from urllib.parse import quote from urllib.parse import quote_plus as urlquote @@ -68,6 +68,26 @@ class DorisConnect(RDBMSDatabase): """Get user info.""" return [] + def get_columns(self, table_name: str) -> List[Dict]: + """Get columns. + Args: + table_name (str): str + Returns: + columns: List[Dict], which contains name: str, type: str, default_expression: str, is_in_primary_key: bool, comment: str + eg:[{'name': 'id', 'type': 'UInt64', 'default_expression': '', 'is_in_primary_key': True, 'comment': 'id'}, ...] + """ + fields = self.get_fields(table_name) + return [ + { + "name": field[0], + "type": field[1], + "default": field[2], + "nullable": field[3], + "comment": field[4], + } + for field in fields + ] + def get_fields(self, table_name) -> List[Tuple]: """Get column fields about specified table.""" cursor = self.get_session().execute(