mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-16 17:00:12 +00:00
292 lines
12 KiB
Python
292 lines
12 KiB
Python
from sqlalchemy import text
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
|
|
from private_gpt.components.database.connection_factory import DatabaseDialect
|
|
from private_gpt.components.database.inspected_schema import (
|
|
InspectedFunction,
|
|
InspectedFunctionParams,
|
|
)
|
|
from private_gpt.components.database.inspector_interface import (
|
|
DatabaseObjectInspector,
|
|
DatabaseObjectType,
|
|
)
|
|
|
|
|
|
class DatabaseFunctionsInspector(DatabaseObjectInspector):
|
|
def get_inspector_type(self) -> str:
|
|
return DatabaseObjectType.FUNCTION
|
|
|
|
def get_objects(self, schema: str) -> list[InspectedFunction]: # type: ignore[override]
|
|
try:
|
|
if self.dialect == DatabaseDialect.MSSQL:
|
|
return self._get_sqlserver_functions(schema)
|
|
if self.dialect == DatabaseDialect.DB2:
|
|
return self._get_db2_functions(schema)
|
|
|
|
# TODO: Postgres support is not complete
|
|
# elif self.dialect == DatabaseDialect.POSTGRES:
|
|
# return self._get_postgresql_functions(schema)
|
|
else:
|
|
return []
|
|
except SQLAlchemyError:
|
|
return []
|
|
|
|
def _get_postgresql_functions(self, schema: str) -> list[InspectedFunction]:
|
|
# TODO: review this function (it is not finished)
|
|
self._ensure_connected()
|
|
schema_filter = f"AND n.nspname = '{schema}'" if schema else ""
|
|
if self._is_readonly:
|
|
schema_filter += " AND n.nspname NOT IN ('pg_temp_1', 'pg_toast_temp_1')"
|
|
schema_filter += " AND UPPER(pg_get_functiondef(p.oid)) NOT LIKE '%DELETE%'"
|
|
schema_filter += " AND UPPER(pg_get_functiondef(p.oid)) NOT LIKE '%UPDATE%'"
|
|
schema_filter += " AND UPPER(pg_get_functiondef(p.oid)) NOT LIKE '%INSERT%'"
|
|
schema_filter += " AND UPPER(pg_get_functiondef(p.oid)) NOT LIKE '%DROP%'"
|
|
schema_filter += " AND UPPER(pg_get_functiondef(p.oid)) NOT LIKE '%ALTER%'"
|
|
|
|
types = "('f')" # To include only functions
|
|
|
|
query = text(
|
|
f"""
|
|
SELECT
|
|
n.nspname as schema_name,
|
|
p.proname as procedure_name,
|
|
pg_get_function_result(p.oid) as return_type,
|
|
params.param_name,
|
|
format_type(params.param_type, NULL) as param_type_name,
|
|
params.param_position
|
|
FROM pg_proc p
|
|
JOIN pg_namespace n ON p.pronamespace = n.oid
|
|
CROSS JOIN LATERAL (
|
|
SELECT *
|
|
FROM ROWS FROM (
|
|
unnest(COALESCE(p.proargnames, ARRAY_FILL(NULL::text, ARRAY[p.pronargs]))),
|
|
unnest(COALESCE(p.proallargtypes, p.proargtypes::oid[])),
|
|
generate_series(1, GREATEST(
|
|
array_length(COALESCE(p.proallargtypes, p.proargtypes::oid[]), 1),
|
|
p.pronargs
|
|
))
|
|
) AS t(param_name, param_type, param_position)
|
|
WHERE t.param_type IS NOT NULL
|
|
) params
|
|
WHERE p.prokind IN {types}
|
|
{schema_filter}
|
|
AND n.nspname NOT IN ('pg_catalog', 'information_schema')
|
|
ORDER BY n.nspname, p.proname, params.param_position
|
|
"""
|
|
)
|
|
|
|
functions = []
|
|
try:
|
|
if not self._engine:
|
|
return []
|
|
conn = self._engine.connect()
|
|
result = conn.execute(query)
|
|
current_function = None
|
|
|
|
for row in result:
|
|
schema, name, return_type, param_name, param_type, param_position = (
|
|
row[0],
|
|
row[1],
|
|
row[2],
|
|
row[3],
|
|
row[4],
|
|
row[5],
|
|
)
|
|
if current_function is None or (
|
|
current_function.schema != schema or current_function.name != name
|
|
):
|
|
if current_function is not None:
|
|
functions.append(current_function)
|
|
current_function = InspectedFunction()
|
|
current_function.schema = schema
|
|
current_function.name = name
|
|
current_function.return_types = [return_type]
|
|
current_function.parameters = []
|
|
if (
|
|
(param_name or param_position)
|
|
and current_function
|
|
and current_function.parameters
|
|
):
|
|
parameter = InspectedFunctionParams()
|
|
parameter.name = (
|
|
param_name if param_name else f"arg{param_position}"
|
|
)
|
|
parameter.data_type = param_type
|
|
current_function.parameters.append(parameter)
|
|
|
|
if current_function is not None:
|
|
functions.append(current_function)
|
|
except Exception:
|
|
return []
|
|
return functions
|
|
|
|
def _get_sqlserver_functions(self, schema: str) -> list[InspectedFunction]:
|
|
|
|
schema_filter = f"AND s.name = '{schema}'" if schema else ""
|
|
if self._is_readonly:
|
|
schema_filter += "AND UPPER(m.definition) NOT LIKE '%DELETE%'"
|
|
schema_filter += "AND UPPER(m.definition) NOT LIKE '%UPDATE%'"
|
|
schema_filter += "AND UPPER(m.definition) NOT LIKE '%INSERT%'"
|
|
schema_filter += "AND UPPER(m.definition) NOT LIKE '%DROP%'"
|
|
schema_filter += "AND UPPER(m.definition) NOT LIKE '%ALTER%'"
|
|
types = "('FN')" # Only stored functions
|
|
query = text(
|
|
f"""
|
|
SELECT
|
|
s.name AS schema_name,
|
|
o.name AS object_name,
|
|
CAST(ep_proc.value AS NVARCHAR(MAX)) AS comment,
|
|
prm.name AS parameter_name,
|
|
t.name AS parameter_type,
|
|
CAST(ep_param.value AS NVARCHAR(MAX)) AS parameter_comment
|
|
FROM sys.objects o
|
|
JOIN sys.schemas s ON o.schema_id = s.schema_id
|
|
LEFT JOIN sys.parameters prm ON o.object_id = prm.object_id
|
|
LEFT JOIN sys.types t ON prm.user_type_id = t.user_type_id
|
|
JOIN sys.sql_modules m ON o.object_id = m.object_id
|
|
LEFT JOIN sys.extended_properties ep_proc ON
|
|
ep_proc.major_id = o.object_id
|
|
AND ep_proc.minor_id = 0
|
|
AND ep_proc.name = 'MS_Description'
|
|
AND ep_proc.class = 1
|
|
LEFT JOIN sys.extended_properties ep_param ON
|
|
ep_param.major_id = o.object_id
|
|
AND ep_param.minor_id = prm.parameter_id
|
|
AND ep_param.name = 'MS_Description'
|
|
AND ep_param.class = 2
|
|
WHERE o.type IN {types}
|
|
{schema_filter}
|
|
ORDER BY s.name, o.name;
|
|
"""
|
|
)
|
|
|
|
functions = []
|
|
try:
|
|
self._ensure_connected()
|
|
if not self._engine:
|
|
return []
|
|
conn = self._engine.connect()
|
|
result = conn.execute(query)
|
|
current_function = None
|
|
|
|
for row in result:
|
|
schema, name, comment, param_name, param_type, parameter_comment = (
|
|
row[0],
|
|
row[1],
|
|
row[2],
|
|
row[3],
|
|
row[4],
|
|
row[5],
|
|
)
|
|
if current_function is None or (
|
|
current_function.schema != schema or current_function.name != name
|
|
):
|
|
if current_function is not None:
|
|
functions.append(current_function)
|
|
current_function = InspectedFunction()
|
|
current_function.schema = schema
|
|
current_function.name = name
|
|
current_function.comment = comment
|
|
current_function.parameters = []
|
|
current_function.return_types = []
|
|
if param_name:
|
|
parameter = InspectedFunctionParams()
|
|
parameter.name = param_name
|
|
parameter.data_type = param_type
|
|
parameter.comment = parameter_comment
|
|
if not current_function.parameters:
|
|
current_function.parameters = []
|
|
current_function.parameters.append(parameter)
|
|
elif param_type:
|
|
if not current_function.return_types:
|
|
current_function.return_types = []
|
|
current_function.return_types.append(param_type)
|
|
|
|
if current_function is not None:
|
|
functions.append(current_function)
|
|
except Exception:
|
|
return []
|
|
return functions
|
|
|
|
def _get_db2_functions(self, schema: str) -> list[InspectedFunction]:
|
|
schema_filter = (
|
|
f"AND LOWER(r.routineschema) = LOWER('{schema}')" if schema else ""
|
|
)
|
|
|
|
if self._is_readonly:
|
|
schema_filter += " AND UPPER(r.text) NOT LIKE '%DELETE%'"
|
|
schema_filter += " AND UPPER(r.text) NOT LIKE '%UPDATE%'"
|
|
schema_filter += " AND UPPER(r.text) NOT LIKE '%INSERT%'"
|
|
schema_filter += " AND UPPER(r.text) NOT LIKE '%DROP%'"
|
|
schema_filter += " AND UPPER(r.text) NOT LIKE '%ALTER%'"
|
|
|
|
# Only stored functions (F = Function)
|
|
query = text(
|
|
f"""
|
|
SELECT
|
|
r.routineschema AS schema_name,
|
|
r.routinename AS object_name,
|
|
r.remarks AS comment,
|
|
p.parmname AS parameter_name,
|
|
p.typename AS parameter_type,
|
|
p.rowtype AS param_direction
|
|
FROM syscat.routines r
|
|
LEFT JOIN syscat.routineparms p ON
|
|
r.routineschema = p.routineschema
|
|
AND r.specificname = p.specificname
|
|
WHERE r.routinetype = 'F'
|
|
AND r.routineschema NOT LIKE 'SYS%'
|
|
{schema_filter}
|
|
ORDER BY r.routineschema, r.routinename, p.ordinal
|
|
"""
|
|
)
|
|
|
|
functions = []
|
|
try:
|
|
self._ensure_connected()
|
|
if not self._engine:
|
|
return []
|
|
conn = self._engine.connect()
|
|
result = conn.execute(query)
|
|
current_function = None
|
|
|
|
for row in result:
|
|
(schema, name, comment, param_name, param_type, param_direction,) = (
|
|
row[0],
|
|
row[1],
|
|
row[2],
|
|
row[3],
|
|
row[4],
|
|
row[5],
|
|
)
|
|
|
|
if current_function is None or (
|
|
current_function.schema != schema or current_function.name != name
|
|
):
|
|
if current_function is not None:
|
|
functions.append(current_function)
|
|
current_function = InspectedFunction()
|
|
current_function.schema = schema
|
|
current_function.name = name
|
|
current_function.comment = comment
|
|
current_function.parameters = []
|
|
current_function.return_types = []
|
|
|
|
if param_name and param_direction in ("P", "O", "B"):
|
|
# P = IN, O = OUT, B = INOUT
|
|
parameter = InspectedFunctionParams()
|
|
parameter.name = param_name
|
|
parameter.data_type = param_type
|
|
parameter.comment = None
|
|
current_function.parameters.append(parameter)
|
|
elif param_direction == "R" and param_type:
|
|
current_function.return_types.append(
|
|
f"{param_name if param_name else''} {param_type}"
|
|
)
|
|
|
|
if current_function is not None:
|
|
functions.append(current_function)
|
|
except Exception:
|
|
return []
|
|
return functions
|