Toggle Menu
Prev Up Next

# Source code for langchain_community.document_loaders.tidb

    from typing import Any, Dict, Iterator, List, Optional

    from langchain_core.documents import Document

    from langchain_community.document_loaders.base import BaseLoader

    class TiDBLoader(BaseLoader):
        """Load documents from TiDB."""

        def __init__(
            self,
            connection_string: str,
            query: str,
            page_content_columns: Optional[List[str]] = None,
            metadata_columns: Optional[List[str]] = None,
            engine_args: Optional[Dict[str, Any]] = None,
        ) -> None:
            """Initialize TiDB document loader.

            Args:
                connection_string (str): The connection string for the TiDB database,
                    format: "mysql+pymysql://root@127.0.0.1:4000/test".
                query: The query to run in TiDB.
                page_content_columns: Optional. Columns written to Document `page_content`,
                    default(None) to all columns.
                metadata_columns: Optional. Columns written to Document `metadata`,
                    default(None) to no columns.
                engine_args: Optional. Additional arguments to pass to sqlalchemy engine.
            """
            self.connection_string = connection_string
            self.query = query
            self.page_content_columns = page_content_columns
            self.metadata_columns = metadata_columns if metadata_columns is not None else []
            self.engine_args = engine_args

        def lazy_load(self) -> Iterator[Document]:
            """Lazy load TiDB data into document objects."""

            from sqlalchemy import create_engine
            from sqlalchemy.engine import Engine
            from sqlalchemy.sql import text

            # use sqlalchemy to create db connection
            engine: Engine = create_engine(
                self.connection_string, **(self.engine_args or {})
            )

            # execute query
            with engine.connect() as conn:
                result = conn.execute(text(self.query))

                # convert result to Document objects
                column_names = list(result.keys())
                for row in result:
                    # convert row to dict{column:value}
                    row_data = {
                        column_names[index]: value for index, value in enumerate(row)
                    }
                    page_content = "\n".join(
                        f"{k}: {v}"
                        for k, v in row_data.items()
                        if self.page_content_columns is None
                        or k in self.page_content_columns
                    )
                    metadata = {col: row_data[col] for col in self.metadata_columns}
                    yield Document(page_content=page_content, metadata=metadata)