PolarDB:Build a Data-Agent using an external data source

This example uses SelectDB as the source database. The script retrieves DDL for all tables and inserts them into the schema_meta table in PolarDB for AI.

1. Install the dependency: pip install pymysql.

2. Save the following files in your working directory.

   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:
           # Parse the JSON file into a Python object
           data = json.load(f)
   
           # Extract the DDL statement from each table entry
           for table_info in data:
               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):
       """
       Insert DDL statements into the schema_meta table in 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:
           # Adjust the column name based on the source database type.
           # For PolarDB, use 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 endpoint>"
       selectdb_user = "<SelectDB username>"
       selectdb_password = "<SelectDB password>"
       selectdb_database = "<SelectDB database>"
       selectdb_port = <SelectDB port>
   
       polar_host = "<PolarDB endpoint>"
       polar_user = "<PolarDB username>"
       polar_password = "<PolarDB password>"
       polar_database = "<PolarDB database>"
       polar_port = <PolarDB port>
   
       # Step 1: Retrieve DDL for all tables from SelectDB
       # selectdb_connector = MysqlConnector(selectdb_host, selectdb_port, selectdb_user, selectdb_password, selectdb_database)
       # show_index(selectdb_connector, selectdb_database)
   
       # Step 2: Parse the exported JSON file to get DDL statements
       # results = parse_ddl_from_json('table_ddls.json')
       #
       # Step 3: Insert DDL into PolarDB schema_meta
       # polar_mysql_connector = MysqlConnector(polar_host, polar_port, polar_user, polar_password, polar_database)
       # 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'):
           """
           Initialize the database connection.
           :param host: MySQL host address
           :param user: Username
           :param password: Password
           :param database: Database name
           :param port: Port number, default is 3306
           :param charset: Character set, default is utf8mb4
           """
           self.connection = pymysql.connect(
               host=host,
               user=user,
               password=password,
               database=database,
               port=port,
               charset=charset,
               cursorclass=DictCursor  # Return results as dictionaries
           )
   
       def execute_query(self, sql, params=None):
           """
           Execute a SELECT query.
           :param sql: SQL statement
           :param params: Parameter dictionary or tuple
           :return: List of query results (each element is a dictionary)
           """
           try:
               with self.connection.cursor() as cursor:
                   cursor.execute(sql, params)
                   result = cursor.fetchall()
                   return result
           except Error as e:
               print(f"[ERROR] Query execution failed: {e}")
               return None
   
       def execute_update(self, sql, params=None):
           """
           Execute UPDATE, INSERT, or DELETE statements.
           :param sql: SQL statement
           :param params: Parameter dictionary or tuple
           :return: Number of affected rows
           """
           try:
               with self.connection.cursor() as cursor:
                   cursor.execute(sql, params)
               self.connection.commit()
               return cursor.rowcount
           except Error as e:
               print(f"[ERROR] Update execution failed: {e}")
               self.connection.rollback()
               return 0
   
       def execute_update_multirows(self, sql, params_list=None):
           """
           Execute batch INSERT, UPDATE, or DELETE statements.
           :param sql: SQL statement
           :param params_list: List of parameter tuples
           :return: Number of affected rows
           """
           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] Update execution failed: {e}")
               self.connection.rollback()
               return 0
   
       def begin_transaction(self):
           """ Start a transaction. """
           try:
               self.connection.begin()
           except Error as e:
               print(f"[ERROR] Failed to start transaction: {e}")
   
       def commit_transaction(self):
           """ Commit the transaction. """
           try:
               self.connection.commit()
           except Error as e:
               print(f"[ERROR] Failed to commit transaction: {e}")
   
       def rollback_transaction(self):
           """ Roll back the transaction. """
           try:
               self.connection.rollback()
           except Error as e:
               print(f"[ERROR] Failed to roll back transaction: {e}")
   
       def close(self):
           """ Close the database connection. """
           if self.connection:
               self.connection.close()
   
       def __enter__(self):
           return self
   
       def __exit__(self, exc_type, exc_val, exc_tb):
           self.close()

3. Update the connection parameters for SelectDB and PolarDB in selectdb_data_transfer.py.

4. Run the script. First retrieve and review the DDL for all tables to confirm correctness, then run the batch import.

1. Save the following code as es_data_transfer.py.

2. import base64
   import http.client
   import json
   
   # Cluster connection settings
   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)
       # Shard count matches the source cluster index
       number_of_shards = settingsDict[index]["settings"]["index"]["number_of_shards"]
       # Replica count defaults to 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)
   # Uncomment to review system indexes that are automatically skipped
   # if (len(systemIndex) > 0) :
   #     for index in systemIndex:
   #         print (index + " might be a system index and will not be recreated. Handle it separately if needed.")

3. Update the Elasticsearch connection parameters in the script.

4. Run the script to export index creation statements. The output looks like this:

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"}}}
}

1. Insert the exported statements into the schema_meta table in PolarDB for AI.