129 lines
4.8 KiB
Python
129 lines
4.8 KiB
Python
from datetime import datetime, timezone
|
||
from flask import current_app
|
||
from databricks import sql
|
||
|
||
|
||
class DatabricksQuery:
|
||
def __init__(self):
|
||
"""纯Python连接Databricks SQL，无JVM依赖"""
|
||
try:
|
||
# 读取配置（参数和原代码一致，无需修改）
|
||
self.host = current_app.config["DATABRICKS_HOST"]
|
||
self.token = current_app.config["DATABRICKS_TOKEN"]
|
||
self.cluster_id = current_app.config["CLUSTER_ID"]
|
||
|
||
# 校验参数
|
||
if not all([self.host, self.token, self.cluster_id]):
|
||
raise ValueError(
|
||
"必须配置DATABRICKS_HOST、DATABRICKS_TOKEN、CLUSTER_ID"
|
||
)
|
||
|
||
print("Databricks SQL 连接器初始化成功")
|
||
except Exception as e:
|
||
print(f"初始化失败: {e}")
|
||
raise
|
||
|
||
# 复用原有参数校验逻辑
|
||
@staticmethod
|
||
def _validate_table_name(table_name):
|
||
if not table_name or not isinstance(table_name, str):
|
||
raise ValueError(f"无效的表名: {table_name}")
|
||
if ";" in table_name:
|
||
raise ValueError(f"表名包含非法字符 ';': {table_name}")
|
||
parts = table_name.split(".")
|
||
if len(parts) != 3:
|
||
raise ValueError(f"表名格式应为 catalog.schema.table: {table_name}")
|
||
|
||
@staticmethod
|
||
def _validate_vin(vin):
|
||
if vin is None:
|
||
return True
|
||
if not isinstance(vin, str):
|
||
raise ValueError(f"无效的 VIN: {vin} (必须是17位字母数字)")
|
||
return True
|
||
|
||
def query_table(
|
||
self, table_name, vin=None, start_time=None, end_time=None,limit=None
|
||
):
|
||
"""纯Python查询表数据，无JVM依赖"""
|
||
try:
|
||
# 验证参数
|
||
self._validate_table_name(table_name)
|
||
self._validate_vin(vin)
|
||
|
||
# 构建SQL查询语句（替代原PySpark DataFrame操作）
|
||
select_fields = """
|
||
Incident__Name AS incident_name,
|
||
Incident__description AS incident_description,
|
||
Meaning AS meaning,
|
||
Signal__Name As signal_name,
|
||
incident_time,
|
||
oneid,
|
||
speed,
|
||
int_value,
|
||
idc_tickcount_ms,
|
||
latitude,
|
||
longitude,
|
||
mux_data,
|
||
gps_heading,
|
||
float_value
|
||
"""
|
||
|
||
sql_query = f"SELECT {select_fields} FROM {table_name}"
|
||
where_conditions = []
|
||
|
||
# 拼接过滤条件
|
||
if vin:
|
||
where_conditions.append(f"oneid = '{vin}'")
|
||
if start_time:
|
||
start_dt = self._parse_iso_time(start_time).isoformat()
|
||
where_conditions.append(f"incident_time >= '{start_dt}'")
|
||
if end_time:
|
||
end_dt = self._parse_iso_time(end_time).isoformat()
|
||
where_conditions.append(f"incident_time <= '{end_dt}'")
|
||
|
||
# 拼接WHERE子句
|
||
if where_conditions:
|
||
sql_query += " WHERE " + " AND ".join(where_conditions)
|
||
# 拼接LIMIT
|
||
if limit is not None:
|
||
sql_query += f" LIMIT {limit}"
|
||
|
||
# sql_query="select * from hive_metastore.taf_level_two_plus.accident_report_data WHERE oneid = 'LE4LG4GB6SL256937' limit 1"
|
||
print(f"查询sql:{sql_query}")
|
||
|
||
# 执行SQL（纯Python，无JVM）
|
||
with sql.connect(
|
||
server_hostname=self.host,
|
||
http_path=f"sql/protocolv1/o/3035612432650494/{self.cluster_id}", # 注意路径格式
|
||
access_token=self.token,
|
||
) as connection:
|
||
with connection.cursor() as cursor:
|
||
cursor.execute(sql_query)
|
||
# 获取列名和数据，转换为字典格式（和原代码返回格式一致）
|
||
columns = [desc[0] for desc in cursor.description]
|
||
results = []
|
||
for row in cursor.fetchall():
|
||
results.append(dict(zip(columns, row)))
|
||
return results
|
||
|
||
except Exception as e:
|
||
print(f"查询失败: {str(e)}")
|
||
raise
|
||
|
||
# 复用原有时间解析逻辑
|
||
def _parse_iso_time(self, time_str):
|
||
try:
|
||
if time_str.endswith("Z"):
|
||
time_str = time_str[:-1] + "+00:00"
|
||
dt = datetime.fromisoformat(time_str)
|
||
if dt.tzinfo is None:
|
||
dt = dt.replace(tzinfo=timezone.utc)
|
||
return dt
|
||
except Exception as e:
|
||
raise ValueError(f"无效的时间格式 '{time_str}': {str(e)}")
|
||
|
||
def stop(self):
|
||
# 纯Python连接器无需关闭SparkSession，此处保留方法以兼容原有调用逻辑
|
||
print("Databricks SQL 连接已释放")
|