Spaces:
Sleeping
Sleeping
File size: 8,904 Bytes
c96b98a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 | from typing import Optional, Dict, Any
try:
import psycopg2
except ImportError:
raise ImportError(
"`psycopg2` not installed. Please install using `pip install psycopg2`. If you face issues, try `pip install psycopg2-binary`."
)
from phi.tools import Toolkit
from phi.utils.log import logger
class PostgresTools(Toolkit):
"""A basic tool to connect to a PostgreSQL database and perform read-only operations on it."""
def __init__(
self,
connection: Optional[psycopg2.extensions.connection] = None,
db_name: Optional[str] = None,
user: Optional[str] = None,
password: Optional[str] = None,
host: Optional[str] = None,
port: Optional[int] = None,
run_queries: bool = True,
inspect_queries: bool = False,
summarize_tables: bool = True,
export_tables: bool = False,
table_schema: str = "public",
):
super().__init__(name="postgres_tools")
self._connection: Optional[psycopg2.extensions.connection] = connection
self.db_name: Optional[str] = db_name
self.user: Optional[str] = user
self.password: Optional[str] = password
self.host: Optional[str] = host
self.port: Optional[int] = port
self.table_schema: str = table_schema
self.register(self.show_tables)
self.register(self.describe_table)
if inspect_queries:
self.register(self.inspect_query)
if run_queries:
self.register(self.run_query)
if summarize_tables:
self.register(self.summarize_table)
if export_tables:
self.register(self.export_table_to_path)
@property
def connection(self) -> psycopg2.extensions.connection:
"""
Returns the Postgres psycopg2 connection.
:return psycopg2.extensions.connection: psycopg2 connection
"""
if self._connection is None:
connection_kwargs: Dict[str, Any] = {}
if self.db_name is not None:
connection_kwargs["database"] = self.db_name
if self.user is not None:
connection_kwargs["user"] = self.user
if self.password is not None:
connection_kwargs["password"] = self.password
if self.host is not None:
connection_kwargs["host"] = self.host
if self.port is not None:
connection_kwargs["port"] = self.port
if self.table_schema is not None:
connection_kwargs["options"] = f"-c search_path={self.table_schema}"
self._connection = psycopg2.connect(**connection_kwargs)
self._connection.set_session(readonly=True)
return self._connection
def show_tables(self) -> str:
"""Function to show tables in the database
:return: List of tables in the database
"""
stmt = f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{self.table_schema}';"
tables = self.run_query(stmt)
logger.debug(f"Tables: {tables}")
return tables
def describe_table(self, table: str) -> str:
"""Function to describe a table
:param table: Table to describe
:return: Description of the table
"""
stmt = f"SELECT column_name, data_type, character_maximum_length FROM information_schema.columns WHERE table_name = '{table}' AND table_schema = '{self.table_schema}';"
table_description = self.run_query(stmt)
logger.debug(f"Table description: {table_description}")
return f"{table}\n{table_description}"
def summarize_table(self, table: str) -> str:
"""Function to compute a number of aggregates over a table.
The function launches a query that computes a number of aggregates over all columns,
including min, max, avg, std and approx_unique.
:param table: Table to summarize
:return: Summary of the table
"""
stmt = f"""WITH column_stats AS (
SELECT
column_name,
data_type
FROM
information_schema.columns
WHERE
table_name = '{table}'
AND table_schema = '{self.table_schema}'
)
SELECT
column_name,
data_type,
COUNT(COALESCE(column_name::text, '')) AS non_null_count,
COUNT(*) - COUNT(COALESCE(column_name::text, '')) AS null_count,
SUM(COALESCE(column_name::numeric, 0)) AS sum,
AVG(COALESCE(column_name::numeric, 0)) AS mean,
MIN(column_name::numeric) AS min,
MAX(column_name::numeric) AS max,
STDDEV(COALESCE(column_name::numeric, 0)) AS stddev
FROM
column_stats,
LATERAL (
SELECT
*
FROM
{table}
) AS tbl
WHERE
data_type IN ('integer', 'numeric', 'real', 'double precision')
GROUP BY
column_name, data_type
UNION ALL
SELECT
column_name,
data_type,
COUNT(COALESCE(column_name::text, '')) AS non_null_count,
COUNT(*) - COUNT(COALESCE(column_name::text, '')) AS null_count,
NULL AS sum,
NULL AS mean,
NULL AS min,
NULL AS max,
NULL AS stddev
FROM
column_stats,
LATERAL (
SELECT
*
FROM
{table}
) AS tbl
WHERE
data_type NOT IN ('integer', 'numeric', 'real', 'double precision')
GROUP BY
column_name, data_type;
"""
table_summary = self.run_query(stmt)
logger.debug(f"Table summary: {table_summary}")
return table_summary
def inspect_query(self, query: str) -> str:
"""Function to inspect a query and return the query plan. Always inspect your query before running them.
:param query: Query to inspect
:return: Query plan
"""
stmt = f"EXPLAIN {query};"
explain_plan = self.run_query(stmt)
logger.debug(f"Explain plan: {explain_plan}")
return explain_plan
def export_table_to_path(self, table: str, path: Optional[str] = None) -> str:
"""Save a table in CSV format.
If the path is provided, the table will be saved under that path.
Eg: If path is /tmp, the table will be saved as /tmp/table.csv
Otherwise it will be saved in the current directory
:param table: Table to export
:param path: Path to export to
:return: None
"""
logger.debug(f"Exporting Table {table} as CSV to path {path}")
if path is None:
path = f"{table}.csv"
else:
path = f"{path}/{table}.csv"
export_statement = f"COPY {self.table_schema}.{table} TO '{path}' DELIMITER ',' CSV HEADER;"
result = self.run_query(export_statement)
logger.debug(f"Exported {table} to {path}/{table}")
return result
def run_query(self, query: str) -> str:
"""Function that runs a query and returns the result.
:param query: SQL query to run
:return: Result of the query
"""
# -*- Format the SQL Query
# Remove backticks
formatted_sql = query.replace("`", "")
# If there are multiple statements, only run the first one
formatted_sql = formatted_sql.split(";")[0]
try:
logger.info(f"Running: {formatted_sql}")
cursor = self.connection.cursor()
cursor.execute(query)
query_result = cursor.fetchall()
result_output = "No output"
if query_result is not None:
try:
results_as_python_objects = query_result
result_rows = []
for row in results_as_python_objects:
if len(row) == 1:
result_rows.append(str(row[0]))
else:
result_rows.append(",".join(str(x) for x in row))
result_data = "\n".join(result_rows)
result_output = ",".join(query_result.columns) + "\n" + result_data
except AttributeError:
result_output = str(query_result)
logger.debug(f"Query result: {result_output}")
return result_output
except Exception as e:
return str(e)
|