add workflow 泰克产品相关性分析,dev
This commit is contained in:
root
parent
6f30e692a5
commit
a13f79ce29
|
|
@ -0,0 +1,239 @@
|
|||
import requests
|
||||
import json
|
||||
import psycopg2
|
||||
from psycopg2.extras import RealDictCursor
|
||||
import time
|
||||
from typing import List, Dict, Any
|
||||
import logging
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class APIClient:
|
||||
"""简化的API客户端类"""
|
||||
|
||||
def __init__(self, api_url: str, auth_token: str):
|
||||
self.api_url = api_url
|
||||
self.headers = {
|
||||
'Authorization': auth_token,
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
def call_api_with_inputs(self, inputs: Dict[str, Any] = None, query: str = None) -> Dict[str, Any]:
|
||||
"""直接调用带inputs参数的API
|
||||
|
||||
Args:
|
||||
inputs: inputs参数字典
|
||||
|
||||
Returns:
|
||||
API响应数据
|
||||
"""
|
||||
payload = json.dumps({
|
||||
"inputs": inputs,
|
||||
"query": query,
|
||||
"response_mode": "blocking",
|
||||
"user": "admin"
|
||||
})
|
||||
|
||||
try:
|
||||
logger.info("调用带inputs参数的API")
|
||||
response = requests.post(
|
||||
self.api_url,
|
||||
headers=self.headers,
|
||||
data=payload,
|
||||
timeout=1200
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
logger.info(f"API调用成功，状态码: {response.status_code}")
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"status_code": response.status_code,
|
||||
"data": response.json()
|
||||
}
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"API调用失败: {e}")
|
||||
return {
|
||||
"success": False,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
def call_api_without_inputs(self, other_params: Dict[str, Any] = None, query:str = None) -> Dict[str, Any]:
|
||||
"""直接调用不带inputs参数的API
|
||||
|
||||
Args:
|
||||
other_params: 其他参数字典
|
||||
|
||||
Returns:
|
||||
API响应数据
|
||||
"""
|
||||
payload = {
|
||||
"inputs": {},
|
||||
"query": "",
|
||||
"response_mode": "blocking",
|
||||
"user": "admin"
|
||||
}
|
||||
|
||||
# 添加其他参数
|
||||
if other_params:
|
||||
payload.update(other_params)
|
||||
|
||||
try:
|
||||
logger.info("调用不带inputs参数的API")
|
||||
response = requests.post(
|
||||
self.api_url,
|
||||
headers=self.headers,
|
||||
data=json.dumps(payload),
|
||||
timeout=1200
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
logger.info(f"API调用成功，状态码: {response.status_code}")
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"status_code": response.status_code,
|
||||
"data": response.json()
|
||||
}
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"API调用失败: {e}")
|
||||
return {
|
||||
"success": False,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
class DatabaseManager:
|
||||
"""数据库管理类"""
|
||||
|
||||
def __init__(self, db_config: Dict[str, str]):
|
||||
self.db_config = db_config
|
||||
self.connection = None
|
||||
|
||||
def connect(self):
|
||||
"""连接数据库"""
|
||||
try:
|
||||
self.connection = psycopg2.connect(
|
||||
host=self.db_config.get('host', 'localhost'),
|
||||
port=self.db_config.get('port', '5432'),
|
||||
database=self.db_config.get('database', 'postgres'),
|
||||
user=self.db_config.get('user', 'postgres'),
|
||||
password=self.db_config.get('password', ''),
|
||||
cursor_factory=RealDictCursor
|
||||
)
|
||||
logger.info("数据库连接成功")
|
||||
except Exception as e:
|
||||
logger.error(f"数据库连接失败: {e}")
|
||||
raise
|
||||
|
||||
def disconnect(self):
|
||||
"""断开数据库连接"""
|
||||
if self.connection:
|
||||
self.connection.close()
|
||||
logger.info("数据库连接已关闭")
|
||||
|
||||
def get_urls_from_database(self, query: str = None) -> List[str]:
|
||||
"""从数据库获取URL列表"""
|
||||
if not self.connection:
|
||||
self.connect()
|
||||
|
||||
# 默认查询语句，根据实际情况修改
|
||||
if query is None:
|
||||
query = "SELECT 1"
|
||||
|
||||
try:
|
||||
with self.connection.cursor() as cursor:
|
||||
cursor.execute(query)
|
||||
results = cursor.fetchall()
|
||||
|
||||
# 处理结果，返回字符串列表
|
||||
datas = [list(row.values())[0] for row in results]
|
||||
return datas
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"查询数据库失败: {e}")
|
||||
return []
|
||||
#DATABASE_URL=postgresql://dbuser_dba:EmBRxnmmjnE3@124.221.232.219:5432/daas_mpp
|
||||
#DATABASE_SCHEMA=p70_ai_intelligence
|
||||
def main():
|
||||
"""主函数"""
|
||||
# 数据库配置
|
||||
db_config = {
|
||||
'host': '124.221.232.219',
|
||||
'port': '5432',
|
||||
'database': 'daas_mpp',
|
||||
'user': 'dbuser_dba',
|
||||
'password': 'EmBRxnmmjnE3',
|
||||
'schema': 'p70_ai_intelligence'
|
||||
}
|
||||
|
||||
# API配置
|
||||
api_config = {
|
||||
'url': 'https://agent.idgvalue.com/v1/chat-messages',
|
||||
'auth_token': 'Bearer app-yAPAzkPbAaV8l1SuZBjXHTU4'
|
||||
}
|
||||
|
||||
api_client = APIClient(api_config['url'], api_config['auth_token'])
|
||||
|
||||
try:
|
||||
flag = True
|
||||
|
||||
if flag:
|
||||
# 初始化
|
||||
db_manager = DatabaseManager(db_config)
|
||||
custom_query = """ select "客户名称" from p99_temp.tk_product_prefer_list """
|
||||
|
||||
try:
|
||||
# 从数据库获取URL列表
|
||||
list = db_manager.get_urls_from_database(custom_query)
|
||||
|
||||
if not list:
|
||||
logger.warning("未获取到URL")
|
||||
return
|
||||
|
||||
# 遍历每个URL调用API
|
||||
for i, text in enumerate(list, 1):
|
||||
logger.info(f"处理第 {i}/{len(list)} 个数据: {text}")
|
||||
|
||||
# 方法1: 使用带inputs参数的调用
|
||||
inputs_data = {}
|
||||
|
||||
|
||||
result = api_client.call_api_with_inputs(inputs_data,query)
|
||||
|
||||
if result['success']:
|
||||
logger.info(f"URL {text} 处理成功")
|
||||
else:
|
||||
logger.error(f"URL {text} 处理失败: {result.get('error')}")
|
||||
|
||||
# 可选：添加延迟避免请求过于频繁
|
||||
if i < len(text):
|
||||
time.sleep(1)
|
||||
except Exception as e:
|
||||
logger.error(f"程序执行失败: {e}")
|
||||
|
||||
finally:
|
||||
db_manager.disconnect()
|
||||
|
||||
else:
|
||||
logger.info("调用不带inputs参数的API示例")
|
||||
result2 = api_client.call_api_without_inputs()
|
||||
|
||||
if result2['success']:
|
||||
logger.info("不带inputs参数的API调用成功")
|
||||
else:
|
||||
logger.error(f"不带inputs参数的API调用失败: {result2.get('error')}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"初始化失败: {e}")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,53 @@
|
|||
#!/usr/bin/python
|
||||
# -*- encoding=utf-8 -*-
|
||||
from airflow import DAG
|
||||
from datetime import datetime, timedelta
|
||||
from airflow.contrib.hooks.ssh_hook import SSHHook
|
||||
from airflow.contrib.operators.ssh_operator import SSHOperator
|
||||
from airflow.sensors.external_task_sensor import ExternalTaskSensor
|
||||
import json
|
||||
|
||||
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
|
||||
from airflow.operators.email_operator import EmailOperator
|
||||
from airflow.utils.trigger_rule import TriggerRule
|
||||
|
||||
|
||||
sshHook = SSHHook(ssh_conn_id ='ssh_air')
|
||||
default_args = {
|
||||
'owner': 'tek_newsletter@163.com',
|
||||
'email_on_failure': True,
|
||||
'email_on_retry':True,
|
||||
'start_date': datetime(2024, 1, 1),
|
||||
'depends_on_past': False,
|
||||
'retries': 6,
|
||||
'retry_delay': timedelta(minutes=10),
|
||||
}
|
||||
|
||||
dag = DAG('wf_dag_product_infer_list', default_args=default_args,
|
||||
schedule=None,
|
||||
catchup=False,
|
||||
dagrun_timeout=timedelta(minutes=160),
|
||||
max_active_runs=3,
|
||||
tags=['64cbd1bbced14209b5e3a879f89e8ab1','TK_Cust','泰克产品相关性分析']
|
||||
)
|
||||
|
||||
task_failed = EmailOperator (
|
||||
dag=dag,
|
||||
trigger_rule=TriggerRule.ONE_FAILED,
|
||||
task_id="task_failed",
|
||||
to=["tek_newsletter@163.com"],
|
||||
cc=[""],
|
||||
subject="product_infer_list_failed",
|
||||
html_content='<h3>您好，product_infer_list作业失败，请及时处理" </h3>')
|
||||
|
||||
|
||||
tk_product_prefer = SSHOperator(
|
||||
ssh_hook=sshHook,
|
||||
task_id='tk_product_prefer',
|
||||
command='cd /data/airflow/etl/agents && python tk_product_prefer.py',
|
||||
params={'my_param':"tk_product_prefer"},
|
||||
depends_on_past=False,
|
||||
retries=1,
|
||||
dag=dag)
|
||||
|
||||
tk_product_prefer >> task_failed
|
||||
Loading…
Reference in New Issue