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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: arcade_postgres
3
- Version: 0.2.0
3
+ Version: 0.3.0
4
4
  Summary: Tools to query and explore a postgres database
5
5
  Author-email: evantahler <support@arcade.dev>
6
6
  Requires-Python: >=3.10
@@ -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=["DATABASE_CONNECTION_STRING"])
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("DATABASE_CONNECTION_STRING")
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=["DATABASE_CONNECTION_STRING"])
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("DATABASE_CONNECTION_STRING")
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=["DATABASE_CONNECTION_STRING"])
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("DATABASE_CONNECTION_STRING")
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=["DATABASE_CONNECTION_STRING"])
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("DATABASE_CONNECTION_STRING")
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
 
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "arcade_postgres"
7
- version = "0.2.0"
7
+ version = "0.3.0"
8
8
  description = "Tools to query and explore a postgres database"
9
9
  requires-python = ">=3.10"
10
10
  dependencies = [
@@ -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
- DATABASE_CONNECTION_STRING = (
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(key="DATABASE_CONNECTION_STRING", value=DATABASE_CONNECTION_STRING)
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
- DATABASE_CONNECTION_STRING.replace("postgresql", "postgresql+asyncpg").split("?")[0]
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) == ["users", "messages"]
71
+ assert await discover_tables(mock_context) == ["messages", "users"]
68
72
 
69
73
 
70
74
  @pytest.mark.asyncio