from impala.dbapi import connect
import requests
import json
import os

# 配置参数
IMPALA_HOST = os.getenv('IMPALA_HOST', '192.168.0.1')
IMPALA_PORT = int(os.getenv('IMPALA_PORT', 21050))
IMPALA_USER = os.getenv('IMPALA_USER', 'hive')
IMPALA_PWD = os.getenv('IMPALA_PWD', '11111111')
WEBHOOK_URL = os.getenv('WEBHOOK_URL', 'https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=替换为自己的key')

def query_impala():
    try:
        conn = connect(
            host=IMPALA_HOST,
            port=IMPALA_PORT,
            auth_mechanism='LDAP',
            user=IMPALA_USER,
            password=IMPALA_PWD,
            timeout=300
        )
        cursor = conn.cursor()
        
        # 检查数据库和表权限
        cursor.execute("SHOW DATABASES LIKE 'test'")
        if not cursor.fetchall():
            raise Exception("数据库 'test' 不存在或无访问权限")
        
        cursor.execute("SHOW TABLES IN test LIKE 'yarn_apps'")
        if not cursor.fetchall():
            raise Exception("表 'test.yarn_apps' 不存在或无访问权限")
        
        sql = """
        SELECT id, name,
               from_unixtime(CAST(startedTime/1000 AS BIGINT)) AS start_time,
               from_unixtime(CAST(finishedTime/1000 AS BIGINT)) AS finish_time,
               ROUND((finishedTime - startedTime)/60000.0, 2) AS duration_minutes
        FROM test.yarn_apps
        WHERE 
          (finishedTime - startedTime) > 30 * 60 * 1000
          AND to_date(from_unixtime(CAST(startedTime/1000 AS BIGINT))) = to_date(current_timestamp())
          AND date_part('hour', from_unixtime(CAST(startedTime/1000 AS BIGINT))) BETWEEN 0 AND 8
          AND finishedTime IS NOT NULL
          AND startedTime IS NOT NULL
          AND finishedTime > startedTime
        ORDER BY duration_minutes DESC;
        """
        cursor.execute(sql)
        columns = [desc[0] for desc in cursor.description]
        results = [list(row) for row in cursor.fetchall()]
        
        return columns, results
    except Exception as e:
        raise Exception("Impala query failed: " + str(e))
    finally:
        if 'cursor' in locals() and cursor is not None:
            cursor.close()
        if 'conn' in locals() and conn is not None:
            conn.close()

# 先使用文本消息测试,确保Webhook能正常工作
def send_webhook_alert(data):
    headers = {'Content-Type': 'application/json'}
    
    # 构建简单文本消息(避免Markdown格式问题)
    text_lines = ["⚠️ YARN长任务告警:"]
    text_lines.append(f"检测到 {len(data['rows'])} 个超时任务(运行时间超过30分钟)")
    text_lines.append("---")
    
    for i, row in enumerate(data['rows'], 1):
        text_lines.append(f"任务{i}:")
        text_lines.append(f"  ID: {row[0]}")
        text_lines.append(f"  名称: {row[1]}")
        text_lines.append(f"  开始时间: {row[2]}")
        text_lines.append(f"  结束时间: {row[3]}")
        text_lines.append(f"  持续时间(分钟): {row[4]}")
        text_lines.append("---")
    
    # 合并为单行文本(企业微信文本消息支持换行符)
    alert_text = "\n".join(text_lines)
    
    # 使用文本消息类型而非markdown
    payload = {
        "msgtype": "text",
        "text": {
            "content": alert_text
        }
    }
    
    try:
        print(f"发送的Webhook内容: {json.dumps(payload, ensure_ascii=False)}")
        response = requests.post(
            WEBHOOK_URL,
            headers=headers,
            data=json.dumps(payload, ensure_ascii=False).encode('utf-8'),
            timeout=10
        )
        
        response_json = response.json()
        print(f"Webhook响应: {response_json}")
        
        if response.status_code != 200 or response_json.get('errcode') != 0:
            raise Exception(
                f"Webhook发送失败: HTTP {response.status_code}, "
                f"错误码: {response_json.get('errcode')}, "
                f"错误信息: {response_json.get('errmsg')}"
            )
    except json.JSONDecodeError:
        raise Exception(f"Webhook响应格式错误,内容: {response.text}")
    except Exception as e:
        raise Exception(f"Webhook发送错误: {str(e)}")

if __name__ == "__main__":
    try:
        columns, rows = query_impala()
        if rows:
            send_webhook_alert({'columns': columns, 'rows': rows})
            print(f"✅ 已成功发送 {len(rows)} 条告警")
        else:
            print("🟢 无超时任务,无需告警")
    except Exception as e:
        print(f"❌ 执行失败: {str(e)}")