refactor(auth): simplify flow retrieval by removing settings_service dependency (#8956)
* refactor: Simplify flow reading logic by removing unnecessary settings service dependency * 🛠️ (flows.py): Remove settings_service parameter from _read_flow and related functions to streamline flow retrieval logic. * 🔧 (flows.py): Adjust query to directly filter by user_id, enhancing clarity and reducing complexity. * test: Enhance flow retrieval tests with user isolation checks * 🧪 (test_flows.py): Remove hardcoded user_id from flow creation tests to ensure user-specific data handling. * ✨ (test_flows.py): Add a new test to verify that users can only access their own flows, ensuring proper user isolation in flow retrieval. * 🔧 (test_flows.py): Implement user creation and cleanup logic to maintain test integrity and avoid side effects.
This commit is contained in:
Gabriel Luiz Freitas Almeida
GitHub
parent
73c1f203b0
commit
d437d018ce
2 changed files with 137 additions and 17 deletions
|
|
@ -36,7 +36,6 @@ from langflow.services.database.models.flow.utils import get_webhook_component_i
|
|||
from langflow.services.database.models.folder.constants import DEFAULT_FOLDER_NAME
|
||||
from langflow.services.database.models.folder.model import Folder
|
||||
from langflow.services.deps import get_settings_service
|
||||
from langflow.services.settings.service import SettingsService
|
||||
from langflow.utils.compression import compress_response
|
||||
|
||||
# build router
|
||||
|
|
@ -270,17 +269,10 @@ async def _read_flow(
|
|||
session: AsyncSession,
|
||||
flow_id: UUID,
|
||||
user_id: UUID,
|
||||
settings_service: SettingsService,
|
||||
):
|
||||
"""Read a flow."""
|
||||
auth_settings = settings_service.auth_settings
|
||||
stmt = select(Flow).where(Flow.id == flow_id)
|
||||
if auth_settings.AUTO_LOGIN:
|
||||
# If auto login is enable user_id can be current_user.id or None
|
||||
# so write an OR
|
||||
stmt = stmt.where(
|
||||
(Flow.user_id == user_id) | (Flow.user_id == None) # noqa: E711
|
||||
)
|
||||
stmt = select(Flow).where(Flow.id == flow_id).where(Flow.user_id == user_id)
|
||||
|
||||
return (await session.exec(stmt)).first()
|
||||
|
||||
|
||||
|
|
@ -292,7 +284,7 @@ async def read_flow(
|
|||
current_user: CurrentActiveUser,
|
||||
):
|
||||
"""Read a flow."""
|
||||
if user_flow := await _read_flow(session, flow_id, current_user.id, get_settings_service()):
|
||||
if user_flow := await _read_flow(session, flow_id, current_user.id):
|
||||
return user_flow
|
||||
raise HTTPException(status_code=404, detail="Flow not found")
|
||||
|
||||
|
|
@ -327,7 +319,6 @@ async def update_flow(
|
|||
session=session,
|
||||
flow_id=flow_id,
|
||||
user_id=current_user.id,
|
||||
settings_service=settings_service,
|
||||
)
|
||||
|
||||
if not db_flow:
|
||||
|
|
@ -393,7 +384,6 @@ async def delete_flow(
|
|||
session=session,
|
||||
flow_id=flow_id,
|
||||
user_id=current_user.id,
|
||||
settings_service=get_settings_service(),
|
||||
)
|
||||
if not flow:
|
||||
raise HTTPException(status_code=404, detail="Flow not found")
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@ async def test_create_flow(client: AsyncClient, logged_in_headers):
|
|||
"webhook": False,
|
||||
"endpoint_name": "string",
|
||||
"tags": ["string"],
|
||||
"user_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
"folder_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
"fs_path": str(flow_file),
|
||||
}
|
||||
|
|
@ -79,7 +78,6 @@ async def test_read_flow(client: AsyncClient, logged_in_headers):
|
|||
"webhook": False,
|
||||
"endpoint_name": "string",
|
||||
"tags": ["string"],
|
||||
"user_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
"folder_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
}
|
||||
response_ = await client.post("api/v1/flows/", json=basic_case, headers=logged_in_headers)
|
||||
|
|
@ -118,7 +116,6 @@ async def test_update_flow(client: AsyncClient, logged_in_headers):
|
|||
"webhook": False,
|
||||
"endpoint_name": "string",
|
||||
"tags": ["string"],
|
||||
"user_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
"folder_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
}
|
||||
basic_case["name"] = name
|
||||
|
|
@ -167,7 +164,6 @@ async def test_create_flows(client: AsyncClient, logged_in_headers):
|
|||
"is_component": False,
|
||||
"webhook": False,
|
||||
"tags": ["string"],
|
||||
"user_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
"folder_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
}
|
||||
cases = []
|
||||
|
|
@ -192,3 +188,137 @@ async def test_read_basic_examples(client: AsyncClient, logged_in_headers):
|
|||
assert response.status_code == status.HTTP_200_OK
|
||||
assert isinstance(result, list), "The result must be a list"
|
||||
assert len(result) > 0, "The result must have at least one flow"
|
||||
|
||||
|
||||
async def test_read_flows_user_isolation(client: AsyncClient, logged_in_headers, active_user):
|
||||
"""Test that read_flows returns only flows from the current user."""
|
||||
from uuid import uuid4
|
||||
|
||||
from langflow.services.auth.utils import get_password_hash
|
||||
from langflow.services.database.models.user.model import User
|
||||
from langflow.services.deps import session_scope
|
||||
|
||||
# Create a second user
|
||||
other_user_id = uuid4()
|
||||
async with session_scope() as session:
|
||||
other_user = User(
|
||||
id=other_user_id,
|
||||
username="other_test_user",
|
||||
password=get_password_hash("testpassword"),
|
||||
is_active=True,
|
||||
is_superuser=False,
|
||||
)
|
||||
session.add(other_user)
|
||||
await session.commit()
|
||||
await session.refresh(other_user)
|
||||
|
||||
# Login as the other user to get headers
|
||||
login_data = {"username": "other_test_user", "password": "testpassword"}
|
||||
response = await client.post("api/v1/login", data=login_data)
|
||||
assert response.status_code == 200
|
||||
tokens = response.json()
|
||||
other_user_headers = {"Authorization": f"Bearer {tokens['access_token']}"}
|
||||
|
||||
# Create flows for the first user (active_user)
|
||||
flow_user1_1 = {
|
||||
"name": "user1_flow_1",
|
||||
"description": "Flow 1 for user 1",
|
||||
"icon": "string",
|
||||
"icon_bg_color": "#ff00ff",
|
||||
"gradient": "string",
|
||||
"data": {},
|
||||
"is_component": False,
|
||||
"webhook": False,
|
||||
"endpoint_name": "user1_flow_1_endpoint",
|
||||
"tags": ["user1"],
|
||||
"folder_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
}
|
||||
|
||||
flow_user1_2 = {
|
||||
"name": "user1_flow_2",
|
||||
"description": "Flow 2 for user 1",
|
||||
"icon": "string",
|
||||
"icon_bg_color": "#00ff00",
|
||||
"gradient": "string",
|
||||
"data": {},
|
||||
"is_component": False,
|
||||
"webhook": False,
|
||||
"endpoint_name": "user1_flow_2_endpoint",
|
||||
"tags": ["user1"],
|
||||
"folder_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
}
|
||||
|
||||
# Create flows for the second user
|
||||
flow_user2_1 = {
|
||||
"name": "user2_flow_1",
|
||||
"description": "Flow 1 for user 2",
|
||||
"icon": "string",
|
||||
"icon_bg_color": "#0000ff",
|
||||
"gradient": "string",
|
||||
"data": {},
|
||||
"is_component": False,
|
||||
"webhook": False,
|
||||
"endpoint_name": "user2_flow_1_endpoint",
|
||||
"tags": ["user2"],
|
||||
"folder_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
|
||||
}
|
||||
|
||||
# Create flows using the appropriate user headers
|
||||
response1 = await client.post("api/v1/flows/", json=flow_user1_1, headers=logged_in_headers)
|
||||
assert response1.status_code == status.HTTP_201_CREATED
|
||||
|
||||
response2 = await client.post("api/v1/flows/", json=flow_user1_2, headers=logged_in_headers)
|
||||
assert response2.status_code == status.HTTP_201_CREATED
|
||||
|
||||
response3 = await client.post("api/v1/flows/", json=flow_user2_1, headers=other_user_headers)
|
||||
assert response3.status_code == status.HTTP_201_CREATED
|
||||
|
||||
# Test read_flows for user 1 - should only return user 1's flows
|
||||
params = {
|
||||
"remove_example_flows": True, # Exclude example flows to focus on our test flows
|
||||
"components_only": False,
|
||||
"get_all": True,
|
||||
"header_flows": False,
|
||||
"page": 1,
|
||||
"size": 50,
|
||||
}
|
||||
|
||||
response_user1 = await client.get("api/v1/flows/", params=params, headers=logged_in_headers)
|
||||
result_user1 = response_user1.json()
|
||||
|
||||
assert response_user1.status_code == status.HTTP_200_OK
|
||||
assert isinstance(result_user1, list), "The result must be a list"
|
||||
|
||||
# Verify only user 1's flows are returned
|
||||
user1_flow_names = [flow["name"] for flow in result_user1]
|
||||
assert "user1_flow_1" in user1_flow_names, "User 1's first flow should be returned"
|
||||
assert "user1_flow_2" in user1_flow_names, "User 1's second flow should be returned"
|
||||
assert "user2_flow_1" not in user1_flow_names, "User 2's flow should not be returned for user 1"
|
||||
|
||||
# Verify all returned flows belong to user 1
|
||||
for flow in result_user1:
|
||||
assert str(flow["user_id"]) == str(active_user.id), f"Flow {flow['name']} should belong to user 1"
|
||||
|
||||
# Test read_flows for user 2 - should only return user 2's flows
|
||||
response_user2 = await client.get("api/v1/flows/", params=params, headers=other_user_headers)
|
||||
result_user2 = response_user2.json()
|
||||
|
||||
assert response_user2.status_code == status.HTTP_200_OK
|
||||
assert isinstance(result_user2, list), "The result must be a list"
|
||||
|
||||
# Verify only user 2's flows are returned
|
||||
user2_flow_names = [flow["name"] for flow in result_user2]
|
||||
assert "user2_flow_1" in user2_flow_names, "User 2's flow should be returned"
|
||||
assert "user1_flow_1" not in user2_flow_names, "User 1's first flow should not be returned for user 2"
|
||||
assert "user1_flow_2" not in user2_flow_names, "User 1's second flow should not be returned for user 2"
|
||||
|
||||
# Verify all returned flows belong to user 2
|
||||
for flow in result_user2:
|
||||
assert str(flow["user_id"]) == str(other_user_id), f"Flow {flow['name']} should belong to user 2"
|
||||
|
||||
# Cleanup: Delete the other user
|
||||
async with session_scope() as session:
|
||||
user = await session.get(User, other_user_id)
|
||||
if user:
|
||||
await session.delete(user)
|
||||
await session.commit()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue