llama-index

1
"""Database Reader."""
2

3
from typing import Any, List, Optional
4

5
from sqlalchemy import text
6
from sqlalchemy.engine import Engine
7

8
from llama_index.legacy.readers.base import BaseReader
9
from llama_index.legacy.schema import Document
10
from llama_index.legacy.utilities.sql_wrapper import SQLDatabase
11

12

13
class DatabaseReader(BaseReader):
14
    """Simple Database reader.
15

16
    Concatenates each row into Document used by LlamaIndex.
17

18
    Args:
19
        sql_database (Optional[SQLDatabase]): SQL database to use,
20
            including table names to specify.
21
            See :ref:`Ref-Struct-Store` for more details.
22

23
        OR
24

25
        engine (Optional[Engine]): SQLAlchemy Engine object of the database connection.
26

27
        OR
28

29
        uri (Optional[str]): uri of the database connection.
30

31
        OR
32

33
        scheme (Optional[str]): scheme of the database connection.
34
        host (Optional[str]): host of the database connection.
35
        port (Optional[int]): port of the database connection.
36
        user (Optional[str]): user of the database connection.
37
        password (Optional[str]): password of the database connection.
38
        dbname (Optional[str]): dbname of the database connection.
39

40
    Returns:
41
        DatabaseReader: A DatabaseReader object.
42
    """
43

44
    def __init__(
45
        self,
46
        sql_database: Optional[SQLDatabase] = None,
47
        engine: Optional[Engine] = None,
48
        uri: Optional[str] = None,
49
        scheme: Optional[str] = None,
50
        host: Optional[str] = None,
51
        port: Optional[str] = None,
52
        user: Optional[str] = None,
53
        password: Optional[str] = None,
54
        dbname: Optional[str] = None,
55
        *args: Any,
56
        **kwargs: Any,
57
    ) -> None:
58
        """Initialize with parameters."""
59
        if sql_database:
60
            self.sql_database = sql_database
61
        elif engine:
62
            self.sql_database = SQLDatabase(engine, *args, **kwargs)
63
        elif uri:
64
            self.uri = uri
65
            self.sql_database = SQLDatabase.from_uri(uri, *args, **kwargs)
66
        elif scheme and host and port and user and password and dbname:
67
            uri = f"{scheme}://{user}:{password}@{host}:{port}/{dbname}"
68
            self.uri = uri
69
            self.sql_database = SQLDatabase.from_uri(uri, *args, **kwargs)
70
        else:
71
            raise ValueError(
72
                "You must provide either a SQLDatabase, "
73
                "a SQL Alchemy Engine, a valid connection URI, or a valid "
74
                "set of credentials."
75
            )
76

77
    def load_data(self, query: str) -> List[Document]:
78
        """Query and load data from the Database, returning a list of Documents.
79

80
        Args:
81
            query (str): Query parameter to filter tables and rows.
82

83
        Returns:
84
            List[Document]: A list of Document objects.
85
        """
86
        documents = []
87
        with self.sql_database.engine.connect() as connection:
88
            if query is None:
89
                raise ValueError("A query parameter is necessary to filter the data")
90
            else:
91
                result = connection.execute(text(query))
92

93
            for item in result.fetchall():
94
                # fetch each item
95
                doc_str = ", ".join(
96
                    [f"{col}: {entry}" for col, entry in zip(result.keys(), item)]
97
                )
98
                documents.append(Document(text=doc_str))
99
        return documents
100