{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/data_connectors/DatabaseReaderDemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Database Reader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-readers-database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "\n",
    "# My OpenAI Key\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
    "\n",
    "from llama_index.readers.database import DatabaseReader\n",
    "from llama_index.core import VectorStoreIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize DatabaseReader object with the following parameters:\n",
    "\n",
    "db = DatabaseReader(\n",
    "    scheme=\"postgresql\",  # Database Scheme\n",
    "    host=\"localhost\",  # Database Host\n",
    "    port=\"5432\",  # Database Port\n",
    "    user=\"postgres\",  # Database User\n",
    "    password=\"FakeExamplePassword\",  # Database Password\n",
    "    dbname=\"postgres\",  # Database Name\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### DatabaseReader class ###\n",
    "# db is an instance of DatabaseReader:\n",
    "print(type(db))\n",
    "# DatabaseReader available method:\n",
    "print(type(db.load_data))\n",
    "\n",
    "### SQLDatabase class ###\n",
    "# db.sql is an instance of SQLDatabase:\n",
    "print(type(db.sql_database))\n",
    "# SQLDatabase available methods:\n",
    "print(type(db.sql_database.from_uri))\n",
    "print(type(db.sql_database.get_single_table_info))\n",
    "print(type(db.sql_database.get_table_columns))\n",
    "print(type(db.sql_database.get_usable_table_names))\n",
    "print(type(db.sql_database.insert_into_table))\n",
    "print(type(db.sql_database.run_sql))\n",
    "# SQLDatabase available properties:\n",
    "print(type(db.sql_database.dialect))\n",
    "print(type(db.sql_database.engine))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Testing DatabaseReader\n",
    "### from SQLDatabase, SQLAlchemy engine and Database URI:\n",
    "\n",
    "# From SQLDatabase instance:\n",
    "print(type(db.sql_database))\n",
    "db_from_sql_database = DatabaseReader(sql_database=db.sql_database)\n",
    "print(type(db_from_sql_database))\n",
    "\n",
    "# From SQLAlchemy engine:\n",
    "print(type(db.sql_database.engine))\n",
    "db_from_engine = DatabaseReader(engine=db.sql_database.engine)\n",
    "print(type(db_from_engine))\n",
    "\n",
    "# From Database URI:\n",
    "print(type(db.uri))\n",
    "db_from_uri = DatabaseReader(uri=db.uri)\n",
    "print(type(db_from_uri))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The below SQL Query example returns a list values of each row\n",
    "# with concatenated text from the name and age columns\n",
    "# from the users table where the age is greater than or equal to 18\n",
    "\n",
    "query = f\"\"\"\n",
    "    SELECT\n",
    "        CONCAT(name, ' is ', age, ' years old.') AS text\n",
    "    FROM public.users\n",
    "    WHERE age >= 18\n",
    "    \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Please refer to llama_index.utilities.sql_wrapper\n",
    "# SQLDatabase.run_sql method\n",
    "texts = db.sql_database.run_sql(command=query)\n",
    "\n",
    "# Display type(texts) and texts\n",
    "# type(texts) must return <class 'list'>\n",
    "print(type(texts))\n",
    "\n",
    "# Documents must return a list of Tuple objects\n",
    "print(texts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Please refer to llama_index.readers.database.DatabaseReader.load_data\n",
    "# DatabaseReader.load_data method\n",
    "documents = db.load_data(query=query)\n",
    "\n",
    "# Display type(documents) and documents\n",
    "# type(documents) must return <class 'list'>\n",
    "print(type(documents))\n",
    "\n",
    "# Documents must return a list of Document objects\n",
    "print(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  },
  "vscode": {
   "interpreter": {
    "hash": "bd5508c2ffc7f17f7d31cf4086cc872f89e96996a08987e995649e5fbe85a3a4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}