arcade-postgres 0.1.0__py3-none-any.whl

arcade_postgres/database_engine.py

@@ -0,0 +1,104 @@
+from typing import Any, ClassVar
+from urllib.parse import urlparse
+
+from arcade_tdk.errors import RetryableToolError
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+
+MAX_ROWS_RETURNED = 1000
+TEST_QUERY = "SELECT 1"
+
+
+class DatabaseEngine:
+    _instance: ClassVar[None] = None
+    _engines: ClassVar[dict[str, AsyncEngine]] = {}
+
+    @classmethod
+    async def get_instance(cls, connection_string: str) -> AsyncEngine:
+        parsed_url = urlparse(connection_string)
+
+        # TODO: something strange with sslmode= and friends
+        # query_params = parse_qs(parsed_url.query)
+        # query_params = {
+        #     k: v[0] for k, v in query_params.items()
+        # }  # assume one value allowed for each query param
+
+        async_connection_string = f"{parsed_url.scheme.replace('postgresql', 'postgresql+asyncpg')}://{parsed_url.netloc}{parsed_url.path}"
+        key = f"{async_connection_string}"
+        if key not in cls._engines:
+            cls._engines[key] = create_async_engine(async_connection_string)
+
+        # try a simple query to see if the connection is valid
+        try:
+            async with cls._engines[key].connect() as connection:
+                await connection.execute(text(TEST_QUERY))
+            return cls._engines[key]
+        except Exception:
+            await cls._engines[key].dispose()
+
+            # try again
+            try:
+                async with cls._engines[key].connect() as connection:
+                    await connection.execute(text(TEST_QUERY))
+                return cls._engines[key]
+            except Exception as e:
+                raise RetryableToolError(
+                    f"Connection failed: {e}",
+                    developer_message="Connection to postgres failed.",
+                    additional_prompt_content="Check the connection string and try again.",
+                ) from e
+
+    @classmethod
+    async def get_engine(cls, connection_string: str) -> Any:
+        engine = await cls.get_instance(connection_string)
+
+        class ConnectionContextManager:
+            def __init__(self, engine: AsyncEngine) -> None:
+                self.engine = engine
+
+            async def __aenter__(self) -> AsyncEngine:
+                return self.engine
+
+            async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+                # Connection cleanup is handled by the async context manager
+                pass
+
+        return ConnectionContextManager(engine)
+
+    @classmethod
+    async def cleanup(cls) -> None:
+        """Clean up all cached engines. Call this when shutting down."""
+        for engine in cls._engines.values():
+            await engine.dispose()
+        cls._engines.clear()
+
+    @classmethod
+    def clear_cache(cls) -> None:
+        """Clear the engine cache without disposing engines. Use with caution."""
+        cls._engines.clear()
+
+    @classmethod
+    def sanitize_query(cls, query: str) -> str:
+        """
+        Sanitize a query to not break our read-only session.
+        THIS IS REALLY UNSAFE AND SHOULD NOT BE USED IN PRODUCTION. USE A DATABASE CONNECTION WITH A READ-ONLY USER AND PREPARE STATEMENTS.
+        There are also valid reasons for the ";" character, and this prevents that.
+        """
+
+        parts = query.split(";")
+        if len(parts) > 1:
+            raise RetryableToolError(
+                "Multiple statements are not allowed in a single query.",
+                developer_message="Multiple statements are not allowed in a single query.",
+                additional_prompt_content="Split your query into multiple queries and try again.",
+            )
+
+        words = parts[0].split(" ")
+        if words[0].upper().strip() != "SELECT":
+            raise RetryableToolError(
+                "Only SELECT queries are allowed.",
+                developer_message="Only SELECT queries are allowed.",
+                additional_prompt_content="Use the <DiscoverTables> and <GetTableSchema> tools to discover the tables and try again.",
+            )
+
+        return f"{query}"

arcade_postgres/tools/postgres.py

@@ -0,0 +1,178 @@
+from typing import Annotated, Any
+
+from arcade_tdk import ToolContext, tool
+from arcade_tdk.errors import RetryableToolError
+from sqlalchemy import inspect, text
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from ..database_engine import MAX_ROWS_RETURNED, DatabaseEngine
+
+
+@tool(requires_secrets=["DATABASE_CONNECTION_STRING"])
+async def discover_schemas(
+    context: ToolContext,
+) -> list[str]:
+    """Discover all the schemas in the postgres database."""
+    async with await DatabaseEngine.get_engine(
+        context.get_secret("DATABASE_CONNECTION_STRING")
+    ) as engine:
+        schemas = await _get_schemas(engine)
+        return schemas
+
+
+@tool(requires_secrets=["DATABASE_CONNECTION_STRING"])
+async def discover_tables(
+    context: ToolContext,
+    schema_name: Annotated[
+        str, "The database schema to discover tables in (default value: 'public')"
+    ] = "public",
+) -> list[str]:
+    """Discover all the tables in the postgres database when the list of tables is not known.
+
+    THIS TOOL SHOULD ALWAYS BE USED BEFORE ANY OTHER TOOL THAT REQUIRES A TABLE NAME.
+    """
+    async with await DatabaseEngine.get_engine(
+        context.get_secret("DATABASE_CONNECTION_STRING")
+    ) as engine:
+        tables = await _get_tables(engine, schema_name)
+        return tables
+
+
+@tool(requires_secrets=["DATABASE_CONNECTION_STRING"])
+async def get_table_schema(
+    context: ToolContext,
+    schema_name: Annotated[str, "The database schema to get the table schema of"],
+    table_name: Annotated[str, "The table to get the schema of"],
+) -> list[str]:
+    """
+    Get the schema/structure of a postgres table in the postgres database when the schema is not known, and the name of the table is provided.
+
+    THIS TOOL SHOULD ALWAYS BE USED BEFORE EXECUTING ANY QUERY.  ALL TABLES IN THE QUERY MUST BE DISCOVERED FIRST USING THE <DiscoverTables> TOOL.
+    """
+    async with await DatabaseEngine.get_engine(
+        context.get_secret("DATABASE_CONNECTION_STRING")
+    ) as engine:
+        return await _get_table_schema(engine, schema_name, table_name)
+
+
+@tool(requires_secrets=["DATABASE_CONNECTION_STRING"])
+async def execute_query(
+    context: ToolContext,
+    query: Annotated[str, "The postgres SQL query to execute.  Only SELECT queries are allowed."],
+) -> list[str]:
+    """
+    You have a connection to a postgres database.
+    Execute a query and return the results against the postgres database.
+
+    ONLY USE THIS TOOL IF YOU HAVE ALREADY LOADED THE SCHEMA OF THE TABLES YOU NEED TO QUERY.  USE THE <GetTableSchema> TOOL TO LOAD THE SCHEMA IF NOT ALREADY KNOWN.
+
+    When running queries, follow these rules which will help avoid errors:
+    * Always use case-insensitive queries to match strings in the query.
+    * Always trim strings in the query.
+    * Prefer LIKE queries over direct string matches or regex queries.
+    * Only join on columns that are indexed or the primary key.  Do not join on arbitrary columns.
+
+    Only SELECT queries are allowed.  Do not use INSERT, UPDATE, DELETE, or other DML statements.  This tool will reject them.
+
+    Unless otherwise specified, ensure that query has a LIMIT of 100 for all results.  This tool will enforce that no more than 1000 rows are returned at maximum.
+    """
+    async with await DatabaseEngine.get_engine(
+        context.get_secret("DATABASE_CONNECTION_STRING")
+    ) as engine:
+        try:
+            return await _execute_query(engine, query)
+        except Exception as e:
+            raise RetryableToolError(
+                f"Query failed: {e}",
+                developer_message=f"Query '{query}' failed.",
+                additional_prompt_content="Load the database schema <GetTableSchema> or use the <DiscoverTables> tool to discover the tables and try again.",
+                retry_after_ms=10,
+            ) from e
+
+
+async def _get_schemas(engine: AsyncEngine) -> list[str]:
+    """Get all the schemas in the database"""
+    async with engine.connect() as conn:
+
+        def get_schema_names(sync_conn: Any) -> list[str]:
+            return list(inspect(sync_conn).get_schema_names())
+
+        schemas: list[str] = await conn.run_sync(get_schema_names)
+        schemas = [schema for schema in schemas if schema != "information_schema"]
+
+        return schemas
+
+
+async def _get_tables(engine: AsyncEngine, schema_name: str) -> list[str]:
+    """Get all the tables in the database"""
+    async with engine.connect() as conn:
+
+        def get_schema_names(sync_conn: Any) -> list[str]:
+            return list(inspect(sync_conn).get_schema_names())
+
+        schemas: list[str] = await conn.run_sync(get_schema_names)
+        tables = []
+        for schema in schemas:
+            if schema == schema_name:
+
+                def get_table_names(sync_conn: Any, s: str = schema) -> list[str]:
+                    return list(inspect(sync_conn).get_table_names(schema=s))
+
+                these_tables = await conn.run_sync(get_table_names)
+                tables.extend(these_tables)
+        return tables
+
+
+async def _get_table_schema(engine: AsyncEngine, schema_name: str, table_name: str) -> list[str]:
+    """Get the schema of a table"""
+    async with engine.connect() as connection:
+
+        def get_columns(sync_conn: Any, t: str = table_name, s: str = schema_name) -> list[Any]:
+            return list(inspect(sync_conn).get_columns(t, s))
+
+        columns_table = await connection.run_sync(get_columns)
+
+        # Get primary key information
+        pk_constraint = await connection.run_sync(
+            lambda sync_conn: inspect(sync_conn).get_pk_constraint(table_name, schema_name)
+        )
+        primary_keys = set(pk_constraint.get("constrained_columns", []))
+
+        # Get index information
+        indexes = await connection.run_sync(
+            lambda sync_conn: inspect(sync_conn).get_indexes(table_name, schema_name)
+        )
+        indexed_columns = set()
+        for index in indexes:
+            indexed_columns.update(index.get("column_names", []))
+
+        results = []
+        for column in columns_table:
+            column_name = column["name"]
+            column_type = column["type"].python_type.__name__
+
+            # Build column description
+            description = f"{column_name}: {column_type}"
+
+            # Add primary key indicator
+            if column_name in primary_keys:
+                description += " (PRIMARY KEY)"
+
+            # Add index indicator
+            if column_name in indexed_columns:
+                description += " (INDEXED)"
+
+            results.append(description)
+
+        return results[:MAX_ROWS_RETURNED]
+
+
+async def _execute_query(
+    engine: AsyncEngine, query: str, params: dict[str, Any] | None = None
+) -> list[str]:
+    """Execute a query and return the results."""
+    async with engine.connect() as connection:
+        result = await connection.execute(text(DatabaseEngine.sanitize_query(query)), params)
+        rows = result.fetchall()
+        results = [str(row) for row in rows]
+        return results[:MAX_ROWS_RETURNED]

arcade_postgres-0.1.0.dist-info/METADATA

@@ -0,0 +1,23 @@
+Metadata-Version: 2.4
+Name: arcade_postgres
+Version: 0.1.0
+Summary: Tools to query and explore a postgres database
+Author-email: evantahler <support@arcade.dev>
+Requires-Python: >=3.10
+Requires-Dist: arcade-tdk<3.0.0,>=2.0.0
+Requires-Dist: asyncpg>=0.30.0
+Requires-Dist: greenlet>=3.2.3
+Requires-Dist: psycopg2-binary>=2.9.10
+Requires-Dist: pydantic>=2.11.7
+Requires-Dist: sqlalchemy>=2.0.41
+Provides-Extra: dev
+Requires-Dist: arcade-ai[evals]<3.0.0,>=2.0.0; extra == 'dev'
+Requires-Dist: arcade-serve<3.0.0,>=2.0.0; extra == 'dev'
+Requires-Dist: mypy<1.6.0,>=1.5.1; extra == 'dev'
+Requires-Dist: pre-commit<3.5.0,>=3.4.0; extra == 'dev'
+Requires-Dist: pytest-asyncio<0.25.0,>=0.24.0; extra == 'dev'
+Requires-Dist: pytest-cov<4.1.0,>=4.0.0; extra == 'dev'
+Requires-Dist: pytest-mock<3.12.0,>=3.11.1; extra == 'dev'
+Requires-Dist: pytest<8.4.0,>=8.3.0; extra == 'dev'
+Requires-Dist: ruff<0.8.0,>=0.7.4; extra == 'dev'
+Requires-Dist: tox<4.12.0,>=4.11.1; extra == 'dev'

arcade_postgres-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+arcade_postgres/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arcade_postgres/database_engine.py,sha256=-HDhOk2Zgb-4F1QeYRCZdrmvv2JkaGcXFawE9tjn2nM,4088
+arcade_postgres/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arcade_postgres/tools/postgres.py,sha256=E4Z8Be08UbWI08iJynfS3VTM3vP7Cca9Ymz1Op7e_N4,7020
+arcade_postgres-0.1.0.dist-info/METADATA,sha256=mWL4dI8vwmCvO0UvLo7EcPwEdeTbRFAHqNBofR5VmZY,968
+arcade_postgres-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+arcade_postgres-0.1.0.dist-info/RECORD,,

arcade_postgres-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