pgsql-mcp-server 0.1.0__py3-none-any.whl

pgsql_mcp_server/app.py

@@ -0,0 +1,147 @@
+import os
+from contextlib import asynccontextmanager
+from textwrap import dedent
+from typing import AsyncIterator, Optional
+from dataclasses import dataclass
+import click
+from sqlmodel import create_engine, Session, text
+from sqlalchemy.engine import Engine
+from mcp.server.fastmcp import FastMCP, Context
+from sqlalchemy import inspect, Result
+
+
+@dataclass
+class AppContext:
+    engine: Engine
+
+
+@asynccontextmanager
+async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
+    _dsn = os.getenv("DATABASE_URL")
+    engine = create_engine(_dsn, echo=False)
+
+    try:
+        yield AppContext(engine=engine)
+    finally:
+        engine.dispose()
+
+
+mcp = FastMCP("PostgreSQL Explorer", lifespan=app_lifespan)
+
+
+@mcp.tool()
+def get_schema_names(ctx: Context) -> str:
+    """Get all schema names."""
+    engine = ctx.request_context.lifespan_context.engine
+    inspector = inspect(engine)
+
+    schema_names = inspector.get_schema_names()
+    if not inspector.get_schema_names():
+        return "No schemas found."
+    else:
+        return f"All schemas: {', '.join(schema_names)}"
+
+
+@mcp.tool()
+def get_tables(ctx: Context, schema_name: Optional[str] = "public") -> str:
+    """Get all tables in a schema.
+
+    Args:
+        schema_name: The name of the schema to get tables from, defaults to "public".
+    """
+    try:
+        schema_name = schema_name or "public"
+        engine = ctx.request_context.lifespan_context.engine
+        inspector = inspect(engine)
+
+        table_names = inspector.get_table_names(schema=schema_name)
+        if table_names:
+            return f"All tables from schema {schema_name}: {', '.join(table_names)}"
+        else:
+            return f"No tables found in schema {schema_name}."
+
+    except Exception as e:
+        return f"查询表时发生错误: {str(e)}"
+
+
+@mcp.tool()
+def get_columns(ctx: Context, table: str, schema_name: Optional[str] = "public") -> str:
+    """Get all columns in a table.
+
+    Args:
+        table: The name of the table to get columns from.
+        schema_name: The name of the schema to get columns from, defaults to "public".
+    """
+    try:
+        schema_name = schema_name or "public"
+        engine = ctx.request_context.lifespan_context.engine
+        inspector = inspect(engine)
+        columns = inspector.get_columns(table, schema=schema_name)
+
+        if not columns:
+            return "No columns found in table."
+        else:
+            result = []
+            for column in columns:
+                col = dedent(f"""\
+                column_name: {column["name"]}
+                type: {column["type"]}
+                nullable: {column["nullable"]}
+                default: {column.get("default")}
+                autoincrement: {column.get("autoincrement")} 
+                comment: {column.get("comment")} \
+                """)
+                result.append(col)
+            return "\n\n".join(result)
+
+    except Exception as e:
+        return f"查询表时发生错误: {str(e)}"
+
+
+@mcp.tool()
+def get_indexes(ctx: Context, table: str, schema_name: Optional[str] = "public") -> str:
+    """Get all indexes in a table.
+
+    Args:
+        table: The name of the table to get indexes from.
+        schema_name: The name of the schema to get indexes from, defaults to "public".
+    """
+    try:
+        schema_name = schema_name or "public"
+        engine = ctx.request_context.lifespan_context.engine
+        inspector = inspect(engine)
+        indexes = inspector.get_indexes(table, schema=schema_name)
+
+        if not indexes:
+            return "No indexes found in table."
+        else:
+            return str(indexes)
+
+    except Exception as e:
+        return f"查询表时发生错误: {str(e)}"
+
+
+@mcp.tool()
+def run_dql_query(ctx: Context, raw_sql_query: str) -> str:
+    """Run a raw DQL SQL query, like SELECT, SHOW, DESCRIBE, EXPLAIN, etc.
+
+    Args:
+        raw_sql_query: The raw SQL query to run.
+    """
+    engine = ctx.request_context.lifespan_context.engine
+    with Session(engine) as session:
+        raw_sql_select = text(raw_sql_query)
+        result: Result = session.execute(raw_sql_select)
+        rows = result.fetchall()
+        return str(rows)
+
+
+@click.command()
+@click.option("--dsn", "-d", type=str, help="Database connection string")
+@click.option("-v", "--verbose", count=True)
+def serve(dsn: str, verbose: bool):
+    os.environ["DATABASE_URL"] = dsn
+    mcp.run()
+
+if __name__ == "__main__":
+    serve()

pgsql_mcp_server-0.1.0.dist-info/METADATA

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: pgsql-mcp-server
+Version: 0.1.0
+Summary: PostgreSQL MCP Server
+Requires-Python: >=3.12
+Requires-Dist: click>=8.1.8
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: mcp[cli]>=1.6.0
+Requires-Dist: psycopg2-binary>=2.9.10
+Requires-Dist: python-dotenv>=1.1.0
+Requires-Dist: sqlmodel>=0.0.24

pgsql_mcp_server-0.1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+pgsql_mcp_server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pgsql_mcp_server/app.py,sha256=I6YOGNyLebNnm424lbDD0eO3iGBV6pi03U2BNC40kKQ,4437
+pgsql_mcp_server-0.1.0.dist-info/METADATA,sha256=msd5Alk8xe7ru4j14T7GWvEaAMyP8N2_xUjTQlFfQic,310
+pgsql_mcp_server-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pgsql_mcp_server-0.1.0.dist-info/entry_points.txt,sha256=Z3-MsUmk3TG-aQ6uTFcc7jLCjAa38msknabm0kF4axI,64
+pgsql_mcp_server-0.1.0.dist-info/RECORD,,

pgsql_mcp_server-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

pgsql_mcp_server-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+pgsql-mcp-server = pgsql_mcp_server.app:serve