#!/usr/bin/env python
# coding: utf-8

# <a href="https://colab.research.google.com/github/jerryjliu/llama_index/blob/main/docs/docs/examples/data_connectors/DatabaseReaderDemo.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="在 Colab 中打开"/></a>
# 

# # 数据库阅读器
# 

# 如果您在colab上打开这个笔记本，您可能需要安装LlamaIndex 🦙。
# 

# In[ ]:


get_ipython().run_line_magic('pip', 'install llama-index-readers-database')


# In[ ]:


get_ipython().system('pip install llama-index')


# In[ ]:


import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


# In[ ]:


from __future__ import absolute_import

# My OpenAI Key
import os

os.environ["OPENAI_API_KEY"] = ""

from llama_index.readers.database import DatabaseReader
from llama_index.core import VectorStoreIndex


# In[ ]:


# 使用以下参数初始化DatabaseReader对象：

db = DatabaseReader(
    scheme="postgresql",  # 数据库方案
    host="localhost",  # 数据库主机
    port="5432",  # 数据库端口
    user="postgres",  # 数据库用户
    password="FakeExamplePassword",  # 数据库密码
    dbname="postgres",  # 数据库名称
)


# In[ ]:


### DatabaseReader 类 ###
# db 是 DatabaseReader 的一个实例：
print(type(db))
# DatabaseReader 可用的方法：
print(type(db.load_data))

### SQLDatabase 类 ###
# db.sql 是 SQLDatabase 的一个实例：
print(type(db.sql_database))
# SQLDatabase 可用的方法：
print(type(db.sql_database.from_uri))
print(type(db.sql_database.get_single_table_info))
print(type(db.sql_database.get_table_columns))
print(type(db.sql_database.get_usable_table_names))
print(type(db.sql_database.insert_into_table))
print(type(db.sql_database.run_sql))
# SQLDatabase 可用的属性：
print(type(db.sql_database.dialect))
print(type(db.sql_database.engine))


# In[ ]:


### 测试 DatabaseReader
### 来自 SQLDatabase，SQLAlchemy 引擎和数据库 URI：

# 从 SQLDatabase 实例中：
print(type(db.sql_database))
db_from_sql_database = DatabaseReader(sql_database=db.sql_database)
print(type(db_from_sql_database))

# 从 SQLAlchemy 引擎中：
print(type(db.sql_database.engine))
db_from_engine = DatabaseReader(engine=db.sql_database.engine)
print(type(db_from_engine))

# 从数据库 URI 中：
print(type(db.uri))
db_from_uri = DatabaseReader(uri=db.uri)
print(type(db_from_uri))


# In[ ]:


# 以下SQL查询示例返回用户表中年龄大于或等于18岁的每行的名称和年龄列的连接文本值列表

query = f"""
    SELECT
        CONCAT(name, ' is ', age, ' years old.') AS text
    FROM public.users
    WHERE age >= 18
    """


# In[ ]:


# 请参考llama_index.utilities.sql_wrapper
# SQLDatabase.run_sql方法
texts = db.sql_database.run_sql(command=query)

# 显示type(texts)和texts
# type(texts)必须返回<class 'list'>
print(type(texts))

# 文档必须返回一个元组对象的列表
print(texts)


# In[ ]:


# 请参考llama_index.readers.database.DatabaseReader.load_data
# DatabaseReader.load_data方法
documents = db.load_data(query=query)

# 显示type(documents)和documents
# type(documents)必须返回<class 'list'>
print(type(documents))

# documents必须返回一个Document对象的列表
print(documents)


# In[ ]:


index = VectorStoreIndex.from_documents(documents)