WIP: Extend DB for controls table and a mapper table to map controls… #42
Conversation
WalkthroughThe pull request introduces significant changes across multiple files, primarily focusing on enhancements to database configuration and schema management within the Admyral project. Key modifications include the addition of new environment variables for test databases, updates to the handling of database URLs, and the introduction of new schemas and relationships for controls and workflows. Additionally, a new PostgreSQL service configuration for unit testing is added in the Docker setup, along with various test functions to validate the new functionalities. Changes
Possibly related PRs
Suggested reviewers
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
|
The latest updates on your projects. Learn more about Vercel for Git ↗︎
|
| # primary keys | ||
| control_id: int = Field( | ||
| primary_key=True, | ||
| index=True, |
There was a problem hiding this comment.
no need to build an index on a PK
| ) | ||
| op.drop_column("workflow_controls_results", "control_id") |
There was a problem hiding this comment.
Why is this here? I can't see any removal in the diff
| # foreign keys | ||
| user_id: str = Field(sa_type=TEXT(), index=True) |
There was a problem hiding this comment.
Lets make user_id part of the primary key, so we can easily swap it with an org_id/workspace_id later
| ), | ||
| sa.PrimaryKeyConstraint("control_id", "user_id", "workflow_id"), | ||
| ) | ||
| op.drop_column("workflow_controls_results", "control_id") |
There was a problem hiding this comment.
Why do we need this? The workflow_controls_results table does not contain a control_id column (see here)
| from sqlalchemy import text | ||
|
|
||
| # Use SQLite for testing | ||
| DATABASE_URL = "sqlite:///:memory:" |
There was a problem hiding this comment.
The unit tests already expects that our postgres DB is running. I would suggest to use the postgres instead of sqlite. We should then also generalize the session fixture to use the same one across tests. See tests/workers/conftest.py for the fixture used by the other tests.
| # Create a control | ||
| control = ControlSchema( | ||
| control_name="Test Control", | ||
| control_description="A test control", | ||
| user_id="user123", | ||
| ) | ||
| session.add(control) | ||
|
|
||
| # Create a workflow | ||
| workflow = WorkflowSchema( | ||
| workflow_id="workflow123", | ||
| workflow_name="Test Workflow", | ||
| user_id="user123", | ||
| workflow_dag={}, | ||
| is_active=True, | ||
| ) | ||
| session.add(workflow) | ||
| session.commit() | ||
|
|
||
| # Create the mapping | ||
| mapping = ControlsWorkflowsMappingSchema( | ||
| control_id=control.control_id, workflow_id=workflow.workflow_id | ||
| ) | ||
| session.add(mapping) |
There was a problem hiding this comment.
Maybe lets combine the test_create_control_workflow_mapping and test_cascade_delete_workflow test because this here is code duplication and there's no need to test it twice imo.
There was a problem hiding this comment.
As an alternative, we could factor this out in a helper function for setting up the scenario.
leon-schmid
force-pushed
the
ADM-398/Extend-DB-Schema-For-Controls
branch
from
f5eb776 to
bd8f26d
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 14
🧹 Outside diff range and nitpick comments (12)
admyral/models/control.py (2)
1-2: Consider adding validation imports for robust data handling.Since this file defines data models that will interact with a database, consider importing additional Pydantic validation tools:
-from pydantic import BaseModel +from pydantic import BaseModel, Field, validator, constrThis would allow you to:
- Add length constraints on strings
- Validate control_id as positive
- Add custom validation logic if needed
1-12: Consider architectural improvements for better maintainability.Given that this is part of a larger database extension:
- Consider creating a base class for common validation logic:
class AdmyralBaseModel(BaseModel): """Base class for all Admyral models with common validation logic.""" class Config: validate_assignment = True extra = "forbid"
- Document the relationship between Pydantic models and SQLModel schemas:
""" This module defines Pydantic models for data validation. These models are paired with SQLModel schemas in admyral/db/schemas/control_schemas.py for database operations. Note: These models should be kept in sync with their corresponding SQLModel schemas. """Would you like me to create a GitHub issue to track these architectural improvements?
tests/conftest.py (1)
Line range hint
9-22: Well-structured fixture with appropriate scope and cleanup.The session-scoped fixture with proper event loop handling and sequential cleanup is a good design. Consider adding error handling around the cleanup operations to ensure they execute even if tests fail.
Consider wrapping the cleanup in a try-finally block:
@pytest.fixture(scope="session", autouse=True) async def store(event_loop): store = await AdmyralStore.create_store(database_url=CONFIG.test_database_url) yield store - await store.clean_up_workflow_data_of(TEST_USER_ID) - await store.clean_up_controls_data() + try: + yield store + finally: + await store.clean_up_workflow_data_of(TEST_USER_ID) + await store.clean_up_controls_data()admyral/db/schemas/workflow_control_results_schemas.py (1)
35-35: Document the relationship purposeConsider adding a docstring to explain the relationship's purpose and the one-to-many nature between workflows and their control results.
- workflow: "WorkflowSchema" = Relationship(back_populates="control_results") + workflow: "WorkflowSchema" = Relationship( + back_populates="control_results", + description="Reference to the parent workflow that owns these control results" + )docker-compose.dev.yml (2)
30-35: Consider optimizing health check for faster test feedback.The current health check configuration might be too conservative for a test database. Consider reducing the intervals for faster feedback during testing.
healthcheck: test: ["CMD-SHELL", "pg_isready", "-d", "admyral"] - interval: 30s - timeout: 60s - retries: 5 - start_period: 80s + interval: 10s + timeout: 20s + retries: 3 + start_period: 30s
23-29: Add resource constraints for test database.To prevent the test database from consuming excessive resources during CI/CD pipelines or local development, consider adding memory and CPU limits.
postgresql-unit-tests: container_name: postgresql-unit-tests + deploy: + resources: + limits: + memory: 512M + cpus: '0.5' environment:admyral/db/alembic/versions/8fbcc7f42a0b_add_controls_and_controls_workflows_.py (1)
1-85: Consider additional architectural improvementsWhile the basic schema is sound, consider these architectural suggestions:
- Add appropriate CHECK constraints for data validation (e.g., non-empty control_name)
- Consider adding a status column to track control lifecycle
- Consider adding metadata columns for audit purposes (e.g., last_modified_by)
These suggestions would make the schema more robust and future-proof. Would you like me to provide specific implementation details for any of these improvements?
admyral/db/alembic/database_manager.py (2)
Line range hint
13-13: Address critical TODOsSeveral TODO comments need attention, particularly those related to:
- Filtering of alembic_version table (security implications?)
- Database drop functionality (important for testing)
- Database migration safety (backup requirements)
Would you like me to help create GitHub issues to track these TODOs and propose implementations for them?
Also applies to: 73-75, 89-93
Line range hint
26-126: Consider splitting database management and migration concernsThe
DatabaseManagerclass currently handles both database administration (creation, existence checks) and schema migrations. Consider splitting these into separate classes following the Single Responsibility Principle:
DatabaseAdministrator: For database creation/deletionMigrationManager: For handling schema migrationsThis separation would:
- Improve testability
- Make the code more maintainable
- Allow for better error handling specific to each concern
- Make it easier to modify one aspect without affecting the other
Would you like me to propose a detailed design for this separation?
tests/db/test_controls_workflows_mapping.py (1)
17-19: Consider improving the control ID generation strategy.The current implementation uses a fixed range (1-999999) which could potentially conflict with real control IDs in production. Consider:
- Using a dedicated range for test control IDs
- Adding validation for negative numbers
def generate_random_control_id() -> int: - return random.randint(1, 999999) + TEST_CONTROL_ID_MIN = 1_000_000 # Dedicated range for test IDs + TEST_CONTROL_ID_MAX = 1_999_999 + control_id = random.randint(TEST_CONTROL_ID_MIN, TEST_CONTROL_ID_MAX) + if control_id <= 0: + raise ValueError("Control ID must be positive") + return control_idadmyral/config/config.py (1)
Line range hint
118-154: Consider environment-specific configuration managementThe current approach of maintaining separate production and test database URLs might benefit from a more structured environment-based configuration system. This would be particularly valuable as the project grows and potentially needs to support more environments (staging, development, etc.).
Consider implementing an environment-based configuration system:
- Define an Environment enum (PRODUCTION, TEST, STAGING, etc.)
- Create environment-specific configuration classes
- Use factory pattern to load appropriate configuration based on ADMYRAL_ENV
This would make it easier to:
- Manage environment-specific settings
- Validate configurations per environment
- Add new environments in the future
- Test configuration loading logic
admyral/db/schemas/control_schemas.py (1)
64-64: Remove redundantUniqueConstrainton primary key columns.The
UniqueConstrainton("user_id", "control_id")is redundant because these fields are already defined as primary keys inControlSchema. Primary key columns are inherently unique, so this constraint is unnecessary.Apply this diff to remove the redundant unique constraint:
__table_args__ = ( - UniqueConstraint("user_id", "control_id", name="unique_control_id"), ForeignKeyConstraint( ["user_id"], ["User.id"], ondelete="CASCADE", ), )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (14)
admyral/config/config.py(2 hunks)admyral/db/admyral_store.py(5 hunks)admyral/db/alembic/database_manager.py(2 hunks)admyral/db/alembic/versions/8fbcc7f42a0b_add_controls_and_controls_workflows_.py(1 hunks)admyral/db/schemas/__init__.py(2 hunks)admyral/db/schemas/auth_schemas.py(2 hunks)admyral/db/schemas/control_schemas.py(1 hunks)admyral/db/schemas/workflow_control_results_schemas.py(1 hunks)admyral/db/schemas/workflow_schemas.py(2 hunks)admyral/models/__init__.py(2 hunks)admyral/models/control.py(1 hunks)docker-compose.dev.yml(1 hunks)tests/conftest.py(2 hunks)tests/db/test_controls_workflows_mapping.py(1 hunks)
🔇 Additional comments (24)
admyral/models/control.py (1)
10-12: 🛠️ Refactor suggestion
Add validation and documentation for the mapping model.
The mapping model needs:
- Documentation explaining its purpose
- Validation for workflow_id format
- Potential unique constraint consideration
Consider this implementation:
+"""
+Maps the relationship between security controls and workflows.
+
+This mapping table establishes many-to-many relationships between
+controls and workflows, allowing a single control to be associated
+with multiple workflows and vice versa.
+
+Attributes:
+ control_id: Reference to the Control entity
+ workflow_id: Reference to the Workflow entity
+"""
class ControlsWorkflowsMapping(BaseModel):
- control_id: int
- workflow_id: str
+ control_id: int = Field(gt=0, description="Reference to the Control entity")
+ workflow_id: str = Field(
+ regex="^[a-zA-Z0-9_-]+$",
+ min_length=1,
+ max_length=255,
+ description="Reference to the Workflow entity"
+ )
+
+ class Config:
+ schema_extra = {
+ "example": {
+ "control_id": 1,
+ "workflow_id": "workflow_123"
+ }
+ }Let's verify the workflow_id format in existing code:
tests/conftest.py (3)
3-3: LGTM: Import statement is correctly placed and necessary.
The import of CONFIG is required for accessing the test database configuration.
21-21: Verify controls cleanup implementation.
The addition of clean_up_controls_data() is good for test isolation, but let's verify its implementation.
✅ Verification successful
Based on the results, I can generate the final response as we have sufficient information about the implementation:
Controls cleanup implementation is properly handled
The cleanup implementation in admyral/db/admyral_store.py follows good practices:
- Uses transaction handling via async context manager
- Cleans up all related tables in a logical order
- Commits the transaction after all deletions
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the controls cleanup implementation
# Check the implementation of clean_up_controls_data
ast-grep --pattern 'clean_up_controls_data($$$) {
$$$
}'
# Look for transaction handling in the cleanup
rg -A 10 "clean_up_controls_data"
# Check for related schema definitions
ast-grep --pattern 'class $_(Schema|Table)($$$) {
$$$
}'
Length of output: 1688
18-18: Verify test database configuration setup.
The store creation now depends on CONFIG.test_database_url. Let's ensure this configuration is properly set up.
✅ Verification successful
Test database configuration is properly set up
The verification shows that the test database configuration is correctly implemented:
- The
ADMYRAL_TEST_DATABASE_URLenvironment variable is properly defined inadmyral/config/config.py - A dedicated PostgreSQL test database service (
postgresql-unit-tests) is configured in Docker, running on port 5433 - The configuration follows the project's transition to PostgreSQL support
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the test database configuration setup
# Check if CONFIG.test_database_url is defined in the config file
rg -A 5 "test_database_url"
# Check for any existing test database configuration in environment files
fd -e .env -e .env.example -e .env.test --exec cat {} \;
# Look for Docker test database service configuration
fd -g "*docker*" --exec grep -l "test.*database" {} \;
Length of output: 956
Script:
#!/bin/bash
# Let's check for environment variable definitions and their usage
# Check for environment variable definitions
rg "ADMYRAL_TEST_DATABASE_URL"
# Check for any test database configuration in Docker files
fd -e yml -e yaml -g "*docker*" --exec cat {} \;
# Look for any test database setup in test files
rg -g "test_*.py" "database.*setup|database.*config"
Length of output: 10391
admyral/db/schemas/__init__.py (2)
22-25: LGTM! Clean and well-structured import statement.
The new imports follow the established pattern and are logically grouped with other schema imports.
43-44: LGTM! Proper exposure of new schemas.
The new schemas are correctly added to the __all__ list, maintaining alphabetical order and following the module's public interface pattern.
Let's verify the schema definitions exist:
✅ Verification successful
Schema definitions are properly implemented and imported
The verification confirms that:
- Both
ControlSchemaandControlsWorkflowsMappingSchemaare properly defined inadmyral/db/schemas/control_schemas.py - The schemas are correctly imported by other modules (
auth_schemas.pyandworkflow_schemas.py) - The schemas follow the project's pattern by inheriting from
BaseSchemawithtable=True
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify that the control schema definitions exist and are properly defined
# Check for the existence and content of control_schemas.py
ast-grep --pattern 'class ControlSchema:
$$$
'
ast-grep --pattern 'class ControlsWorkflowsMappingSchema:
$$$
'
# Verify no circular imports exist
rg -l "from admyral.db.schemas import.*ControlSchema" | grep -v "__init__.py"
Length of output: 196
Script:
#!/bin/bash
# Let's try a different approach to locate these schema definitions
# Search for files containing these schema classes
fd -e py . | xargs rg "class (ControlSchema|ControlsWorkflowsMappingSchema)"
# Look for any imports of these schemas in the codebase
rg "from.*import.*(ControlSchema|ControlsWorkflowsMappingSchema)"
# Check content of potential schema files
fd "control.*schema.*\.py" -I -i
Length of output: 774
docker-compose.dev.yml (1)
23-43: LGTM! Well-structured test database configuration.
The new PostgreSQL service for unit tests is well-configured with:
- Proper isolation from production (separate ports, volumes)
- Consistent environment variable handling
- Appropriate health checks
admyral/db/alembic/versions/8fbcc7f42a0b_add_controls_and_controls_workflows_.py (1)
80-85: Verify referenced table dependencies
The migration references the User and workflows tables. Let's verify these dependencies exist.
admyral/models/__init__.py (2)
66-66: LGTM! Import statement follows project conventions.
The new import for control-related models is properly placed and follows the established import pattern.
132-133: Verify readiness for public API exposure.
The new control-related models are being added to the public API. Since this is marked as WIP, please ensure that:
- The model interfaces are stable and well-documented
- All necessary validation rules are in place
- The models properly handle all edge cases
Let's verify the model definitions:
✅ Verification successful
Models are well-defined and ready for public API exposure
The Control-related models are properly implemented with:
- Clear and simple interfaces in
admyral/models/control.py - Comprehensive database schemas with proper constraints and relationships
- Extensive test coverage validating core functionality
- Proper cascade delete behavior and relationship handling
- Working integration with the workflow system
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check the completeness of Control models
# Expected: Find model definitions with proper documentation and validation
# Check for model class definitions and their docstrings
ast-grep --pattern 'class Control:
$$$
'
# Look for validation rules
rg -A 5 'validator|Field|constraint' admyral/models/control.py
Length of output: 111
Script:
#!/bin/bash
# Let's try a different approach to find the model definitions
# First, let's find the relevant model files
fd "control.*\.py$"
# Then let's examine their content
fd "control.*\.py$" -x cat {}
# Also check for any test files that might give us insights into the model usage
fd "test.*control.*\.py$" -x cat {}
Length of output: 21993
admyral/db/schemas/workflow_schemas.py (3)
17-17: LGTM! Clean import organization
The new imports are well-organized and follow good practices, particularly the use of TYPE_CHECKING for avoiding circular imports.
Also applies to: 25-25
71-78: Verify impact on existing code
The relationship changes (replacing direct controls with control_workflow_mappings, control_results, and controls) might impact existing code that relies on these relationships.
✅ Verification successful
No impact on existing code from relationship changes
The verification shows that the existing usage of workflow.controls in graph_conversion.py refers to workflow_dag.controls, which is a different attribute from the SQLModel relationship being modified. The relationship changes in workflow_schemas.py are safe as there are no direct usages of these SQLModel relationships in queries or joins across the codebase.
- The only match found in
admyral/editor/graph_conversion.pyis accessingworkflow.workflow_dag.controls, not the SQLModel relationship - No SQL joins or queries were found that rely on these relationships
- The new relationship structure properly maintains the connections through
ControlsWorkflowsMappingSchema
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for existing usage of workflow.controls
rg -l "workflow.*\.controls" --type py
# Search for potential ORM queries using these relationships
rg -l "\.join.*controls.*workflow" --type py
Length of output: 122
Script:
#!/bin/bash
# Let's check the actual usage context in graph_conversion.py
rg -A 5 -B 5 "workflow.*\.controls" admyral/editor/graph_conversion.py
# Let's also check for any other potential relationships between workflows and controls
ast-grep --pattern 'workflow.$_: list[$_] = Relationship($$$)'
# Check for any direct SQL or ORM queries involving workflows and controls
rg -A 5 "from.*workflow.*import|select.*workflow.*control|where.*workflow.*control" --type py
Length of output: 16209
77-78: Verify the corresponding relationship in ControlSchema
The relationship looks correct, but we should verify that ControlSchema has a matching workflows relationship with the same link_model.
✅ Verification successful
The relationship is correctly defined in both schemas
The workflows relationship in ControlSchema is properly set up with back_populates="controls" and link_model=ControlsWorkflowsMappingSchema, which matches exactly with the corresponding relationship in WorkflowSchema. The many-to-many relationship between workflows and controls is correctly implemented on both sides.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for the corresponding relationship in ControlSchema
ast-grep --pattern 'class ControlSchema {
$$$
workflows: list["WorkflowSchema"] = Relationship(
back_populates="controls",
link_model=ControlsWorkflowsMappingSchema
)
$$$
}'
Length of output: 191
Script:
#!/bin/bash
# Let's first find the ControlSchema definition
rg "class ControlSchema" -A 10
Length of output: 746
Script:
#!/bin/bash
# Let's check the entire ControlSchema file for any relationship definitions
cat admyral/db/schemas/control_schemas.py
Length of output: 2918
admyral/db/schemas/auth_schemas.py (2)
16-16: LGTM: Import statement follows project conventions
The import statement is properly placed with other schema imports and follows the project's import style.
57-59: Relationship configuration looks good, verify migration strategy
The relationship configuration follows the established pattern and correctly implements cascading delete for user-owned controls. However, since this is a WIP PR adding a new table relationship, please ensure:
- A migration plan exists for handling existing users
- The corresponding Control table is properly set up with the reciprocal relationship
Let's verify the migration and control table setup:
tests/db/test_controls_workflows_mapping.py (3)
72-72: Verify database and session configuration.
Based on previous review comments:
- Ensure tests are using PostgreSQL instead of SQLite
- Consider using the shared session fixture from
tests/workers/conftest.py
Let's verify the current configuration:
✅ Verification successful
Based on the shell script results, I can now generate the final response:
Database and session configuration is properly set up
The test is correctly configured:
- Using PostgreSQL via
CONFIG.test_database_urlin the session-scoped store fixture - Properly utilizing
AdmyralStoreand its async session maker - Following the project's database testing patterns with proper cleanup
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check database configuration
rg -A 5 "async_session_maker|AdmyralStore" --type py
# Check session fixture usage
rg -A 5 "pytest.fixture.*session" --type py
Length of output: 19508
94-115: 🛠️ Refactor suggestion
Refactor test cases to reduce duplication and improve coverage.
Several improvements are needed:
-
The setup code is duplicated between tests
-
Cleanup should be consistent across all tests
-
Manual cascade delete could be replaced with database CASCADE
-
Missing negative test cases
-
Add a fixture for cleanup:
@pytest.fixture(autouse=True)
async def cleanup(store: AdmyralStore):
yield
async with store.async_session_maker() as session:
await store.clean_up_workflow_data_of(TEST_USER_ID)
await store.clean_up_controls_data()- Add negative test cases:
@pytest.mark.asyncio
async def test_create_control_workflow_mapping_invalid_ids(store: AdmyralStore):
async with store.async_session_maker() as session:
with pytest.raises(IntegrityError):
mapping = ControlsWorkflowsMappingSchema(
control_id=-1, # Invalid ID
workflow_id="invalid",
user_id=TEST_USER_ID,
)
session.add(mapping)
await session.commit()- Use database CASCADE instead of manual deletion:
Let's verify if CASCADE is configured:
Also applies to: 117-176
38-50: Add validation for workflow DAG structure.
The workflow DAG is initialized as an empty dictionary without any schema validation. This could lead to invalid DAG structures in production.
Let's verify the DAG schema requirements:
admyral/config/config.py (2)
118-118: LGTM!
The new environment variable constant follows the established naming convention and is appropriately placed with other database-related constants.
154-154: LGTM!
The new test_database_url attribute is well-integrated into the GlobalConfig class, maintaining consistency with other database-related configurations.
admyral/db/admyral_store.py (3)
42-43: LGTM: New schema imports follow conventions
The new schema imports for controls are properly organized with other schema imports and follow the established naming conventions.
95-98: LGTM: Backward compatible method signature update
The create_store method changes maintain backward compatibility by making database_url optional and defaulting to CONFIG.database_url. This is a good approach for supporting both new and existing code.
79-83: Verify the impact of constructor signature change
The constructor signature has changed from accepting GlobalConfig to database_url, which is a breaking change. This could affect existing code that instantiates AdmyralStore.
✅ Verification successful
Let's examine the test files to understand the usage patterns and verify the impact.
Let's try a different approach to find the usage patterns.
Let's check one more thing to see how the store is being created.
Based on the evidence from the test files, I can now provide a final response:
Constructor change is properly handled across the codebase
The change from GlobalConfig to direct database_url parameter is safe because:
- All instantiations in tests use the
create_storefactory method withdatabase_urlparameter - The main instantiation pattern seen in
conftest.pyusesAdmyralStore.create_store(database_url=CONFIG.test_database_url) - No direct instantiations using the old
GlobalConfigparameter were found in the codebase
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Find all direct instantiations of AdmyralStore to verify impact
# Test: Search for AdmyralStore instantiations. Expect: Instances using the new signature.
# Search for direct instantiations
rg -p "AdmyralStore\(" --type py
# Search for potential test files that might need updates
fd -e py -x grep -l "AdmyralStore" {} \; | grep "test"
Length of output: 333
Script:
#!/bin/bash
# Check test files for AdmyralStore instantiation patterns
rg -p "AdmyralStore\(" --type py -A 3 ./tests/
Length of output: 48
Script:
#!/bin/bash
# Search for AdmyralStore usage in test files with more context
for file in tests/workers/utils.py tests/workers/test_workflow_executor.py tests/conftest.py tests/compiler/test_workflow_compiler.py tests/db/test_controls_workflows_mapping.py; do
echo "=== $file ==="
cat "$file" | grep -A 5 -B 5 "AdmyralStore"
done
# Also check if there are any imports of GlobalConfig in the main file
rg "GlobalConfig" admyral/db/admyral_store.py
Length of output: 11675
Script:
#!/bin/bash
# Check the create_store method implementation
ast-grep --pattern 'async def create_store' -A 10 admyral/db/admyral_store.py
Length of output: 79
admyral/db/schemas/control_schemas.py (1)
67-67: Verify the correctness of the foreign key reference to User.id.
Ensure that the foreign key constraint references the correct table and column. If the actual table name is users or differs from User, this could lead to database errors.
Please confirm that "User.id" correctly references the primary key of the User table.
admyral/db/alembic/versions/8fbcc7f42a0b_add_controls_and_controls_workflows_.py
Outdated
Show resolved
Hide resolved
admyral/db/alembic/versions/8fbcc7f42a0b_add_controls_and_controls_workflows_.py
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (8)
admyral/db/alembic/versions/1239762f7ef5_add_controls_and_controls_workflows_.py (1)
39-42: Consider adding constraints and optimizing column types.A few suggestions to improve the table design:
- Consider making
control_idauto-incrementing (SERIAL)- Consider adding length limits to TEXT columns to prevent excessive storage usage
- sa.Column("control_id", sa.Integer(), nullable=False), - sa.Column("user_id", sa.TEXT(), nullable=False), - sa.Column("control_name", sa.TEXT(), nullable=False), - sa.Column("control_description", sa.TEXT(), nullable=False), + sa.Column("control_id", sa.Integer(), sa.Sequence('control_id_seq'), nullable=False), + sa.Column("user_id", sa.String(36), nullable=False), + sa.Column("control_name", sa.String(255), nullable=False), + sa.Column("control_description", sa.String(1000), nullable=False),admyral/db/admyral_store.py (2)
80-84: Consider making SQL statement logging configurableThe SQLAlchemy engine is created with
echo=True, which logs all SQL statements. This might generate excessive logs in production.Consider making this configurable:
- def __init__(self, database_url: str) -> None: + def __init__(self, database_url: str, debug_sql: bool = False) -> None: self.database_url = database_url self.engine = create_async_engine( - database_url, echo=True, future=True, pool_pre_ping=True + database_url, echo=debug_sql, future=True, pool_pre_ping=True )
1085-1102: Optimize transaction management for multiple delete operationsThe method performs multiple delete operations with separate commits. This can be optimized by using a single transaction.
Consider this optimization:
async def clean_up_controls_data(self, user_id: str) -> None: """Delete all data from controls-related tables. Should only be used in testing environments.""" async with self._get_async_session() as db: - # Delete workflow control results where the workflow belongs to the user - await db.exec( - delete(WorkflowControlResultsSchema).where( - WorkflowControlResultsSchema.workflow_id.in_( - sa_select(WorkflowSchema.workflow_id).where( - WorkflowSchema.user_id == user_id - ) - ) - ) - ) - await db.exec( - delete(ControlsWorkflowsMappingSchema).where( - ControlsWorkflowsMappingSchema.user_id == user_id - ) - ) - await db.exec(delete(ControlSchema).where(ControlSchema.user_id == user_id)) - await db.commit() + try: + # Delete workflow control results where the workflow belongs to the user + await db.exec( + delete(WorkflowControlResultsSchema).where( + WorkflowControlResultsSchema.workflow_id.in_( + sa_select(WorkflowSchema.workflow_id).where( + WorkflowSchema.user_id == user_id + ) + ) + ) + ) + await db.exec( + delete(ControlsWorkflowsMappingSchema).where( + ControlsWorkflowsMappingSchema.user_id == user_id + ) + ) + await db.exec(delete(ControlSchema).where(ControlSchema.user_id == user_id)) + await db.commit() + except Exception: + await db.rollback() + raisetests/db/test_controls_workflows_mapping.py (5)
17-23: Add type annotations to helper functions for clarityConsider adding type annotations to the
generate_random_control_idandgenerate_random_workflow_idfunctions to improve code readability and facilitate static type checking.Apply this diff:
def generate_random_control_id() -> int: return random.randint(1, 999999) def generate_random_workflow_id() -> str: return f"workflow-{uuid.uuid4()}"
25-35: Add type annotations tocreate_controlfunctionAdding type annotations to the
create_controlfunction parameters enhances code readability and helps with static analysis tools.Apply this diff:
+from sqlmodel import AsyncSession async def create_control(session: AsyncSession, control_id: int = 1): control = ControlSchema( control_name="Test Control", control_description="A test control",Ensure that
AsyncSessionis imported appropriately.
38-49: Add type annotations tocreate_workflowfunctionSimilarly, add type annotations to the
create_workflowfunction for clarity.Apply this diff:
+from sqlmodel import AsyncSession async def create_workflow(session: AsyncSession, workflow_id: str = "workflow123"): workflow = WorkflowSchema( workflow_id=workflow_id, workflow_name=f"Test Workflow-{workflow_id}",Ensure that
AsyncSessionis imported appropriately.
118-176: Leverage cascading deletes to simplifytest_cascade_delete_workflowCurrently,
test_cascade_delete_workflowmanually deletes related entries before deleting the workflow. If the database schema is configured with proper cascading deletes, you can simplify the test by deleting the workflow directly and asserting that related entries are automatically removed. This ensures that the cascading behavior is correctly implemented and tested.For example, you can modify the test as follows:
- Delete the workflow first.
- Commit the session.
- Verify that related entries are automatically deleted.
- Adjust assertions accordingly.
Apply this diff:
# Delete the workflow await session.delete(workflow) await session.commit() # Verify cascading deletes mapping_result = await session.exec(select(ControlsWorkflowsMappingSchema)) mappings = mapping_result.all() assert len(mappings) == 0 results_result = await session.exec(select(WorkflowControlResultsSchema)) results = results_result.all() assert len(results) == 0Ensure that the database schema is configured with
ON DELETE CASCADEfor the foreign key relationships.
143-160: Refactor deletion logic into a helper functionThe code blocks for deleting entries are similar. Consider refactoring the deletion of entries into a helper function to reduce code duplication and improve maintainability.
For example:
async def delete_entries(session, model, **filters): entries_to_delete = await session.exec(select(model).filter_by(**filters)) for entry in entries_to_delete: await session.delete(entry)Then, you can use this helper function:
await delete_entries(session, WorkflowControlResultsSchema, workflow_id=workflow.workflow_id) await delete_entries(session, ControlsWorkflowsMappingSchema, workflow_id=workflow.workflow_id)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (5)
admyral/db/admyral_store.py(6 hunks)admyral/db/alembic/versions/1239762f7ef5_add_controls_and_controls_workflows_.py(1 hunks)admyral/db/schemas/control_schemas.py(1 hunks)tests/conftest.py(2 hunks)tests/db/test_controls_workflows_mapping.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- admyral/db/schemas/control_schemas.py
- tests/conftest.py
🔇 Additional comments (10)
admyral/db/alembic/versions/1239762f7ef5_add_controls_and_controls_workflows_.py (3)
1-20: LGTM! Migration metadata is well-structured.
The migration metadata follows best practices with proper typing hints and required imports.
80-85: LGTM! Downgrade function is properly implemented.
The tables and indexes are dropped in the correct order, and cascade deletes will handle cleanup automatically.
1-85: Verify PostgreSQL compatibility.
Since the project now exclusively supports PostgreSQL, let's verify the compatibility of all database features used in this migration.
✅ Verification successful
PostgreSQL compatibility confirmed for migration.
The migration uses standard SQLAlchemy features that are fully compatible with PostgreSQL:
- TIMESTAMP with timezone support (standard PostgreSQL feature)
- Standard foreign key constraints with CASCADE delete
- Basic table and index creation operations
- Server-default timestamps using
now()
All these features are well-supported in PostgreSQL and align with the existing database schema, as evidenced by the previous migration that explicitly uses PostgreSQL dialect for timestamp handling.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for any PostgreSQL version requirements or compatibility issues in project docs
rg -i "postgresql.*version|postgres.*version" --type md
# Check for any existing PostgreSQL-specific features in other migrations
fd -e py -d 1 . admyral/db/alembic/versions/ --exec rg -l "postgresql|postgres"
Length of output: 218
Script:
#!/bin/bash
# Check the content of the previous migration that uses PostgreSQL features
cat admyral/db/alembic/versions/7985f1c159a3_switch_to_timezone_aware_datetime_.py
# Let's also check if there are any database configuration or setup files
fd -e py -e sql . admyral/db/ --exec rg -l "postgresql|postgres"
# Check for any SQLAlchemy configuration that might indicate PostgreSQL requirements
rg "create_engine|Engine" admyral/db/
Length of output: 16886
admyral/db/admyral_store.py (3)
96-99: LGTM!
The changes to create_store maintain backward compatibility while adding flexibility to specify a custom database URL.
113-113: LGTM!
The DatabaseManager initialization is correctly updated to use both the engine and database URL.
1083-1084: Add environment check for test-only method
The docstring indicates this method should only be used in testing environments, but there's no programmatic enforcement of this restriction.
tests/db/test_controls_workflows_mapping.py (4)
52-69: Consider adding error handling and docstrings to setup_scenario function
Adding error handling and a docstring to setup_scenario will enhance its robustness and readability.
71-82: Test for create_control is properly implemented
The test function test_create_control accurately tests the creation of a control, and all assertions are appropriate.
83-93: Test for create_workflow is properly implemented
The test function test_create_workflow correctly tests the creation of a workflow, with appropriate assertions.
95-115: Test for control-workflow mapping is comprehensive
The test_create_control_workflow_mapping function effectively tests the mapping between controls and workflows, including verification of relationships.
admyral/db/alembic/versions/1239762f7ef5_add_controls_and_controls_workflows_.py
Outdated
Show resolved
Hide resolved
admyral/db/alembic/versions/1239762f7ef5_add_controls_and_controls_workflows_.py
Outdated
Show resolved
Hide resolved
|
|
||
| ADMYRAL_TEST_DATABASE_URL = os.getenv( | ||
| ENV_ADMYRAL_TEST_DATABASE_URL, | ||
| "postgresql://postgres:your-super-secret-and-long-postgres-password@localhost:5433/admyral", |
There was a problem hiding this comment.
The default should point to the test postgres and not the actual postgres database. Can you also adapt .env.example and set the an env var for the test database, as well, so that we're consistent with how we set the normal database url.
docker-compose.dev.yml
Outdated
| POSTGRES_DB: admyral | ||
| image: postgres:${POSTGRES_VERSION:-16.4-bookworm} | ||
| healthcheck: | ||
| test: ["CMD-SHELL", "pg_isready", "-d", "admyral"] |
There was a problem hiding this comment.
Minor: Let's stay consistent and use temporal instead of admyral.
| { | ||
| "control_id": self.control_id, | ||
| "workflow_id": self.workflow_id, | ||
| "user_id": self.user_id, |
There was a problem hiding this comment.
we usually don't return the user_id because it is unnecessary metadata.
There was a problem hiding this comment.
also this will fail since ControlsWorkflowsMapping does not have user_id member.
admyral/db/admyral_store.py
Outdated
| @@ -8,6 +8,7 @@ | |||
| from contextlib import asynccontextmanager | |||
| from uuid import uuid4 | |||
| import json | |||
| from sqlalchemy import select as sa_select | |||
There was a problem hiding this comment.
You can just use select from sqlmodel. sqlmodel is a wrapper around sqlalchemy.
| return f"workflow-{uuid.uuid4()}" | ||
|
|
||
|
|
||
| async def create_control(session, control_id=1): |
There was a problem hiding this comment.
The default is not needed because we always generate a random id. Same holds for create_workflow.
| results_to_delete = await session.exec( | ||
| select(WorkflowControlResultsSchema).where( | ||
| WorkflowControlResultsSchema.workflow_id == workflow.workflow_id | ||
| ) | ||
| ) | ||
| for result in results_to_delete.all(): | ||
| await session.delete(result) |
There was a problem hiding this comment.
You can just do
| results_to_delete = await session.exec( | |
| select(WorkflowControlResultsSchema).where( | |
| WorkflowControlResultsSchema.workflow_id == workflow.workflow_id | |
| ) | |
| ) | |
| for result in results_to_delete.all(): | |
| await session.delete(result) | |
| await session.exec( | |
| delete(WorkflowControlResultsSchema).where( | |
| WorkflowControlResultsSchema.workflow_id == workflow.workflow_id | |
| ) | |
| ) |
There is no need for retrieving everything first.
| # Then delete the workflow mappings | ||
| mappings_to_delete = await session.exec( | ||
| select(ControlsWorkflowsMappingSchema).where( | ||
| ControlsWorkflowsMappingSchema.workflow_id == workflow.workflow_id | ||
| ) | ||
| ) | ||
| for mapping in mappings_to_delete.all(): | ||
| await session.delete(mapping) |
There was a problem hiding this comment.
maybe let's delete the workflows for the test user instead
| mapping_result = await session.exec(select(ControlsWorkflowsMappingSchema)) | ||
| mappings = mapping_result.all() | ||
| assert len(mappings) == 0 | ||
|
|
||
| results_result = await session.exec(select(WorkflowControlResultsSchema)) | ||
| results = results_result.all() | ||
| assert len(results) == 0 | ||
|
|
||
| control_result = await session.exec(select(ControlSchema)) |
There was a problem hiding this comment.
I would filter for the TEST_USER_ID
|
|
||
| control_result = await session.exec(select(ControlSchema)) | ||
| controls = control_result.all() | ||
| assert len(controls) == 1 |
There was a problem hiding this comment.
Can we also test the other direction? Delete a control, then it should remove the corresponding rows from the mapping table
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Outside diff range and nitpick comments (5)
.env.example (1)
7-7: LGTM! Consider adding clarifying comments.The addition of a separate test database URL is a good practice. However, consider adding comments to:
- Indicate this is specifically for PostgreSQL databases
- Explain the purpose of the test database URL
- Document that this supports the new test database Docker service
ADMYRAL_DATABASE_URL="postgresql://postgres:your-super-secret-and-long-postgres-password@localhost:5432/admyral" +# Test database URL for running unit tests with the test database Docker service +# Note: PostgreSQL is now the only supported database ADMYRAL_TEST_DATABASE_URL="postgresql://postgres:your-super-secret-and-long-postgres-password@localhost:5433/admyral"admyral/db/admyral_store.py (1)
Line range hint
1072-1079: Add input parameter validationThe method should validate that the input parameters are not None or empty strings before proceeding with the database operation.
async def store_workflow_control_result( self, workflow_id: str, run_id: str, result: str, ) -> None: + if not workflow_id or not run_id or result is None: + raise ValueError("workflow_id, run_id, and result must not be empty") async with self._get_async_session() as db: await db.exec(admyral/db/alembic/versions/3967e08c0051_add_controls_and_controls_workflows_.py (1)
13-13: Unnecessary import ofsqlmodelThe
sqlmodelmodule is imported but not used in this script. Consider removing the import statement to clean up the code.tests/db/test_controls_workflows_mapping.py (2)
20-21: Ensure consistent ID generation for workflowsThe
generate_random_workflow_idfunction adds a'workflow-'prefix to the UUID, whereasgenerate_random_control_iddoes not add a prefix. For consistency and to adhere to standard UUID usage, consider removing the prefix fromgenerate_random_workflow_id.Apply this diff to update the function:
-def generate_random_workflow_id() -> str: - return f"workflow-{uuid.uuid4()}" +def generate_random_workflow_id() -> str: + return str(uuid.uuid4())
37-41: Adjust workflow name to include a descriptive prefixWith the removal of the
'workflow-'prefix from theworkflow_id, consider adjusting theworkflow_nameto maintain a clear identification of test workflows.Apply this diff to update the workflow name:
- workflow_name=f"Test Workflow-{workflow_id}", + workflow_name=f"Test Workflow-workflow-{workflow_id}",
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (8)
.env.example(1 hunks)admyral/config/config.py(2 hunks)admyral/db/admyral_store.py(5 hunks)admyral/db/alembic/database_manager.py(2 hunks)admyral/db/alembic/versions/3967e08c0051_add_controls_and_controls_workflows_.py(1 hunks)admyral/db/schemas/control_schemas.py(1 hunks)docker-compose.dev.yml(1 hunks)tests/db/test_controls_workflows_mapping.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- admyral/db/schemas/control_schemas.py
🔇 Additional comments (22)
.env.example (1)
6-7: Consider using distinct database names for production and test environments.
Both database URLs use the same database name "admyral" but different ports. This could lead to confusion or potential data corruption if the wrong port is used. Consider:
- Using distinct database names (e.g., "admyral_test" for testing)
- Using different placeholder passwords to emphasize these are separate environments
docker-compose.dev.yml (4)
42-43: Use consistent environment variable naming
For consistency with other environment variables in the project, we should use POSTGRES_TEST_VOLUME_PATH instead of TEST_POSTGRES_VOLUME_PATH.
Apply this diff:
volumes:
- - ${TEST_POSTGRES_VOLUME_PATH:-.admyral/postgres-test}:/var/lib/postgresql/data
+ - ${POSTGRES_TEST_VOLUME_PATH:-.admyral/postgres-test}:/var/lib/postgresql/data28-28: Use consistent database naming
For consistency, we should use temporal instead of admyral as the database name, following the project's naming convention.
Apply this diff:
- POSTGRES_DB: admyral
+ POSTGRES_DB: temporal30-31:
Fix healthcheck configuration
The healthcheck is configured to check the temporal database, but this service is configured to use the admyral database. This mismatch could lead to incorrect health status reporting.
Apply this diff:
healthcheck:
- test: ["CMD-SHELL", "pg_isready", "-d", "temporal"]
+ test: ["CMD-SHELL", "pg_isready", "-d", "admyral"]Likely invalid or redundant comment.
23-44: Verify database initialization scripts
The new PostgreSQL service for unit tests might require specific initialization scripts or schema setup for the controls and workflows tables mentioned in the PR description.
Let's check for any database initialization scripts:
✅ Verification successful
Let me search for any database-related files or migrations that might contain schema information.
Let me check the Prisma schema to understand the database structure and requirements.
No initialization scripts needed for unit tests database
The PostgreSQL container for unit tests doesn't require additional initialization scripts. The database schema is managed by Prisma ORM, which automatically handles schema creation and migrations during application startup. The schema.prisma file contains all necessary table definitions, and Prisma will create these tables when the application connects to the database.
- The
POSTGRES_DB: admyralenvironment variable in the docker-compose file correctly sets up the initial database - Prisma will handle the schema creation and migrations for all required tables (users, workflows, secrets, etc.)
- The healthcheck is configured appropriately to ensure database availability
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Look for SQL initialization scripts that might need to be mounted
fd -e sql -e sh --type f
Length of output: 186
Script:
#!/bin/bash
# Look for migration files or database setup files
rg -g '!node_modules' -l -i "(migration|schema|database|sql|postgres)" --type-add 'prisma:*.prisma' --type prisma
Length of output: 140
Script:
#!/bin/bash
# Check the Prisma schema content
cat web/prisma/schema.prisma
Length of output: 7107
admyral/db/alembic/database_manager.py (5)
26-28:
Add URL validation and sanitization
The change from GlobalConfig to direct database_url string simplifies the interface but removes potential validation. Consider adding URL validation to prevent SQL injection and ensure proper URL format.
43-44:
Use URLlib for robust URL parsing
The current string manipulation approach for URL parsing is fragile and could break with complex URLs containing parameters or special characters.
48-58:
Improve error handling and use dynamic database name
The current implementation has several issues:
- Hardcoded database name instead of using the one from URL
- Over-broad exception handling
- Potential resource leaks in error cases
61-64:
Add error handling and SQL injection protection
The database creation lacks error handling and uses a hardcoded database name, making it vulnerable to SQL injection.
61-64:
Add check for race conditions in database creation
The create_database method should check if the database exists before attempting creation. In a concurrent environment, another process might create the database between the last existence check and the creation attempt.
Consider wrapping the creation in a transaction and handling potential race conditions:
async def create_database(self) -> None:
engine = self._get_postgres_setup_engine()
async with engine.connect() as conn:
- await conn.execute(text("commit"))
- await conn.execute(text("create database admyral"))
+ await conn.execute(text("begin"))
+ try:
+ # Check again within transaction
+ result = await conn.execute(
+ text("select exists (select 1 from pg_database where datname = :name)")
+ .bindparams(name=db_name)
+ )
+ if result.scalar():
+ await conn.execute(text("rollback"))
+ raise ValueError(f"Database {db_name} already exists")
+
+ await conn.execute(
+ text("create database :name").bindparams(name=db_name)
+ )
+ await conn.execute(text("commit"))
+ except Exception as e:
+ await conn.execute(text("rollback"))
+ raise RuntimeError(f"Failed to create database: {e}") from eLikely invalid or redundant comment.
admyral/config/config.py (4)
118-118: LGTM!
The new environment variable constant follows the established naming convention and is appropriately placed with other database-related constants.
154-154: LGTM!
The new field follows the class's type safety practices and properly uses the configured test database URL.
124-126:
Port 5433 might cause connection issues
- The default port 5433 is non-standard for PostgreSQL (standard is 5432). This might cause connection issues for users with default PostgreSQL installations.
- The URL transformation still lacks validation as mentioned in previous reviews.
Consider:
- Using the standard port 5432 for the main database
- Adding URL validation as suggested in previous reviews
- Documenting the port requirement if 5433 is intentional
- "postgresql://postgres:your-super-secret-and-long-postgres-password@localhost:5433/admyral",
+ "postgresql://postgres:your-super-secret-and-long-postgres-password@localhost:5432/admyral",127-131:
Test database should use a different database name
Using the same database URL for both main and test configurations is risky:
- Could lead to data corruption if tests modify production data
- Makes it harder to run tests in parallel with the main application
Consider:
- Using a different database name for tests (e.g.,
admyral_test) - Updating
.env.exampleas suggested in previous reviews - Adding URL validation as previously recommended
- "postgresql://postgres:your-super-secret-and-long-postgres-password@localhost:5433/admyral",
+ "postgresql://postgres:your-super-secret-and-long-postgres-password@localhost:5433/admyral_test",admyral/db/admyral_store.py (5)
42-42: LGTM!
The import of ControlSchema is correctly placed and necessary for the new controls functionality.
94-97: LGTM!
The changes to create_store maintain backward compatibility while adding flexibility through the optional database_url parameter.
111-111: LGTM!
The DatabaseManager initialization is correctly updated to use both the engine and database_url parameters.
1081-1095: Add environment check for test-only method
While the docstring indicates this method should only be used in testing environments, there's no programmatic enforcement of this restriction.
78-82: Verify the impact of constructor signature change
The constructor now takes database_url instead of GlobalConfig. While this change simplifies the dependency, it's a breaking change that needs careful verification.
✅ Verification successful
Constructor change is safely abstracted by factory method
The constructor signature change from GlobalConfig to database_url is safely handled by the create_store factory method, which:
- Maintains backward compatibility by defaulting to
CONFIG.database_urlwhen no URL is provided - Is consistently used across the codebase for instantiation (no direct constructor calls found)
- Already supports the new
database_urlparameter in all usage locations
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Find all instantiations of AdmyralStore to ensure they're updated
# Test: Search for direct instantiations and factory method calls
rg -A 2 "AdmyralStore\("
rg -A 2 "create_store\("
Length of output: 1090
Script:
#!/bin/bash
# Let's check the create_store factory method implementation to understand how it handles the database_url
ast-grep --pattern 'async def create_store($$$) -> $_: $$$'
# Also check if there are any other references to GlobalConfig in the store
rg "GlobalConfig" admyral/db/admyral_store.py
Length of output: 688
admyral/db/alembic/versions/3967e08c0051_add_controls_and_controls_workflows_.py (2)
43-43:
Verify the table name in the foreign key constraint for user_id
The foreign key constraint references ["User.id"], but table names are typically lowercase and may be pluralized. Ensure that the referenced table is correctly named User, or update it to match the actual table name (e.g., users).
Run the following script to confirm the table name:
✅ Verification successful
The foreign key constraint referencing User table is correct
From the migration files output, we can see that the table is indeed created with uppercase name User in the same codebase:
op.create_table(
"User",
sa.Column(
"created_at",
sa.TIMESTAMP(),
server_default=sa.text("now()"),This confirms that the foreign key constraint ["User.id"] is using the correct table name.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the existence of the 'User' table in the codebase.
# Search for the creation of a table named 'User'
rg "op.create_table\(['\"]User['\"]" admyral/db/alembic/
# Search for a model class named 'User'
rg "class User" admyral/db/models/
Length of output: 158
Script:
#!/bin/bash
# Let's try to find any SQL table definitions or models across the codebase
# First, let's find all alembic migration files
fd -e py . admyral/db/alembic/versions/ -x cat {} | grep -A 5 "create_table"
# Also search for any model definitions across the codebase
fd -t f -e py . -x grep -l "class.*Base"
# And search for SQLAlchemy model definitions with 'user' in filename
fd -t f -e py . -x grep -l "user" {}
Length of output: 9525
65-75: Confirm necessity of user_id in controls_workflows primary key and foreign keys
Including user_id in the primary key and foreign keys of the controls_workflows table may complicate the schema. Verify that this is required for your data model. If control_id and workflow_id uniquely identify a record, you might simplify the schema by excluding user_id.
Run the following script to analyze the usage of user_id in related tables:
✅ Verification successful
Keep user_id in the controls_workflows primary key and foreign keys
The inclusion of user_id in the primary key and foreign keys of the controls_workflows table is necessary and correctly designed based on the following evidence:
- The
controlstable uses a composite primary key of(control_id, user_id)and has a unique constraint on these fields - Controls are user-specific as shown by the foreign key from
controls.user_idtoUser.id - The
controls_workflowstable maintains referential integrity by including both fields in its foreign key constraint to match the parent table's primary key structure
This design ensures:
- Each control belongs to a specific user
- The same control_id can be used by different users
- Proper cascading deletion when a user or control is removed
- Referential integrity across the user-control-workflow relationship chain
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Analyze the necessity of 'user_id' in the 'controls_workflows' table.
# Check how 'user_id' is used in the 'controls' table definition
rg -A 15 "op.create_table\(['\"]controls['\"]" admyral/db/alembic/
# Check how 'user_id' is used in the 'controls_workflows' table definition
rg -A 20 "op.create_table\(['\"]controls_workflows['\"]" admyral/db/alembic/
# Search for references to 'user_id' in models and schemas
rg "user_id" admyral/db/models/
rg "user_id" admyral/db/schemas/
Length of output: 1971
Script:
#!/bin/bash
# Let's check the schema definitions and relationships more thoroughly
# Check for any table definitions that reference controls or workflows
rg -A 20 "op.create_table" admyral/db/alembic/
# Check for any foreign key constraints involving controls or workflows
rg -B 2 -A 2 "ForeignKeyConstraint" admyral/db/alembic/
# Look for any unique constraints that might indicate data relationships
rg -B 2 -A 2 "UniqueConstraint" admyral/db/alembic/
# Check for any SQLAlchemy model definitions that might show the relationships
rg -A 10 "class.*Table" admyral/db/schemas/
Length of output: 42543
tests/db/test_controls_workflows_mapping.py (1)
51-68: Enhance setup_scenario with error handling and documentation
Consider adding error handling and a docstring to the setup_scenario function to improve its robustness and maintainability. This will make the function more understandable and easier to modify in the future.
admyral/db/alembic/versions/3967e08c0051_add_controls_and_controls_workflows_.py
Show resolved
Hide resolved
admyral/db/alembic/versions/3967e08c0051_add_controls_and_controls_workflows_.py
Show resolved
Hide resolved
This PR adds the following:
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Tests
Documentation