fix: store SSH public key on rental to enable SSH after key deletion #290
Conversation
WalkthroughReplaces usage of an SSH key foreign-key (ssh_key_id) with direct storage and propagation of the SSH public key (ssh_public_key) across DB schema/migrations, API extractors and routes, aggregator models/service, SDK types, and CLI handlers/helpers; updates queries, migrations, models, and call sites to persist and carry ssh_public_key end-to-end. Changes
sequenceDiagram
autonumber
participant CLI as CLI (user)
participant SDK as Basilica SDK
participant API as Basilica API
participant DB as Database
participant Validator as Validator Service
CLI->>SDK: list/status request
SDK->>API: HTTP request (list/status)
API->>DB: SELECT rentals (r.ssh_public_key ...)
DB-->>API: rows with ssh_public_key
API->>Validator: validator request (includes ssh_public_key)
Validator-->>API: status + ssh_credentials
API-->>SDK: response (status, ssh_credentials, ssh_public_key)
SDK-->>CLI: return data
Note over CLI,API: CLI uses ssh_public_key to resolve local private key for SSH ops
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
crates/basilica-api/migrations/017_drop_ssh_key_fk_constraint.sql (1)
14-15: Migration is irreversible - ensure proper backup/rollback strategy.The
DROP COLUMN ssh_key_idstatements are destructive operations. While the backfill preserves the public key data, there's no way to recreate thessh_key_idreference if a rollback is needed. This is acceptable for a forward-only migration, but ensure:
- Database backups are in place before deployment
- The team understands this migration cannot be rolled back via SQL
Also applies to: 25-25
crates/basilica-cli/src/cli/handlers/gpu_rental.rs (1)
1838-1863: Consider extracting duplicated private key resolution logic.The pattern for resolving the private key from
ssh_public_keyis duplicated acrosshandle_exec,handle_ssh, andhandle_cp:let public_key = ssh_public_key.ok_or_else(|| { CliError::Internal( eyre!("No SSH public key stored for this rental") .suggestion("The rental was created without an SSH key") .note("Create a new rental with an SSH key to enable SSH access"), ) })?; crate::ssh::find_private_key_for_public_key(&public_key).map_err(CliError::Internal)?This could be extracted to a helper function like
resolve_private_key_from_rental(ssh_public_key: Option<String>) -> Result<PathBuf, CliError>to reduce duplication and ensure consistent error messaging.crates/basilica-cli/src/cli/handlers/gpu_rental_helpers.rs (1)
639-648: Unused return value - consider using the fetchedssh_public_keyor simplifying the call.
fetch_community_ssh_inforeturns(ssh_command, ssh_public_key)but the second element is discarded. Instead,rental.ssh_public_keyfrom the list response is used. This works but is slightly redundant sincefetch_community_ssh_infomakes an additional API call toget_rental_statuswhich also contains thessh_public_key.Consider either:
- Using the
ssh_public_keyfrom the fetch helper (avoiding reliance on the list data)- Or simplifying to only fetch
ssh_credentialsif the list data is preferredif let Ok(ref community) = community_result { if let Some(rental) = community.rentals.iter().find(|r| r.rental_id == target_id) { - let (ssh_command, _) = fetch_community_ssh_info(target_id, api_client).await?; + let (ssh_command, ssh_public_key) = fetch_community_ssh_info(target_id, api_client).await?; return Ok(RentalWithSsh { rental_id: target_id.to_string(), compute_type: ComputeCategory::CommunityCloud, ssh_command, - ssh_public_key: rental.ssh_public_key.clone(), + ssh_public_key, }); } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
crates/basilica-aggregator/src/db.rs(8 hunks)crates/basilica-aggregator/src/models.rs(1 hunks)crates/basilica-aggregator/src/service.rs(1 hunks)crates/basilica-api/migrations/017_drop_ssh_key_fk_constraint.sql(1 hunks)crates/basilica-api/src/api/extractors/ownership.rs(10 hunks)crates/basilica-api/src/api/routes/rentals.rs(6 hunks)crates/basilica-api/src/api/routes/secure_cloud.rs(2 hunks)crates/basilica-cli/src/cli/handlers/gpu_rental.rs(6 hunks)crates/basilica-cli/src/cli/handlers/gpu_rental_helpers.rs(7 hunks)crates/basilica-sdk/src/types.rs(3 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-10-13T15:34:46.544Z
Learnt from: epappas
Repo: one-covenant/basilica PR: 184
File: crates/basilica-validator/src/rental/billing.rs:129-130
Timestamp: 2025-10-13T15:34:46.544Z
Learning: In the basilica-validator crate, rental_id is a String type (not UUID) throughout the system - in RentalInfo struct, database rentals.id column, and all rental operations. The persistence layer methods should accept &str for rental_id parameters, not &Uuid, since they immediately convert to String for SQL binding anyway.
Applied to files:
crates/basilica-sdk/src/types.rscrates/basilica-api/src/api/routes/secure_cloud.rscrates/basilica-api/src/api/extractors/ownership.rscrates/basilica-aggregator/src/db.rscrates/basilica-cli/src/cli/handlers/gpu_rental_helpers.rscrates/basilica-api/src/api/routes/rentals.rs
🧬 Code graph analysis (2)
crates/basilica-cli/src/cli/handlers/gpu_rental.rs (1)
crates/basilica-cli/src/ssh/key_matcher.rs (1)
find_private_key_for_public_key(24-103)
crates/basilica-api/src/api/routes/rentals.rs (3)
crates/basilica-sdk-python/python/basilica/_basilica.pyi (1)
RentalStatusWithSshResponse(184-199)crates/basilica-sdk/src/types.rs (1)
from_validator_response(226-242)crates/basilica-api/src/api/extractors/ownership.rs (1)
store_rental_ownership(120-150)
🔇 Additional comments (27)
crates/basilica-aggregator/src/models.rs (1)
139-139: LGTM - Field rename aligns with the new data model.The rename from
ssh_key_idtossh_public_keycorrectly reflects that this field now stores the actual SSH public key content rather than a foreign key reference. This enables SSH access to persist even after the registered key is deleted.crates/basilica-aggregator/src/service.rs (1)
521-521: LGTM - Stores public key at deployment creation time.The change correctly captures and stores the SSH public key content when creating a deployment, enabling SSH access to persist independently of the registered key lifecycle.
crates/basilica-aggregator/src/db.rs (3)
172-173: LGTM - Consistent column rename in create_deployment.SQL INSERT and binding correctly updated from
ssh_key_idtossh_public_key.Also applies to: 186-186
282-282: LGTM - Row parsing updated in parse_deployment_from_row.
297-297: LGTM - All SELECT queries consistently updated.The column name change is applied uniformly across
get_deployment,get_deployment_by_user,list_deployments_by_user,list_deployments, andget_all_active_deployments.Also applies to: 319-319, 338-338, 366-366, 423-423
crates/basilica-api/src/api/routes/secure_cloud.rs (3)
38-39: LGTM - Query simplified by removing ssh_keys JOIN.The query now reads
ssh_public_keydirectly from the rental record instead of joining with thessh_keystable. This aligns with the PR objective and ensures SSH access works even after key deletion.Also applies to: 150-157
222-222: LGTM - ssh_public_key correctly propagated to response.The public key is properly destructured from the query result and included in the
SecureCloudRentalListItemresponse.Also applies to: 292-292
323-328: LGTM - Archive operation preserves ssh_public_key.The archive query correctly copies
ssh_public_keyfrom active to terminated rental records.crates/basilica-api/src/api/extractors/ownership.rs (4)
23-23: LGTM - Ownership structs extended with ssh_public_key.
UserRentalRowandOwnedRentalcorrectly include the new field, and extraction properly populates it.Also applies to: 38-38, 78-78
99-112: LGTM - Database queries and store function updated.The
get_rental_ownershipquery andstore_rental_ownershipfunction correctly handle the newssh_public_keycolumn.Also applies to: 119-150
176-182: LGTM - RentalWithDetails includes ssh_public_key.The struct and
get_user_rentals_with_detailsfunction properly return the public key for each rental.Also applies to: 211-240
265-276: LGTM - Archive operation preserves ssh_public_key.The archive query correctly copies
ssh_public_keyfrom active to terminated community rental records.crates/basilica-sdk/src/types.rs (2)
97-100: LGTM - ApiRentalListItem extended with ssh_public_key.The new field is appropriately optional and includes
skip_serializing_ifto omit null values from JSON responses.
213-215: LGTM - RentalStatusWithSshResponse gains ssh_public_key.The struct and
from_validator_responseconstructor are correctly updated. The signature change requires all call sites to pass the new parameter, which aligns with the broader refactoring.Also applies to: 225-230, 238-238
crates/basilica-cli/src/cli/handlers/gpu_rental.rs (2)
1682-1717: LGTM - handle_exec uses rental-specific SSH public key.The handler correctly:
- Destructures
ssh_public_keyfrom the resolved rental- Returns a clear error if no key is stored
- Uses
find_private_key_for_public_keyto locate the matching private keyThis enables SSH access even after the globally registered key is deleted.
1739-1779: LGTM - handle_ssh uses rental-specific SSH public key.Same pattern as
handle_exec, correctly implemented.crates/basilica-api/migrations/017_drop_ssh_key_fk_constraint.sql (1)
27-31: Community cloud rentals do not need backfill—they never stored ssh_key_id.The migration correctly omits backfill logic for
user_rentalsandterminated_user_rentals. Unlikesecure_cloud_rentals, which stored anssh_key_idforeign key reference, community rentals capturessh_public_keydirectly at rental creation time by querying the user's SSH key from thessh_keystable (seerentals.rslines 80-95). Since no foreign key exists in these tables, there is no data to backfill from. Existing community rentals already havessh_public_keypopulated if they were created with SSH enabled.crates/basilica-api/src/api/routes/rentals.rs (5)
57-63: LGTM!The
ssh_public_keyis correctly propagated from theOwnedRentalextractor to the response, aligning with the SDK's updatedfrom_validator_responsesignature.
236-246: LGTM!The SSH public key is correctly cloned and passed to the validator request, maintaining the key association with the rental at creation time.
267-279: LGTM!The conversion of empty string to
Nonebefore storage is appropriate, ensuring clean database values. The ownership storage correctly includes the SSH public key.
652-660: LGTM!The
ssh_public_key_mapis correctly built alongside the other maps, enabling efficient lookup when assembling the rental list response.
727-756: LGTM!The SSH public key retrieval using
.and_then(|opt| opt.clone())correctly handles theOption<Option<String>>from the map lookup, and the field is properly included inApiRentalListItem.crates/basilica-cli/src/cli/handlers/gpu_rental_helpers.rs (5)
606-611: LGTM!The
RentalWithSshstruct correctly gains thessh_public_keyfield to support per-rental SSH key storage, enabling SSH access after the registered key is deleted.
651-670: LGTM!The secure cloud path correctly retrieves
ssh_public_keyfrom the rental object and includes it in the result.
678-697: LGTM!The interactive selector path correctly unpacks both
ssh_commandandssh_public_keyfrom the fetch helpers and includes them in theRentalWithSshresult.
700-723: LGTM!The
fetch_community_ssh_infohelper correctly retrieves rental status and returns bothssh_credentialsandssh_public_key. Error handling properly guides users when SSH is unavailable.
725-753: LGTM!The
fetch_secure_ssh_infohelper correctly extractsssh_commandandssh_public_keyfrom the secure cloud rental, with appropriate error handling for missing SSH configuration.
Previously, deleting an SSH key from the API would: 1. Fail with FK constraint error if active rentals existed 2. Break SSH access because CLI looked up the registered key Now, the SSH public key is stored directly on the rental record at creation time. This allows: - SSH key deletion without constraint errors - SSH access using the rental's stored key, independent of registered keys
itzlambda
force-pushed
the
fix-ssh-key-delete-fk
branch
from
164299c to
2f0b80f
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
crates/basilica-api/src/api/extractors/ownership.rs (1)
477-479: Consider adding test coverage for non-nullssh_public_key.The test passes
Nonefor the newssh_public_keyparameter. Consider adding a test case that stores and verifies retrieval of an actual SSH public key value to ensure end-to-end correctness of the new field.crates/basilica-cli/src/cli/handlers/gpu_rental.rs (1)
1852-1863: Consider extracting duplicate error handling logic.The identical error handling for missing SSH public keys is duplicated across
handle_exec(lines 1706-1717),handle_ssh(lines 1768-1779), andhandle_cp(lines 1852-1863). Extracting this into a helper function would improve maintainability.Consider adding a helper function:
/// Resolve the private key path from a rental's SSH public key fn resolve_private_key_from_rental( ssh_public_key: Option<String> ) -> Result<PathBuf, CliError> { let public_key = ssh_public_key.ok_or_else(|| { CliError::Internal( eyre!("No SSH public key available for this rental") .suggestion("The rental may have been created without SSH, or the SSH key was deleted") .note("SSH access requires the original key used during rental creation"), ) })?; crate::ssh::find_private_key_for_public_key(&public_key) .map_err(CliError::Internal) }Then use it in all three handlers:
let private_key_path = resolve_private_key_from_rental(ssh_public_key)?;
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
crates/basilica-aggregator/src/db.rs(8 hunks)crates/basilica-aggregator/src/models.rs(1 hunks)crates/basilica-aggregator/src/service.rs(1 hunks)crates/basilica-api/migrations/017_drop_ssh_key_fk_constraint.sql(1 hunks)crates/basilica-api/src/api/extractors/ownership.rs(10 hunks)crates/basilica-api/src/api/routes/rentals.rs(6 hunks)crates/basilica-api/src/api/routes/secure_cloud.rs(2 hunks)crates/basilica-cli/src/cli/handlers/gpu_rental.rs(6 hunks)crates/basilica-cli/src/cli/handlers/gpu_rental_helpers.rs(7 hunks)crates/basilica-sdk/src/types.rs(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- crates/basilica-api/src/api/routes/secure_cloud.rs
- crates/basilica-sdk/src/types.rs
- crates/basilica-cli/src/cli/handlers/gpu_rental_helpers.rs
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-10-13T15:34:46.544Z
Learnt from: epappas
Repo: one-covenant/basilica PR: 184
File: crates/basilica-validator/src/rental/billing.rs:129-130
Timestamp: 2025-10-13T15:34:46.544Z
Learning: In the basilica-validator crate, rental_id is a String type (not UUID) throughout the system - in RentalInfo struct, database rentals.id column, and all rental operations. The persistence layer methods should accept &str for rental_id parameters, not &Uuid, since they immediately convert to String for SQL binding anyway.
Applied to files:
crates/basilica-api/src/api/routes/rentals.rscrates/basilica-aggregator/src/db.rscrates/basilica-api/src/api/extractors/ownership.rs
🧬 Code graph analysis (1)
crates/basilica-api/src/api/routes/rentals.rs (2)
crates/basilica-sdk/src/types.rs (1)
from_validator_response(226-242)crates/basilica-api/src/api/extractors/ownership.rs (1)
store_rental_ownership(120-150)
🔇 Additional comments (11)
crates/basilica-aggregator/src/service.rs (1)
521-521: LGTM!The change correctly stores the SSH public key directly on the deployment record at creation time, which aligns with the PR objective to decouple rentals from the
ssh_keystable and preserve SSH access after key deletion.crates/basilica-aggregator/src/models.rs (1)
139-139: LGTM!The field rename from
ssh_key_idtossh_public_keycorrectly reflects the new data model where the actual public key is stored instead of a foreign key reference.crates/basilica-aggregator/src/db.rs (2)
166-197: LGTM!The
create_deploymentfunction correctly usesssh_public_keyin both the SQL INSERT statement and the bind parameter, maintaining consistency with the updated model.
272-290: LGTM!The
parse_deployment_from_rowhelper correctly maps thessh_public_keycolumn to the Deployment struct field, ensuring consistent deserialization across all query paths.crates/basilica-api/src/api/extractors/ownership.rs (3)
119-150: LGTM!The
store_rental_ownershipfunction correctly accepts and persists thessh_public_keyparameter, enabling SSH public key storage at rental creation time.
211-240: LGTM!The
get_user_rentals_with_detailsfunction correctly retrieves and maps thessh_public_keyfield, ensuring the public key is available when listing rental details.
265-276: LGTM!The archive operation correctly copies
ssh_public_keyto the terminated rentals table, preserving the SSH access information in historical records.crates/basilica-api/migrations/017_drop_ssh_key_fk_constraint.sql (1)
27-31: The migration handles community and secure cloud rentals correctly according to their respective designs.user_rentalsandterminated_user_rentalsnever had anssh_key_idcolumn—they usessh_credentials(added in migrations 002 and 003). Onlysecure_cloud_rentalsandterminated_secure_cloud_rentalshave thessh_key_idcolumn that requires backfill and removal. No further action needed.Likely an incorrect or invalid review comment.
crates/basilica-api/src/api/routes/rentals.rs (3)
57-62: LGTM!The addition of SSH public key to the response construction is correct. The field is properly sourced from the database ownership record and passed to
from_validator_response, aligning with the new API surface.
239-239: LGTM!The SSH public key propagation through the rental creation flow is correctly implemented:
- The key is cloned and passed to the validator request
- Empty string (used for no-SSH rentals) is properly converted to
Nonefor database storage- The ownership record stores the key alongside SSH credentials
Also applies to: 268-279
652-660: LGTM!The SSH public key retrieval and inclusion in the rentals list follows the established pattern used for SSH status and port mappings. The map-based lookup correctly handles the optional nature of the field, and the key is properly surfaced in the API response.
Also applies to: 727-730, 755-755
Replace anonymous tuple with named RentalDetailsRow struct for rental details query, improving code readability. Update SSH key error messages to better explain that the key may exist on a different machine.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
crates/basilica-cli/src/cli/handlers/gpu_rental.rs (2)
1706-1717: Clarify error message: remove misleading suggestion about missing key.When
ssh_public_keyis None, the error message suggests "the required SSH key is not on this machine," but this scenario only applies whenfind_private_key_for_public_keyfails. If the public key is None, there's no key to look for.Consider simplifying:
let public_key = ssh_public_key.ok_or_else(|| { CliError::Internal( eyre!("No SSH public key available for this rental") - .suggestion("The rental may have been created without SSH, or the required SSH key is not on this machine") - .note("SSH access requires the original key used during rental creation"), + .suggestion("The rental was created without SSH access") + .note("Create a new rental with an SSH key to enable SSH"), ) })?;The scenario "key not on this machine" is already handled by
find_private_key_for_public_keyerror messages.
1768-1779: Apply the same error message clarification here.This error handling has the same issue as in
handle_exec(lines 1706-1717). The message incorrectly suggests the key might not be on the machine whenssh_public_keyis None.Apply the same fix as suggested for
handle_exec.
🧹 Nitpick comments (1)
crates/basilica-api/src/api/extractors/ownership.rs (1)
483-493: Consider adding test coverage for ssh_public_key storage.The test currently passes
Nonefor ssh_public_key and only verifies ssh_credentials. To validate the new functionality, consider adding a test case that stores a public key and verifies it's correctly retrieved.Example enhancement:
// Store ownership with SSH credentials -store_rental_ownership(&db, rental_id, user_id, ssh_creds, None, None) +let ssh_public_key = Some("ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ..."); +store_rental_ownership(&db, rental_id, user_id, ssh_creds, None, ssh_public_key) .await .expect("Failed to store ownership"); // Now user should own the rental and have SSH credentials let ownership = get_rental_ownership(&db, rental_id, user_id) .await .expect("Failed to check ownership"); assert!(ownership.is_some()); let row = ownership.unwrap(); assert_eq!(row.ssh_credentials, ssh_creds.map(String::from)); +assert_eq!(row.ssh_public_key, ssh_public_key.map(String::from));
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
crates/basilica-api/src/api/extractors/ownership.rs(11 hunks)crates/basilica-cli/src/cli/handlers/gpu_rental.rs(6 hunks)crates/basilica-cli/src/cli/handlers/gpu_rental_helpers.rs(7 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- crates/basilica-cli/src/cli/handlers/gpu_rental_helpers.rs
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-10-13T15:34:46.544Z
Learnt from: epappas
Repo: one-covenant/basilica PR: 184
File: crates/basilica-validator/src/rental/billing.rs:129-130
Timestamp: 2025-10-13T15:34:46.544Z
Learning: In the basilica-validator crate, rental_id is a String type (not UUID) throughout the system - in RentalInfo struct, database rentals.id column, and all rental operations. The persistence layer methods should accept &str for rental_id parameters, not &Uuid, since they immediately convert to String for SQL binding anyway.
Applied to files:
crates/basilica-cli/src/cli/handlers/gpu_rental.rscrates/basilica-api/src/api/extractors/ownership.rs
🧬 Code graph analysis (1)
crates/basilica-cli/src/cli/handlers/gpu_rental.rs (1)
crates/basilica-cli/src/ssh/key_matcher.rs (1)
find_private_key_for_public_key(24-103)
🔇 Additional comments (5)
crates/basilica-api/src/api/extractors/ownership.rs (4)
23-23: LGTM: ssh_public_key field additions are correct.The ssh_public_key field is consistently added as
Option<String>across structs and properly mapped in the ownership extractor.Also applies to: 38-38, 78-78
101-101: LGTM: ownership storage correctly updated for ssh_public_key.The query and function signature properly handle the new ssh_public_key field with correct parameter ordering and SQL bindings.
Also applies to: 119-150
176-246: LGTM: rental details correctly include ssh_public_key.The RentalWithDetails struct, query, and mapping are all properly updated to include the ssh_public_key field.
273-274: LGTM: archive operations correctly preserve ssh_public_key.Both community and secure cloud archive operations properly copy ssh_public_key to terminated tables, maintaining historical SSH configuration data.
Also applies to: 329-334
crates/basilica-cli/src/cli/handlers/gpu_rental.rs (1)
1686-1686: LGTM: RentalWithSsh destructuring correctly includes ssh_public_key.All handlers (exec, ssh, cp) properly extract ssh_public_key from the resolved rental.
Also applies to: 1743-1743, 1840-1840
Summary
This fixes two issues:
Summary by CodeRabbit
New Features
Refactor
✏️ Tip: You can customize this high-level summary in your review settings.