使用外部数据源构建Data-Agent

此处以selectDB为例，为您展示如何批量导入至PolarDB for AI的schema_meta表中。

1. 安装依赖库：pip install pymysql。

2. 将以下代码保存在您的业务环境中。

   selectdb_data_transfer.py

   from mysql_connector import MysqlConnector
   
   import json
   
   def parse_ddl_from_json(file_name: str) -> list[str]:
       ddl_statements = []
       with open(file_name, 'r', encoding='utf-8') as f:
           # 使用 json.load() 来解析 JSON 文件内容为一个 Python 对象
           data = json.load(f)
   
           # 遍历列表中的每一个字典
           for table_info in data:
               # 从字典中获取 'Create Table' 键对应的值，即 DDL 语句
               create_table_sql = table_info.get('Create Table')
               if create_table_sql:
                   ddl_statements.append(create_table_sql)
   
       return ddl_statements
   
   def insert_into_schema_meta(db_name, ddls: list[str], connector: MysqlConnector):
       """
       通过connector连接polardb
       """
       insert_sql = """insert into schema_meta (db_name, create_table_statement) values (%s, %s)"""
       for i, ddl in enumerate(ddls):
           rows = connector.execute_update(insert_sql, (db_name, ddl))
           if rows == 0:
               print('insert failed: ', i)
               break
   
   def show_index(connector: MysqlConnector, db_name):
       table_names = connector.execute_query(
           f"select table_name FROM INFORMATION_SCHEMA.TABLES WHERE table_schema = '{db_name}'")
       table_ddls = []
       for result in table_names:
           # 请根据关系性数据库类型，修改返回的参数值。例如，PolarDB应为TABLE_NAME。
           table_name = result['table_name']
           table_ddl = connector.execute_query(f"show create table {table_name}")[0]
           table_ddls.append(table_ddl)
       json.dump(table_ddls, open('table_ddls.json', 'w'), indent=4)
   
   
   if __name__ == '__main__':
       selectdb_host = "<SelectDB连接地址>"
       selectdb_user = "<SelectDB用户名>"
       selectdb_password = "<SelectDB密码>"
       selectdb_database = "<SelectDB数据库>"
       selectdb_port = <SelectDB端口>
   
       polar_host = "<PolarDB连接地址>"
       polar_user = "<PolarDB用户名>"
       polar_password = "<PolarDB密码>"
       polar_database = "<PolarDB数据库>"
       polar_port = <PolarDB端口>
   
       # # 获取所有表的DDL
       # selectdb_connector = MysqlConnector(selectdb_host, selectdb_port, selectdb_user, selectdb_password, selectdb_database)
       # show_index(selectdb_connector, selectdb_database)
       
       # # 解析.json, 获得DDL语句
       # results = parse_ddl_from_json('table_ddls.json')
       # 
       # # 连接数据库
       # polar_mysql_connector = MysqlConnector(polar_host, polar_port, polar_user, polar_password, polar_database)
       # # 批量导入至PolarDB中
       # insert_into_schema_meta(polar_database, results, polar_mysql_connector)

   mysql_connector.py

   import pymysql
   from pymysql import Error
   from pymysql.cursors import DictCursor
   
   class MysqlConnector:
       def __init__(self, host, port, user, password, database, charset='utf8mb4'):
           """
           初始化数据库连接
           :param host: MySQL 主机地址
           :param user: 用户名
           :param password: 密码
           :param database: 数据库名
           :param port: 端口号，默认3306
           :param charset: 字符集，默认 utf8mb4
           """
           self.connection = pymysql.connect(
               host=host,
               user=user,
               password=password,
               database=database,
               port=port,
               charset=charset,
               cursorclass=DictCursor  # 结果以字典形式返回
           )
   
       def execute_query(self, sql, params=None):
           """
           执行 SELECT 查询语句
           :param sql: SQL 语句
           :param params: 参数字典或元组
           :return: 查询结果列表（每个元素为字典）
           """
           try:
               with self.connection.cursor() as cursor:
                   cursor.execute(sql, params)
                   result = cursor.fetchall()
                   return result
           except Error as e:
               print(f"[ERROR] 查询执行失败: {e}")
               return None
   
       def execute_update(self, sql, params=None):
           """
           执行 UPDATE / INSERT / DELETE 等语句
           :param sql: SQL 语句
           :param params: 参数字典或元组
           :return: 受影响的行数
           """
           try:
               with self.connection.cursor() as cursor:
                   cursor.execute(sql, params)
               self.connection.commit()
               return cursor.rowcount
           except Error as e:
               print(f"[ERROR] 更新执行失败: {e}")
               self.connection.rollback()
               return 0
       def execute_update_multirows(self, sql, params_list=None):
           """
           执行 UPDATE / INSERT / DELETE 等语句
           :param sql: SQL 语句
           :param params: 参数字典或元组
           :return: 受影响的行数
           """
           try:
               with self.connection.cursor() as cursor:
                   cursor.executemany(sql, params_list)
               self.connection.commit()
               return cursor.rowcount
           except Error as e:
               print(f"[ERROR] 更新执行失败: {e}")
               self.connection.rollback()
               return 0
   
       def begin_transaction(self):
           """ 开始事务 """
           try:
               self.connection.begin()
           except Error as e:
               print(f"[ERROR] 事务开始失败: {e}")
   
       def commit_transaction(self):
           """ 提交事务 """
           try:
               self.connection.commit()
           except Error as e:
               print(f"[ERROR] 事务提交失败: {e}")
   
       def rollback_transaction(self):
           """ 回滚事务 """
           try:
               self.connection.rollback()
           except Error as e:
               print(f"[ERROR] 事务回滚失败: {e}")
   
       def close(self):
           """ 关闭数据库连接 """
           if self.connection:
               self.connection.close()
   
       def __enter__(self):
           return self
   
       def __exit__(self, exc_type, exc_val, exc_tb):
           self.close()
   

3. 修改脚本中的SelectDB和PolarDB的连接参数。

4. 运行脚本。建议您首先获取所有表的DDL，以确认其正确性，然后再进行批量导入。

1. 将以下代码保存为es_data_transfer.py。

   import base64
   import http.client
   import json
   ## 集群host
   ClusterHost = "<elasticsearch-host>:<elasticsearch-port>"
   ## 集群用户名
   ClusterUserName = "<elasticsearch-username>"
   ## 集群密码
   ClusterPassword = "<elasticsearch-password>"
   DEFAULT_REPLICAS = 0
   
   def httpRequest(method, host, endpoint, params="", username="", password=""):
       conn = http.client.HTTPConnection(host)
       headers = {}
       if username:
           auth_str = f"{username}:{password}".encode('utf-8')
           base64string = base64.b64encode(auth_str).decode('utf-8')
           headers["Authorization"] = "Basic {}".format(base64string)
       if method == "GET":
           headers["Content-Type"] = "application/x-www-form-urlencoded"
           conn.request(method=method, url=endpoint, headers=headers)
       else:
           headers["Content-Type"] = "application/json"
           if params:
               conn.request(method=method, url=endpoint, body=params.encode('utf-8'), headers=headers)
           else:
               conn.request(method=method, url=endpoint, headers=headers)
       response = conn.getresponse()
       res = response.read().decode('utf-8')
       return res
   
   def httpGet(host, endpoint, username="", password=""):
       return httpRequest("GET", host, endpoint, "", username, password)
   def httpPost(host, endpoint, params, username="", password=""):
       return httpRequest("POST", host, endpoint, params, username, password)
   def httpPut(host, endpoint, params, username="", password=""):
       return httpRequest("PUT", host, endpoint, params, username, password)
   def getIndices(host, username="", password=""):
       endpoint = "/_cat/indices"
       indicesResult = httpGet(ClusterHost, endpoint, ClusterUserName, ClusterPassword)
       indicesList = indicesResult.split("\n")
       indexList = []
       for indices in indicesList:
           if (indices.find("open") > 0):
               indexList.append(indices.split()[2])
       return indexList
   def getSettings(index, host, username="", password=""):
       endpoint = "/" + index + "/_settings"
       indexSettings = httpGet(host, endpoint, username, password)
       settingsDict = json.loads(indexSettings)
       ## 分片数默认和源集群索引保持一致。
       number_of_shards = settingsDict[index]["settings"]["index"]["number_of_shards"]
       ## 副本数默认为0。
       number_of_replicas = DEFAULT_REPLICAS
       newSetting = "\"settings\": {\"number_of_shards\": %s, \"number_of_replicas\": %s}" % (number_of_shards, number_of_replicas)
       return newSetting
   def getMapping(index, host, username="", password=""):
       endpoint = "/" + index + "/_mapping"
       indexMapping = httpGet(host, endpoint, username, password)
       mappingDict = json.loads(indexMapping)
       mappings = json.dumps(mappingDict[index]["mappings"])
       newMapping = "\"mappings\" : " + mappings
       return newMapping
   def createIndexStatement(IndexName):
       settingStr = getSettings(IndexName, ClusterHost, ClusterUserName, ClusterPassword)
       mappingStr = getMapping(IndexName, ClusterHost, ClusterUserName, ClusterPassword)
       createstatement = "{\n" + str(settingStr) + ",\n" + str(mappingStr) + "\n}"
       return createstatement
   def createIndex(IndexName, newIndexName=""):
       if (newIndexName == "") :
           newIndexName = IndexName
       createstatement = createIndexStatement(IndexName)
       print ("PUT /" + newIndexName + "\n" + createstatement)
   
   indexList = getIndices(ClusterHost, ClusterUserName, ClusterPassword)
   systemIndex = []
   for index in indexList:
       if (index.startswith(".")):
           systemIndex.append(index)
       else :
           createIndex(index, index)
   # # 取消注释可检查被自动识别和省略的系统索引
   # if (len(systemIndex) > 0) :
   #     for index in systemIndex:
   #         print (index + " 或许是系统索引，不会重新创建，如有需要，请单独处理～")

2. 修改脚本中的Elasticsearch的连接参数。

3. 运行脚本，将导出Elasticsearch中所有索引的创建语句，示例如下：

   PUT /product_info
   {
   "settings": {"number_of_shards": 5, "number_of_replicas": 0},
   "mappings" : {"properties": {"annual_rate": {"type": "keyword"}, "describe": {"type": "text", "analyzer": "ik_smart"}, "productName": {"type": "text", "analyzer": "ik_smart"}}}
   }
   PUT /user_info
   {
   "settings": {"number_of_shards": 5, "number_of_replicas": 0},
   "mappings" : {"properties": {"address": {"type": "text", "analyzer": "ik_smart"}, "describe": {"type": "text", "analyzer": "ik_smart"}, "userName": {"type": "text", "analyzer": "ik_smart"}}}
   }

4. 将其保存并插入到PolarDB for AI的schema_meta表中。