mysql.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. from .base import MetadataProvider
  2. from typing import List, Dict, Any
  3. from sqlalchemy import create_engine, text, inspect
  4. class MySQLMetadataProvider(MetadataProvider):
  5. def __init__(self, config: Dict):
  6. super().__init__(config)
  7. url = f"mysql+pymysql://{config['username']}:{config.get('password','')}@{config['host']}:{config['port']}/{config.get('database','')}?charset=utf8mb4"
  8. self.engine = create_engine(url, echo=True ) # 设为 True 可打印 SQL 日志
  9. self.inspector = inspect(self.engine)
  10. def _get_connection(self):
  11. return self.engine.connect()
  12. # return pymysql.connect(
  13. # host=self.config["host"],
  14. # port=self.config["port"],
  15. # user=self.config["username"],
  16. # password=self.config.get("password", ""),
  17. # database=self.config.get("database", ""), # 可为空
  18. # charset='utf8mb4'
  19. # )
  20. def get_databases(self) -> List[str]:
  21. # with self._get_connection() as conn:
  22. # # cursor = conn.cursor()
  23. # # cursor.execute("SHOW DATABASES")
  24. # # return [row[0] for row in cursor.fetchall() if not row[0].startswith('information_schema')]
  25. # result = conn.execute(text("SHOW DATABASES"))
  26. # return [row[0] for row in result if not row[0].startswith('information_schema')]
  27. return [row for row in self.inspector.get_schema_names() if not row.startswith('information_schema')]
  28. def get_tables(self, database: str) -> List[str]:
  29. return self.inspector.get_table_names(schema=database)
  30. def get_columns(self, table: str, database: str) -> List[Dict[str, Any]]:
  31. return self.inspector.get_columns(table_name=table, schema=database)
  32. # with self._get_connection() as conn:
  33. # cursor = conn.cursor()
  34. # # 查询 information_schema 获取字段详情
  35. # query = """
  36. # SELECT COLUMN_NAME, DATA_TYPE, COLUMN_COMMENT
  37. # FROM information_schema.COLUMNS
  38. # WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
  39. # ORDER BY ORDINAL_POSITION
  40. # """
  41. # cursor.execute(query, (database, table))
  42. # return [
  43. # {
  44. # "name": row[0],
  45. # "type": row[1],
  46. # "comment": row[2] or ""
  47. # }
  48. # for row in cursor.fetchall()
  49. # ]
  50. def has_table(self, table_name:str, schema:str) -> bool:
  51. return self.inspector.has_table(table_name=table_name, schema=schema)
  52. def get_table_comment(self, table_name:str, schema:str) -> str:
  53. return self.inspector.get_table_comment(table_name=table_name, schema=schema)['text']
  54. def get_pk_constraint(self, table_name:str, schema:str) -> Dict:
  55. return self.inspector.get_pk_constraint(table_name=table_name, schema=schema)
  56. def get_unique_constraints(self, table_name:str, schema:str) -> Dict:
  57. return self.inspector.get_unique_constraints(table_name=table_name, schema=schema)
  58. def get_foreign_keys(self, table_name:str, schema:str) -> list[Dict]:
  59. return self.inspector.get_foreign_keys(table_name=table_name, schema=schema)
  60. def fectch_distinct_values(self, column_name:str, table_name:str, schema:str, max_num:int = 5) -> list[str]:
  61. with self._get_connection() as conn:
  62. result = conn.execute(text(f"select distinct `{column_name}` from `{schema}`.`{table_name}` where `{column_name}` is not null limit {max_num}"))
  63. return [row[0] for row in result]