mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
feat(datasource): Hive data source supports adapting to SparkThriftServer metadata parsing #1391
This commit is contained in:
@@ -123,6 +123,7 @@ class DatasourceConf(BaseModel):
|
||||
lowVersion: bool = False
|
||||
ssl: bool = False
|
||||
poolSize: int = 5
|
||||
engine: str = 'hive'
|
||||
|
||||
def to_dict(self):
|
||||
return {
|
||||
@@ -140,7 +141,8 @@ class DatasourceConf(BaseModel):
|
||||
"timeout": self.timeout,
|
||||
"lowVersion": self.lowVersion,
|
||||
"ssl": self.ssl,
|
||||
"poolSize": self.poolSize
|
||||
"poolSize": self.poolSize,
|
||||
"engine": self.engine
|
||||
}
|
||||
|
||||
|
||||
|
||||
+84
-14
@@ -126,6 +126,88 @@ def get_origin_connect(type: str, conf: DatasourceConf):
|
||||
)
|
||||
|
||||
|
||||
def is_spark_thrift(conf: DatasourceConf) -> bool:
|
||||
"""Whether the hive datasource points to a Spark Thrift Server."""
|
||||
engine = getattr(conf, 'engine', None) or 'hive'
|
||||
return equals_ignore_case(engine, 'spark')
|
||||
|
||||
|
||||
def parse_hive_table_rows(res, is_spark: bool):
|
||||
"""Parse SHOW TABLES result rows.
|
||||
|
||||
- HiveServer2 returns a single `tab_name` column.
|
||||
- Spark Thrift Server (2.4+ through 4.x) returns
|
||||
`(database, tableName, isTemporary)`; the plain table name always sits in
|
||||
the second column. Over the legacy Hive-JDBC path a header row named
|
||||
`tableName` may be returned in-band, which we skip defensively.
|
||||
|
||||
Returns a list of (table_name, comment) tuples.
|
||||
"""
|
||||
tables = []
|
||||
seen = set()
|
||||
for item in res:
|
||||
if not item:
|
||||
continue
|
||||
row = [c.decode('utf-8') if isinstance(c, bytes) else c for c in item]
|
||||
row = [str(c) if c is not None else '' for c in row]
|
||||
|
||||
if is_spark:
|
||||
# Spark Thrift Server returns (database, tableName, isTemporary).
|
||||
if len(row) >= 2:
|
||||
# Skip a possible in-band header row over the legacy Hive-JDBC path.
|
||||
if row[1].lower() == 'tablename':
|
||||
continue
|
||||
name = row[1].strip()
|
||||
else:
|
||||
name = row[0].strip()
|
||||
else:
|
||||
# HiveServer2 returns a single tab_name column.
|
||||
name = row[0].strip()
|
||||
|
||||
if not name or name in seen:
|
||||
continue
|
||||
seen.add(name)
|
||||
tables.append((name, ''))
|
||||
return tables
|
||||
|
||||
|
||||
def parse_hive_field_rows(res, is_spark: bool):
|
||||
"""Parse DESCRIBE table result rows for HiveServer2 and Spark Thrift Server
|
||||
(all Spark 2.4+ versions use the same col_name/data_type/comment shape with
|
||||
a `# Partition Information` / `# col_name` section).
|
||||
|
||||
On the legacy Hive-JDBC path the (col_name, data_type, comment) header may
|
||||
also be returned in-band; we skip it defensively.
|
||||
|
||||
Returns a list of (field_name, field_type, field_comment) tuples.
|
||||
"""
|
||||
fields = []
|
||||
seen_fields = set()
|
||||
for item in res:
|
||||
field_name = item[0] if len(item) > 0 else None
|
||||
field_type = item[1] if len(item) > 1 else None
|
||||
field_comment = item[2] if len(item) > 2 else None
|
||||
|
||||
if not field_name or not str(field_name).strip():
|
||||
continue
|
||||
stripped_name = str(field_name).strip()
|
||||
|
||||
type_str = (field_type or '').strip() if isinstance(field_type, str) else (str(field_type).strip() if field_type else '')
|
||||
|
||||
# Legacy Hive-JDBC path may return the column header row as data.
|
||||
if is_spark and stripped_name.lower() == 'col_name':
|
||||
continue
|
||||
|
||||
# DESCRIBE includes section headings and repeats partition columns.
|
||||
if stripped_name.lstrip().startswith('#') and type_str in ('', 'data_type'):
|
||||
continue
|
||||
if stripped_name in seen_fields:
|
||||
continue
|
||||
seen_fields.add(stripped_name)
|
||||
fields.append((field_name, type_str, field_comment))
|
||||
return fields
|
||||
|
||||
|
||||
# use sqlalchemy
|
||||
def get_engine(ds: CoreDatasource, timeout: int = 0, use_pool: bool = False) -> Engine:
|
||||
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if not equals_ignore_case(ds.type,
|
||||
@@ -534,7 +616,7 @@ def get_tables(ds: CoreDatasource):
|
||||
with get_driver_connection(ds) as conn, conn.cursor() as cursor:
|
||||
cursor.execute(sql)
|
||||
res = cursor.fetchall()
|
||||
res_list = [TableSchema(*item) for item in res]
|
||||
res_list = [TableSchema(*item) for item in parse_hive_table_rows(res, is_spark_thrift(conf))]
|
||||
return res_list
|
||||
|
||||
|
||||
@@ -585,19 +667,7 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
|
||||
with get_driver_pool(ds).connection() as conn, conn.cursor() as cursor:
|
||||
cursor.execute(sql)
|
||||
res = cursor.fetchall()
|
||||
res_list = []
|
||||
seen_fields = set()
|
||||
for item in res:
|
||||
field_name = item[0]
|
||||
if not field_name or not field_name.strip():
|
||||
continue
|
||||
# DESCRIBE includes section headings and repeats partition columns.
|
||||
if field_name.lstrip().startswith('#') and (item[1] or '').strip() in ('', 'data_type'):
|
||||
continue
|
||||
if field_name in seen_fields:
|
||||
continue
|
||||
seen_fields.add(field_name)
|
||||
res_list.append(ColumnSchema(*item[:3]))
|
||||
res_list = [ColumnSchema(*item) for item in parse_hive_field_rows(res, is_spark_thrift(conf))]
|
||||
return res_list
|
||||
|
||||
|
||||
|
||||
@@ -41,6 +41,8 @@ def get_hive_fields():
|
||||
}
|
||||
_load_symbol("common/utils/utils.py", "equals_ignore_case", namespace)
|
||||
_load_symbol("apps/datasource/models/datasource.py", "ColumnSchema", namespace)
|
||||
_load_symbol("apps/db/db.py", "is_spark_thrift", namespace)
|
||||
_load_symbol("apps/db/db.py", "parse_hive_field_rows", namespace)
|
||||
get_fields = _load_symbol("apps/db/db.py", "get_fields", namespace)
|
||||
|
||||
def discover(rows):
|
||||
@@ -122,3 +124,106 @@ def test_extra_driver_columns_remain_supported(get_hive_fields):
|
||||
@pytest.mark.parametrize("rows", [[], [("", "", ""), ("# col_name", "data_type", "comment")]])
|
||||
def test_no_real_columns_returns_empty_list(get_hive_fields, rows):
|
||||
assert get_hive_fields(rows) == []
|
||||
|
||||
|
||||
# --- Spark Thrift Server support ---
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def get_hive_spark_fields():
|
||||
cursor = Mock()
|
||||
connection = SimpleNamespace(cursor=lambda: nullcontext(cursor))
|
||||
pool = SimpleNamespace(connection=lambda: nullcontext(connection))
|
||||
engine_probe = Mock(return_value=True)
|
||||
namespace = {
|
||||
"json": json,
|
||||
"CoreDatasource": SimpleNamespace,
|
||||
"DatasourceConf": SimpleNamespace,
|
||||
"aes_decrypt": lambda value: value,
|
||||
"DB": SimpleNamespace(
|
||||
get_db=lambda _: SimpleNamespace(connect_type="driver")
|
||||
),
|
||||
"ConnectType": SimpleNamespace(sqlalchemy="sqlalchemy"),
|
||||
"get_driver_pool": lambda _: pool,
|
||||
"get_field_sql": lambda *args: ("DESCRIBE sample", None, None),
|
||||
"is_spark_thrift": engine_probe,
|
||||
}
|
||||
_load_symbol("common/utils/utils.py", "equals_ignore_case", namespace)
|
||||
_load_symbol("apps/datasource/models/datasource.py", "ColumnSchema", namespace)
|
||||
_load_symbol("apps/db/db.py", "parse_hive_field_rows", namespace)
|
||||
get_fields = _load_symbol("apps/db/db.py", "get_fields", namespace)
|
||||
|
||||
def discover(rows):
|
||||
cursor.fetchall.return_value = rows
|
||||
fields = get_fields(SimpleNamespace(type="hive", configuration="{}"), "sample")
|
||||
return [(field.fieldName, field.fieldType, field.fieldComment) for field in fields]
|
||||
|
||||
return discover
|
||||
|
||||
|
||||
def test_spark_describe_skips_col_name_header_row(get_hive_spark_fields):
|
||||
# Spark Thrift Server returns the (col_name, data_type, comment) header as data.
|
||||
rows = [("col_name", "data_type", "comment"),
|
||||
("id", "bigint", None),
|
||||
("dt", "string", "partition date")]
|
||||
assert get_hive_spark_fields(rows) == [
|
||||
("id", "bigint", None),
|
||||
("dt", "string", "partition date"),
|
||||
]
|
||||
|
||||
|
||||
def test_spark_describe_with_partition_section(get_hive_spark_fields):
|
||||
rows = [
|
||||
("col_name", "data_type", "comment"),
|
||||
("id", "bigint", None),
|
||||
("# Partition Information", "", ""),
|
||||
("# col_name", "data_type", "comment"),
|
||||
("dt", "string", "partition date"),
|
||||
]
|
||||
assert get_hive_spark_fields(rows) == [
|
||||
("id", "bigint", None),
|
||||
("dt", "string", "partition date"),
|
||||
]
|
||||
|
||||
|
||||
# --- SHOW TABLES (parse_hive_table_rows) ---
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def get_hive_tables():
|
||||
namespace = {"DatasourceConf": SimpleNamespace}
|
||||
_load_symbol("apps/db/db.py", "parse_hive_table_rows", namespace)
|
||||
_load_symbol("common/utils/utils.py", "equals_ignore_case", namespace)
|
||||
_load_symbol("apps/db/db.py", "is_spark_thrift", namespace)
|
||||
return namespace["parse_hive_table_rows"]
|
||||
|
||||
|
||||
def test_hs2_show_tables_single_column(get_hive_tables):
|
||||
rows = [("orders",), ("users",), ("orders",)]
|
||||
assert get_hive_tables(rows, False) == [("orders", ""), ("users", "")]
|
||||
|
||||
|
||||
def test_spark_show_tables_three_columns_with_header(get_hive_tables):
|
||||
rows = [
|
||||
("default", "tableName", "isTemporary"), # header row returned as data
|
||||
("default", "orders", "false"),
|
||||
("default", "users", "false"),
|
||||
("default", "orders", "false"),
|
||||
]
|
||||
assert get_hive_tables(rows, True) == [("orders", ""), ("users", "")]
|
||||
|
||||
|
||||
def test_spark_show_tables_without_header(get_hive_tables):
|
||||
rows = [("default", "sales", "false"), ("default", "inventory", "false")]
|
||||
assert get_hive_tables(rows, True) == [("sales", ""), ("inventory", "")]
|
||||
|
||||
|
||||
def test_spark_show_tables_ignores_database_column_and_handles_temp_views(get_hive_tables):
|
||||
# Spark 2.4+/3.x/4.x always returns (database, tableName, isTemporary);
|
||||
# take the second column and tolerate an empty database (temp views).
|
||||
rows = [
|
||||
("salesdb", "orders", "false"),
|
||||
("", "temp_view", "true"),
|
||||
("salesdb", "orders", "false"),
|
||||
]
|
||||
assert get_hive_tables(rows, True) == [("orders", ""), ("temp_view", "")]
|
||||
|
||||
@@ -422,7 +422,10 @@
|
||||
"low_version": "Compatible with lower versions",
|
||||
"ssl": "Enable SSL",
|
||||
"file_path": "File Path",
|
||||
"pool_size": "Connection Pool Size"
|
||||
"pool_size": "Connection Pool Size",
|
||||
"engine": "Engine",
|
||||
"engine_hive": "HiveServer2",
|
||||
"engine_spark": "Spark Thrift Server"
|
||||
},
|
||||
"sync_fields": "Sync Fields",
|
||||
"sync_fields_success": "Sync fields successfully",
|
||||
|
||||
@@ -422,7 +422,10 @@
|
||||
"low_version": "낮은 버전 호환",
|
||||
"ssl": "SSL 활성화",
|
||||
"file_path": "파일 경로",
|
||||
"pool_size": "연결 풀 크기"
|
||||
"pool_size": "연결 풀 크기",
|
||||
"engine": "엔진",
|
||||
"engine_hive": "HiveServer2",
|
||||
"engine_spark": "Spark Thrift Server"
|
||||
},
|
||||
"sync_fields": "동기화된 테이블 구조",
|
||||
"sync_fields_success": "테이블 구조 동기화 성공",
|
||||
|
||||
@@ -422,7 +422,10 @@
|
||||
"low_version": "เข้ากันได้กับเวอร์ชันต่ำกว่า",
|
||||
"ssl": "เปิดใช้งาน SSL",
|
||||
"file_path": "เส้นทางไฟล์",
|
||||
"pool_size": "ขนาดพูลการเชื่อมต่อ"
|
||||
"pool_size": "ขนาดพูลการเชื่อมต่อ",
|
||||
"engine": "เอนจิน",
|
||||
"engine_hive": "HiveServer2",
|
||||
"engine_spark": "Spark Thrift Server"
|
||||
},
|
||||
"sync_fields": "ซิงค์โครงสร้างตาราง",
|
||||
"sync_fields_success": "ซิงค์โครงสร้างตารางสำเร็จ",
|
||||
|
||||
@@ -422,7 +422,10 @@
|
||||
"low_version": "兼容低版本",
|
||||
"ssl": "启用 SSL",
|
||||
"file_path": "文件路径",
|
||||
"pool_size": "连接池大小"
|
||||
"pool_size": "连接池大小",
|
||||
"engine": "引擎",
|
||||
"engine_hive": "HiveServer2",
|
||||
"engine_spark": "Spark Thrift Server"
|
||||
},
|
||||
"sync_fields": "同步表结构",
|
||||
"sync_fields_success": "同步表结构成功",
|
||||
|
||||
@@ -422,7 +422,10 @@
|
||||
"low_version": "相容低版本",
|
||||
"ssl": "啟用 SSL",
|
||||
"file_path": "文件路徑",
|
||||
"pool_size": "連線池大小"
|
||||
"pool_size": "連線池大小",
|
||||
"engine": "引擎",
|
||||
"engine_hive": "HiveServer2",
|
||||
"engine_spark": "Spark Thrift Server"
|
||||
},
|
||||
"sync_fields": "同步表結構",
|
||||
"sync_fields_success": "同步表結構成功",
|
||||
|
||||
@@ -129,6 +129,7 @@ const form = ref<any>({
|
||||
lowVersion: false,
|
||||
ssl: false,
|
||||
poolSize: 5,
|
||||
engine: 'hive',
|
||||
})
|
||||
|
||||
const close = () => {
|
||||
@@ -182,6 +183,7 @@ const initForm = (item: any, editTable: boolean = false) => {
|
||||
configuration.poolSize !== null && configuration.poolSize !== undefined
|
||||
? configuration.poolSize
|
||||
: 5
|
||||
form.value.engine = configuration.engine ? configuration.engine : 'hive'
|
||||
}
|
||||
|
||||
if (editTable) {
|
||||
@@ -256,6 +258,7 @@ const initForm = (item: any, editTable: boolean = false) => {
|
||||
lowVersion: false,
|
||||
ssl: false,
|
||||
poolSize: 5,
|
||||
engine: 'hive',
|
||||
}
|
||||
}
|
||||
dialogVisible.value = true
|
||||
@@ -346,6 +349,7 @@ const buildConf = () => {
|
||||
lowVersion: form.value.lowVersion,
|
||||
ssl: form.value.ssl,
|
||||
poolSize: form.value.poolSize,
|
||||
engine: form.value.engine,
|
||||
})
|
||||
)
|
||||
const obj = JSON.parse(JSON.stringify(form.value))
|
||||
@@ -364,6 +368,7 @@ const buildConf = () => {
|
||||
delete obj.lowVersion
|
||||
delete obj.ssl
|
||||
delete obj.poolSize
|
||||
delete obj.engine
|
||||
return obj
|
||||
}
|
||||
|
||||
@@ -738,6 +743,12 @@ defineExpose({
|
||||
<el-radio value="sid">{{ t('ds.form.mode.sid') }}</el-radio>
|
||||
</el-radio-group>
|
||||
</el-form-item>
|
||||
<el-form-item v-if="form.type === 'hive'" :label="t('ds.form.engine')" prop="engine">
|
||||
<el-radio-group v-model="form.engine">
|
||||
<el-radio value="hive">{{ t('ds.form.engine_hive') }}</el-radio>
|
||||
<el-radio value="spark">{{ t('ds.form.engine_spark') }}</el-radio>
|
||||
</el-radio-group>
|
||||
</el-form-item>
|
||||
<el-form-item
|
||||
v-if="form.type === 'sqlServer'"
|
||||
:label="t('ds.form.low_version')"
|
||||
|
||||
Reference in New Issue
Block a user