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

[docs]class TiDBLoader(BaseLoader): """Load documents from TiDB."""
[docs] 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@". 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
[docs] 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)