|
| 1 | +use axum::{ |
| 2 | + Json, |
| 3 | + extract::Multipart, |
| 4 | +}; |
| 5 | +use dal::{ |
| 6 | + Schema, |
| 7 | + SchemaVariant, |
| 8 | + SchemaVariantId, |
| 9 | + cached_module::CachedModule, |
| 10 | + pkg::{ |
| 11 | + ImportOptions, |
| 12 | + import_funcs_for_module_update, |
| 13 | + import_pkg_from_pkg, |
| 14 | + import_schema_variant, |
| 15 | + }, |
| 16 | +}; |
| 17 | +use sdf_extract::{ |
| 18 | + PosthogEventTracker, |
| 19 | + change_set::ChangeSetDalContext, |
| 20 | +}; |
| 21 | +use serde::{ |
| 22 | + Deserialize, |
| 23 | + Serialize, |
| 24 | +}; |
| 25 | +use serde_json::json; |
| 26 | +use si_events::audit_log::AuditLogKind; |
| 27 | +use si_pkg::{ |
| 28 | + PkgSpec, |
| 29 | + SiPkg, |
| 30 | +}; |
| 31 | +use telemetry::prelude::*; |
| 32 | +use utoipa::{ |
| 33 | + self, |
| 34 | + ToSchema, |
| 35 | +}; |
| 36 | + |
| 37 | +use super::{ |
| 38 | + SchemaError, |
| 39 | + SchemaResult, |
| 40 | +}; |
| 41 | + |
| 42 | +#[derive(Deserialize, Serialize, Debug, ToSchema)] |
| 43 | +#[serde(rename_all = "camelCase")] |
| 44 | +pub struct InstallFromFileV1Response { |
| 45 | + #[schema(value_type = String, example = "01H9ZQD35JPMBGHH69BT0Q79VZ")] |
| 46 | + pub schema_id: si_events::SchemaId, |
| 47 | + #[schema(value_type = String, example = "01H9ZQD35JPMBGHH69BT0Q79VY")] |
| 48 | + pub schema_variant_id: si_events::SchemaVariantId, |
| 49 | + #[schema(example = "AWS::EC2::Instance")] |
| 50 | + pub schema_name: String, |
| 51 | + #[schema(example = "EC2 Instance")] |
| 52 | + pub display_name: String, |
| 53 | + #[schema(example = "AWS::EC2")] |
| 54 | + pub category: String, |
| 55 | +} |
| 56 | + |
| 57 | +/// Install a schema from a PkgSpec file |
| 58 | +/// |
| 59 | +/// Accepts a multipart form with a `pkg_spec` field containing the JSON PkgSpec. |
| 60 | +/// If the schema already exists, it will be upgraded with the new variant. |
| 61 | +#[utoipa::path( |
| 62 | + post, |
| 63 | + path = "/v1/w/{workspace_id}/change-sets/{change_set_id}/schemas/install_from_file", |
| 64 | + params( |
| 65 | + ("workspace_id" = String, Path, description = "Workspace identifier"), |
| 66 | + ("change_set_id" = String, Path, description = "Change Set identifier"), |
| 67 | + ), |
| 68 | + tag = "schemas", |
| 69 | + summary = "Install a schema from a PkgSpec file", |
| 70 | + responses( |
| 71 | + (status = 200, description = "Schema installed successfully", body = InstallFromFileV1Response), |
| 72 | + (status = 400, description = "Bad request - Invalid or missing pkg_spec", body = crate::service::v1::common::ApiError), |
| 73 | + (status = 401, description = "Unauthorized - Invalid or missing token"), |
| 74 | + (status = 422, description = "Validation error - Invalid PkgSpec data", body = crate::service::v1::common::ApiError), |
| 75 | + (status = 500, description = "Internal server error", body = crate::service::v1::common::ApiError) |
| 76 | + ) |
| 77 | +)] |
| 78 | +pub async fn install_from_file( |
| 79 | + ChangeSetDalContext(ref ctx): ChangeSetDalContext, |
| 80 | + tracker: PosthogEventTracker, |
| 81 | + mut multipart: Multipart, |
| 82 | +) -> SchemaResult<Json<InstallFromFileV1Response>> { |
| 83 | + if ctx.change_set_id() == ctx.get_workspace_default_change_set_id().await? { |
| 84 | + return Err(SchemaError::NotPermittedOnHead); |
| 85 | + } |
| 86 | + |
| 87 | + // Extract the pkg_spec field from multipart form |
| 88 | + let mut maybe_module_json = None; |
| 89 | + while let Some(field) = multipart.next_field().await? { |
| 90 | + match field.name() { |
| 91 | + Some("pkg_spec") => { |
| 92 | + maybe_module_json = Some(field.bytes().await?); |
| 93 | + } |
| 94 | + _ => debug!("Unknown multipart form field on module install, skipping..."), |
| 95 | + } |
| 96 | + } |
| 97 | + |
| 98 | + let Some(module_bytes) = maybe_module_json else { |
| 99 | + return Err(SchemaError::PkgFileError("Missing pkg_spec field")); |
| 100 | + }; |
| 101 | + |
| 102 | + let module_string = String::from_utf8_lossy(&module_bytes); |
| 103 | + let spec: PkgSpec = serde_json::from_str(&module_string)?; |
| 104 | + let pkg = SiPkg::load_from_spec(spec)?; |
| 105 | + |
| 106 | + // Validate that the package has exactly one schema |
| 107 | + let schemas = pkg.schemas()?; |
| 108 | + if schemas.len() != 1 { |
| 109 | + return Err(SchemaError::PkgFileError( |
| 110 | + "Pkg must have exactly one schema", |
| 111 | + )); |
| 112 | + } |
| 113 | + |
| 114 | + let schema_spec = schemas |
| 115 | + .first() |
| 116 | + .ok_or(SchemaError::PkgFileError("Pkg has no schemas"))?; |
| 117 | + let schema_name = schema_spec.name(); |
| 118 | + |
| 119 | + // Check if schema exists in change set or module cache |
| 120 | + let schema_exists_in_cache = CachedModule::find_latest_for_schema_name(ctx, schema_name) |
| 121 | + .await? |
| 122 | + .is_some(); |
| 123 | + |
| 124 | + let variant_ids = match Schema::get_by_name_opt(ctx, schema_name).await? { |
| 125 | + Some(existing_schema) => { |
| 126 | + upgrade_schema_from_uploaded_file(ctx, &pkg, existing_schema).await? |
| 127 | + } |
| 128 | + None if schema_exists_in_cache => { |
| 129 | + let installed_schema = Schema::get_or_install_by_name(ctx, schema_name).await?; |
| 130 | + upgrade_schema_from_uploaded_file(ctx, &pkg, installed_schema).await? |
| 131 | + } |
| 132 | + None => { |
| 133 | + let (_, variant_ids, _) = import_pkg_from_pkg( |
| 134 | + ctx, |
| 135 | + &pkg, |
| 136 | + Some(ImportOptions { |
| 137 | + schema_id: None, |
| 138 | + past_module_hashes: None, |
| 139 | + ..Default::default() |
| 140 | + }), |
| 141 | + ) |
| 142 | + .await |
| 143 | + .map_err(SchemaError::Pkg)?; |
| 144 | + variant_ids |
| 145 | + } |
| 146 | + }; |
| 147 | + |
| 148 | + let schema_variant_id = variant_ids |
| 149 | + .first() |
| 150 | + .ok_or(SchemaError::PkgFileError("Pkg has no variants"))?; |
| 151 | + |
| 152 | + let variant = SchemaVariant::get_by_id(ctx, *schema_variant_id).await?; |
| 153 | + let schema = variant.schema(ctx).await?; |
| 154 | + |
| 155 | + tracker.track( |
| 156 | + ctx, |
| 157 | + "api_install_from_file", |
| 158 | + json!({ |
| 159 | + "schema_id": schema.id(), |
| 160 | + "schema_variant_id": variant.id(), |
| 161 | + "display_name": variant.display_name(), |
| 162 | + "category": variant.category(), |
| 163 | + }), |
| 164 | + ); |
| 165 | + |
| 166 | + ctx.write_audit_log( |
| 167 | + AuditLogKind::CreateSchemaVariant { |
| 168 | + schema_id: schema.id(), |
| 169 | + schema_variant_id: variant.id(), |
| 170 | + }, |
| 171 | + variant.display_name().to_string(), |
| 172 | + ) |
| 173 | + .await?; |
| 174 | + |
| 175 | + ctx.commit().await?; |
| 176 | + |
| 177 | + Ok(Json(InstallFromFileV1Response { |
| 178 | + schema_id: schema.id(), |
| 179 | + schema_variant_id: variant.id(), |
| 180 | + schema_name: schema.name.clone(), |
| 181 | + display_name: variant.display_name().to_string(), |
| 182 | + category: variant.category().to_string(), |
| 183 | + })) |
| 184 | +} |
| 185 | + |
| 186 | +async fn upgrade_schema_from_uploaded_file( |
| 187 | + ctx: &dal::DalContext, |
| 188 | + pkg: &SiPkg, |
| 189 | + schema: Schema, |
| 190 | +) -> SchemaResult<Vec<SchemaVariantId>> { |
| 191 | + // Import and update funcs from uploaded pkg |
| 192 | + let mut thing_map = import_funcs_for_module_update(ctx, pkg.funcs()?) |
| 193 | + .await |
| 194 | + .map_err(SchemaError::Pkg)?; |
| 195 | + |
| 196 | + // Get specs from uploaded pkg |
| 197 | + let pkg_schemas = pkg.schemas()?; |
| 198 | + let schema_spec = pkg_schemas |
| 199 | + .first() |
| 200 | + .ok_or(SchemaError::PkgFileError("Pkg has no schemas"))?; |
| 201 | + let variants = schema_spec.variants()?; |
| 202 | + let variant_spec = variants |
| 203 | + .first() |
| 204 | + .ok_or(SchemaError::PkgFileError("Schema has no variants"))?; |
| 205 | + |
| 206 | + // Create new variant from uploaded pkg |
| 207 | + let new_variant = import_schema_variant( |
| 208 | + ctx, |
| 209 | + &schema, |
| 210 | + schema_spec.clone(), |
| 211 | + variant_spec, |
| 212 | + None, |
| 213 | + &mut thing_map, |
| 214 | + None, |
| 215 | + ) |
| 216 | + .await |
| 217 | + .map_err(SchemaError::Pkg)?; |
| 218 | + |
| 219 | + // Set as new default |
| 220 | + schema |
| 221 | + .set_default_variant_id(ctx, new_variant.id()) |
| 222 | + .await |
| 223 | + .map_err(|e| SchemaError::Pkg(e.into()))?; |
| 224 | + |
| 225 | + Ok(vec![new_variant.id()]) |
| 226 | +} |
0 commit comments