codekingpro/portable-devtools

1from typing import Any, Dict, Iterator, List, Optional2 3from langchain_core.documents import Document4 5from langchain_community.document_loaders.base import BaseLoader6 7 8class TiDBLoader(BaseLoader):9    """Load documents from TiDB."""10 11    def __init__(12        self,13        connection_string: str,14        query: str,15        page_content_columns: Optional[List[str]] = None,16        metadata_columns: Optional[List[str]] = None,17        engine_args: Optional[Dict[str, Any]] = None,18    ) -> None:19        """Initialize TiDB document loader.20 21        Args:22            connection_string (str): The connection string for the TiDB database,23                format: "mysql+pymysql://root@127.0.0.1:4000/test".24            query: The query to run in TiDB.25            page_content_columns: Optional. Columns written to Document `page_content`,26                default(None) to all columns.27            metadata_columns: Optional. Columns written to Document `metadata`,28                default(None) to no columns.29            engine_args: Optional. Additional arguments to pass to sqlalchemy engine.30        """31        self.connection_string = connection_string32        self.query = query33        self.page_content_columns = page_content_columns34        self.metadata_columns = metadata_columns if metadata_columns is not None else []35        self.engine_args = engine_args36 37    def lazy_load(self) -> Iterator[Document]:38        """Lazy load TiDB data into document objects."""39 40        from sqlalchemy import create_engine41        from sqlalchemy.engine import Engine42        from sqlalchemy.sql import text43 44        # use sqlalchemy to create db connection45        engine: Engine = create_engine(46            self.connection_string, **(self.engine_args or {})47        )48 49        # execute query50        with engine.connect() as conn:51            result = conn.execute(text(self.query))52 53            # convert result to Document objects54            column_names = list(result.keys())55            for row in result:56                # convert row to dict{column:value}57                row_data = {58                    column_names[index]: value for index, value in enumerate(row)59                }60                page_content = "\n".join(61                    f"{k}: {v}"62                    for k, v in row_data.items()63                    if self.page_content_columns is None64                    or k in self.page_content_columns65                )66                metadata = {col: row_data[col] for col in self.metadata_columns}67                yield Document(page_content=page_content, metadata=metadata)68