arcade-postgres 0.2.0__tar.gz → 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/PKG-INFO +1 -1
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/arcade_postgres/tools/postgres.py +10 -8
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/pyproject.toml +1 -1
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/tests/test_postgres.py +8 -4
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/.gitignore +0 -0
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/Makefile +0 -0
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/arcade_postgres/__init__.py +0 -0
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/arcade_postgres/database_engine.py +0 -0
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/arcade_postgres/tools/__init__.py +0 -0
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/evals/eval_postgres.py +0 -0
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/tests/__init__.py +0 -0
- {arcade_postgres-0.2.0 → arcade_postgres-0.3.0}/tests/dump.sql +0 -0
|
@@ -8,19 +8,19 @@ from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
8
8
|
from ..database_engine import MAX_ROWS_RETURNED, DatabaseEngine
|
|
9
9
|
|
|
10
10
|
|
|
11
|
-
@tool(requires_secrets=["
|
|
11
|
+
@tool(requires_secrets=["POSTGRES_DATABASE_CONNECTION_STRING"])
|
|
12
12
|
async def discover_schemas(
|
|
13
13
|
context: ToolContext,
|
|
14
14
|
) -> list[str]:
|
|
15
15
|
"""Discover all the schemas in the postgres database."""
|
|
16
16
|
async with await DatabaseEngine.get_engine(
|
|
17
|
-
context.get_secret("
|
|
17
|
+
context.get_secret("POSTGRES_DATABASE_CONNECTION_STRING")
|
|
18
18
|
) as engine:
|
|
19
19
|
schemas = await _get_schemas(engine)
|
|
20
20
|
return schemas
|
|
21
21
|
|
|
22
22
|
|
|
23
|
-
@tool(requires_secrets=["
|
|
23
|
+
@tool(requires_secrets=["POSTGRES_DATABASE_CONNECTION_STRING"])
|
|
24
24
|
async def discover_tables(
|
|
25
25
|
context: ToolContext,
|
|
26
26
|
schema_name: Annotated[
|
|
@@ -32,13 +32,13 @@ async def discover_tables(
|
|
|
32
32
|
ALWAYS use this tool before any other tool that requires a table name.
|
|
33
33
|
"""
|
|
34
34
|
async with await DatabaseEngine.get_engine(
|
|
35
|
-
context.get_secret("
|
|
35
|
+
context.get_secret("POSTGRES_DATABASE_CONNECTION_STRING")
|
|
36
36
|
) as engine:
|
|
37
37
|
tables = await _get_tables(engine, schema_name)
|
|
38
38
|
return tables
|
|
39
39
|
|
|
40
40
|
|
|
41
|
-
@tool(requires_secrets=["
|
|
41
|
+
@tool(requires_secrets=["POSTGRES_DATABASE_CONNECTION_STRING"])
|
|
42
42
|
async def get_table_schema(
|
|
43
43
|
context: ToolContext,
|
|
44
44
|
schema_name: Annotated[str, "The database schema to get the table schema of"],
|
|
@@ -50,12 +50,12 @@ async def get_table_schema(
|
|
|
50
50
|
This tool should ALWAYS be used before executing any query. All tables in the query must be discovered first using the <DiscoverTables> tool.
|
|
51
51
|
"""
|
|
52
52
|
async with await DatabaseEngine.get_engine(
|
|
53
|
-
context.get_secret("
|
|
53
|
+
context.get_secret("POSTGRES_DATABASE_CONNECTION_STRING")
|
|
54
54
|
) as engine:
|
|
55
55
|
return await _get_table_schema(engine, schema_name, table_name)
|
|
56
56
|
|
|
57
57
|
|
|
58
|
-
@tool(requires_secrets=["
|
|
58
|
+
@tool(requires_secrets=["POSTGRES_DATABASE_CONNECTION_STRING"])
|
|
59
59
|
async def execute_select_query(
|
|
60
60
|
context: ToolContext,
|
|
61
61
|
select_clause: Annotated[
|
|
@@ -116,7 +116,7 @@ async def execute_select_query(
|
|
|
116
116
|
* Only join on columns that are indexed or the primary key. Do not join on arbitrary columns.
|
|
117
117
|
"""
|
|
118
118
|
async with await DatabaseEngine.get_engine(
|
|
119
|
-
context.get_secret("
|
|
119
|
+
context.get_secret("POSTGRES_DATABASE_CONNECTION_STRING")
|
|
120
120
|
) as engine:
|
|
121
121
|
try:
|
|
122
122
|
return await _execute_query(
|
|
@@ -171,6 +171,8 @@ async def _get_tables(engine: AsyncEngine, schema_name: str) -> list[str]:
|
|
|
171
171
|
|
|
172
172
|
these_tables = await conn.run_sync(get_table_names)
|
|
173
173
|
tables.extend(these_tables)
|
|
174
|
+
|
|
175
|
+
tables.sort()
|
|
174
176
|
return tables
|
|
175
177
|
|
|
176
178
|
|
|
@@ -15,7 +15,7 @@ from arcade_tdk.errors import RetryableToolError
|
|
|
15
15
|
from sqlalchemy import text
|
|
16
16
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
17
17
|
|
|
18
|
-
|
|
18
|
+
POSTGRES_DATABASE_CONNECTION_STRING = (
|
|
19
19
|
environ.get("TEST_POSTGRES_DATABASE_CONNECTION_STRING")
|
|
20
20
|
or "postgresql://evan@localhost:5432/postgres"
|
|
21
21
|
)
|
|
@@ -26,7 +26,9 @@ def mock_context():
|
|
|
26
26
|
context = ToolContext()
|
|
27
27
|
context.secrets = []
|
|
28
28
|
context.secrets.append(
|
|
29
|
-
ToolSecretItem(
|
|
29
|
+
ToolSecretItem(
|
|
30
|
+
key="POSTGRES_DATABASE_CONNECTION_STRING", value=POSTGRES_DATABASE_CONNECTION_STRING
|
|
31
|
+
)
|
|
30
32
|
)
|
|
31
33
|
|
|
32
34
|
return context
|
|
@@ -37,7 +39,9 @@ def mock_context():
|
|
|
37
39
|
async def restore_database():
|
|
38
40
|
with open(f"{os.path.dirname(__file__)}/dump.sql") as f:
|
|
39
41
|
engine = create_async_engine(
|
|
40
|
-
|
|
42
|
+
POSTGRES_DATABASE_CONNECTION_STRING.replace("postgresql", "postgresql+asyncpg").split(
|
|
43
|
+
"?"
|
|
44
|
+
)[0]
|
|
41
45
|
)
|
|
42
46
|
async with engine.connect() as c:
|
|
43
47
|
queries = f.read().split(";")
|
|
@@ -64,7 +68,7 @@ async def test_discover_schemas(mock_context) -> None:
|
|
|
64
68
|
|
|
65
69
|
@pytest.mark.asyncio
|
|
66
70
|
async def test_discover_tables(mock_context) -> None:
|
|
67
|
-
assert await discover_tables(mock_context) == ["
|
|
71
|
+
assert await discover_tables(mock_context) == ["messages", "users"]
|
|
68
72
|
|
|
69
73
|
|
|
70
74
|
@pytest.mark.asyncio
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|