rigg_client/
arm.rs

1//! Azure Resource Manager client for discovering Search and Foundry services
2
3use reqwest::Client;
4use serde::Deserialize;
5use tracing::debug;
6
7use crate::auth::AzCliAuth;
8use crate::error::ClientError;
9
10const ARM_BASE_URL: &str = "https://management.azure.com";
11
12/// Azure Resource Manager client for subscription/service discovery
13pub struct ArmClient {
14    http: Client,
15    token: String,
16}
17
18/// Azure subscription
19#[derive(Debug, Clone, Deserialize)]
20#[serde(rename_all = "camelCase")]
21pub struct Subscription {
22    pub subscription_id: String,
23    pub display_name: String,
24    pub state: String,
25}
26
27impl std::fmt::Display for Subscription {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        write!(f, "{} ({})", self.display_name, self.subscription_id)
30    }
31}
32
33/// Azure AI Search service
34#[derive(Debug, Clone, Deserialize)]
35pub struct SearchService {
36    pub name: String,
37    pub location: String,
38    pub sku: SearchServiceSku,
39    #[serde(default)]
40    pub id: String,
41}
42
43#[derive(Debug, Clone, Deserialize)]
44pub struct SearchServiceSku {
45    pub name: String,
46}
47
48impl std::fmt::Display for SearchService {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        write!(
51            f,
52            "{} ({}, {})",
53            self.name,
54            self.location,
55            self.sku.name.to_uppercase()
56        )
57    }
58}
59
60/// Result of the discovery flow
61#[derive(Debug, Clone)]
62pub struct DiscoveredService {
63    pub name: String,
64    pub subscription_id: String,
65    pub location: String,
66}
67
68/// Azure AI Services account (kind=AIServices)
69#[derive(Debug, Clone, Deserialize)]
70pub struct AiServicesAccount {
71    pub name: String,
72    pub location: String,
73    #[serde(default)]
74    pub kind: String,
75    #[serde(default)]
76    pub id: String,
77    #[serde(default)]
78    pub properties: AiServicesAccountProperties,
79}
80
81#[derive(Debug, Clone, Default, Deserialize)]
82pub struct AiServicesAccountProperties {
83    /// Primary endpoint (e.g., "https://name.cognitiveservices.azure.com/")
84    #[serde(default)]
85    pub endpoint: Option<String>,
86}
87
88impl AiServicesAccount {
89    /// Derive the `.services.ai.azure.com` endpoint for the agents API.
90    ///
91    /// Extracts the custom subdomain from the ARM `properties.endpoint`
92    /// (which may differ from the resource name), then constructs the
93    /// AI services endpoint. Falls back to the resource name.
94    pub fn agents_endpoint(&self) -> String {
95        if let Some(ref endpoint) = self.properties.endpoint {
96            if let Some(subdomain) = extract_subdomain(endpoint) {
97                return format!("https://{}.services.ai.azure.com", subdomain);
98            }
99        }
100        format!("https://{}.services.ai.azure.com", self.name)
101    }
102}
103
104/// Extract the subdomain from an Azure endpoint URL.
105///
106/// `"https://my-svc.cognitiveservices.azure.com/"` → `"my-svc"`
107fn extract_subdomain(endpoint: &str) -> Option<&str> {
108    let host = endpoint.strip_prefix("https://")?.split('/').next()?;
109    host.split('.').next()
110}
111
112impl std::fmt::Display for AiServicesAccount {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        write!(f, "{} ({})", self.name, self.location)
115    }
116}
117
118/// Microsoft Foundry project (sub-resource of AI Services account)
119#[derive(Debug, Clone, Deserialize)]
120pub struct FoundryProject {
121    /// ARM name — may be "accountName/projectName" for sub-resources
122    #[serde(default)]
123    name: String,
124    pub location: String,
125    #[serde(default)]
126    pub id: String,
127    #[serde(default)]
128    pub properties: FoundryProjectProperties,
129}
130
131#[derive(Debug, Clone, Default, Deserialize)]
132#[serde(rename_all = "camelCase")]
133pub struct FoundryProjectProperties {
134    #[serde(default)]
135    pub display_name: String,
136}
137
138impl FoundryProject {
139    /// The project display name (human-friendly, e.g. "proj-default")
140    pub fn display_name(&self) -> &str {
141        if !self.properties.display_name.is_empty() {
142            &self.properties.display_name
143        } else {
144            // Fallback: parse from "account/project" ARM name
145            self.name.rsplit('/').next().unwrap_or(&self.name)
146        }
147    }
148}
149
150impl std::fmt::Display for FoundryProject {
151    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
152        write!(f, "{} ({})", self.display_name(), self.location)
153    }
154}
155
156/// Azure Storage account
157#[derive(Debug, Clone, Deserialize)]
158pub struct StorageAccount {
159    pub name: String,
160    pub location: String,
161    #[serde(default)]
162    pub id: String,
163}
164
165impl std::fmt::Display for StorageAccount {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167        write!(f, "{} ({})", self.name, self.location)
168    }
169}
170
171/// Storage account key
172#[derive(Debug, Clone, Deserialize)]
173struct StorageKey {
174    value: String,
175}
176
177/// Storage account key list response
178#[derive(Debug, Deserialize)]
179struct StorageKeyList {
180    keys: Vec<StorageKey>,
181}
182
183/// Azure OpenAI model deployment
184#[derive(Debug, Clone, Deserialize)]
185pub struct ModelDeployment {
186    pub name: String,
187    #[serde(default)]
188    pub properties: ModelDeploymentProperties,
189    #[serde(default)]
190    pub sku: ModelDeploymentSku,
191}
192
193#[derive(Debug, Clone, Default, Deserialize)]
194pub struct ModelDeploymentProperties {
195    #[serde(default)]
196    pub model: ModelDeploymentModel,
197}
198
199#[derive(Debug, Clone, Default, Deserialize)]
200pub struct ModelDeploymentModel {
201    #[serde(default)]
202    pub name: String,
203    #[serde(default)]
204    pub version: String,
205}
206
207#[derive(Debug, Clone, Default, Deserialize)]
208pub struct ModelDeploymentSku {
209    #[serde(default)]
210    pub name: String,
211    #[serde(default)]
212    pub capacity: u32,
213}
214
215impl std::fmt::Display for ModelDeployment {
216    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
217        write!(
218            f,
219            "{} ({}, {})",
220            self.name, self.properties.model.name, self.sku.name
221        )
222    }
223}
224
225/// A resource's managed identity block.
226#[derive(Debug, Clone)]
227pub struct ResourceIdentity {
228    /// `SystemAssigned`, `UserAssigned`, `SystemAssigned, UserAssigned`, or `None`.
229    pub kind: String,
230    /// System-assigned principal id, when enabled.
231    pub principal_id: Option<String>,
232    /// (resource id, principal id) of attached user-assigned identities.
233    pub user_assigned: Vec<(String, String)>,
234}
235
236impl ResourceIdentity {
237    /// All principal ids this resource can act as.
238    pub fn principal_ids(&self) -> Vec<&str> {
239        let mut ids: Vec<&str> = self.principal_id.iter().map(String::as_str).collect();
240        ids.extend(self.user_assigned.iter().map(|(_, p)| p.as_str()));
241        ids
242    }
243}
244
245/// Deterministic UUID-shaped name from a string (stable role-assignment names).
246fn deterministic_uuid(input: &str) -> String {
247    let mut h1: u64 = 0xcbf29ce484222325;
248    let mut h2: u64 = 0x9e3779b97f4a7c15;
249    for b in input.as_bytes() {
250        h1 ^= u64::from(*b);
251        h1 = h1.wrapping_mul(0x100000001b3);
252        h2 = h2.rotate_left(7) ^ u64::from(*b);
253        h2 = h2.wrapping_mul(0x2545f4914f6cdd1d);
254    }
255    let bytes = [h1.to_be_bytes(), h2.to_be_bytes()].concat();
256    format!(
257        "{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
258        bytes[0],
259        bytes[1],
260        bytes[2],
261        bytes[3],
262        bytes[4],
263        bytes[5],
264        bytes[6],
265        bytes[7],
266        bytes[8],
267        bytes[9],
268        bytes[10],
269        bytes[11],
270        bytes[12],
271        bytes[13],
272        bytes[14],
273        bytes[15]
274    )
275}
276
277/// ARM list response envelope
278#[derive(Debug, Deserialize)]
279struct ArmListResponse<T> {
280    value: Vec<T>,
281}
282
283impl ArmClient {
284    /// Create a new ARM client using Azure CLI credentials
285    pub fn new() -> Result<Self, ClientError> {
286        let token = AzCliAuth::get_arm_token()?;
287        let http = Client::builder()
288            .timeout(std::time::Duration::from_secs(30))
289            .build()?;
290
291        Ok(Self { http, token })
292    }
293
294    /// Read a resource's managed identity block: `GET {id}?api-version=...`.
295    pub async fn get_resource_identity(
296        &self,
297        resource_id: &str,
298        api_version: &str,
299    ) -> Result<Option<ResourceIdentity>, ClientError> {
300        let url = format!("{ARM_BASE_URL}{resource_id}?api-version={api_version}");
301        let response = self
302            .http
303            .get(&url)
304            .header("Authorization", format!("Bearer {}", self.token))
305            .send()
306            .await?;
307        let status = response.status();
308        if !status.is_success() {
309            let body = response.text().await?;
310            return Err(ClientError::from_response(status.as_u16(), &body));
311        }
312        let value: serde_json::Value = response.json().await?;
313        let Some(identity) = value.get("identity") else {
314            return Ok(None);
315        };
316        let kind = identity
317            .get("type")
318            .and_then(|t| t.as_str())
319            .unwrap_or("None")
320            .to_string();
321        let principal_id = identity
322            .get("principalId")
323            .and_then(|p| p.as_str())
324            .map(str::to_string);
325        let user_assigned = identity
326            .get("userAssignedIdentities")
327            .and_then(|u| u.as_object())
328            .map(|map| {
329                map.iter()
330                    .filter_map(|(id, v)| {
331                        v.get("principalId")
332                            .and_then(|p| p.as_str())
333                            .map(|p| (id.clone(), p.to_string()))
334                    })
335                    .collect()
336            })
337            .unwrap_or_default();
338        Ok(Some(ResourceIdentity {
339            kind,
340            principal_id,
341            user_assigned,
342        }))
343    }
344
345    /// Role definition IDs assigned to `principal_id` at (or inherited by) `scope`.
346    pub async fn list_role_assignments(
347        &self,
348        scope: &str,
349        principal_id: &str,
350    ) -> Result<Vec<String>, ClientError> {
351        let url = format!(
352            "{ARM_BASE_URL}{scope}/providers/Microsoft.Authorization/roleAssignments?api-version=2022-04-01&$filter=principalId%20eq%20'{principal_id}'"
353        );
354        let response = self
355            .http
356            .get(&url)
357            .header("Authorization", format!("Bearer {}", self.token))
358            .send()
359            .await?;
360        let status = response.status();
361        if !status.is_success() {
362            let body = response.text().await?;
363            return Err(ClientError::from_response(status.as_u16(), &body));
364        }
365        let value: serde_json::Value = response.json().await?;
366        Ok(value
367            .get("value")
368            .and_then(|v| v.as_array())
369            .map(|arr| {
370                arr.iter()
371                    .filter_map(|a| {
372                        a.get("properties")
373                            .and_then(|p| p.get("roleDefinitionId"))
374                            .and_then(|r| r.as_str())
375                            .map(str::to_string)
376                    })
377                    .collect()
378            })
379            .unwrap_or_default())
380    }
381
382    /// Create a role assignment for a principal at a scope.
383    pub async fn create_role_assignment(
384        &self,
385        scope: &str,
386        principal_id: &str,
387        role_definition_guid: &str,
388    ) -> Result<(), ClientError> {
389        let assignment_name =
390            deterministic_uuid(&format!("{scope}|{principal_id}|{role_definition_guid}"));
391        let url = format!(
392            "{ARM_BASE_URL}{scope}/providers/Microsoft.Authorization/roleAssignments/{assignment_name}?api-version=2022-04-01"
393        );
394        let sub = scope.split('/').nth(2).unwrap_or_default();
395        let body = serde_json::json!({
396            "properties": {
397                "roleDefinitionId": format!(
398                    "/subscriptions/{sub}/providers/Microsoft.Authorization/roleDefinitions/{role_definition_guid}"
399                ),
400                "principalId": principal_id,
401                "principalType": "ServicePrincipal"
402            }
403        });
404        let response = self
405            .http
406            .put(&url)
407            .header("Authorization", format!("Bearer {}", self.token))
408            .json(&body)
409            .send()
410            .await?;
411        let status = response.status();
412        // 409 = already exists → fine
413        if status.is_success() || status.as_u16() == 409 {
414            return Ok(());
415        }
416        let text = response.text().await?;
417        Err(ClientError::from_response(status.as_u16(), &text))
418    }
419
420    /// Enable a system-assigned managed identity on a resource (PATCH).
421    pub async fn enable_system_identity(
422        &self,
423        resource_id: &str,
424        api_version: &str,
425    ) -> Result<(), ClientError> {
426        let url = format!("{ARM_BASE_URL}{resource_id}?api-version={api_version}");
427        let response = self
428            .http
429            .patch(&url)
430            .header("Authorization", format!("Bearer {}", self.token))
431            .json(&serde_json::json!({"identity": {"type": "SystemAssigned"}}))
432            .send()
433            .await?;
434        let status = response.status();
435        if status.is_success() {
436            return Ok(());
437        }
438        let text = response.text().await?;
439        Err(ClientError::from_response(status.as_u16(), &text))
440    }
441
442    /// Find the full ARM resource id of a search service by name.
443    pub async fn find_search_service_id(&self, name: &str) -> Result<String, ClientError> {
444        for sub in self.list_subscriptions().await? {
445            for svc in self.list_search_services(&sub.subscription_id).await? {
446                if svc.name.eq_ignore_ascii_case(name) && !svc.id.is_empty() {
447                    return Ok(svc.id);
448                }
449            }
450        }
451        Err(ClientError::NotFound {
452            kind: "search service".to_string(),
453            name: name.to_string(),
454        })
455    }
456
457    /// The ARM bearer token this client authenticated with.
458    pub fn token(&self) -> &str {
459        &self.token
460    }
461
462    /// List subscriptions the user has access to
463    pub async fn list_subscriptions(&self) -> Result<Vec<Subscription>, ClientError> {
464        let url = format!("{}/subscriptions?api-version=2022-12-01", ARM_BASE_URL);
465        debug!("Listing subscriptions: {}", url);
466
467        let response = self
468            .http
469            .get(&url)
470            .header("Authorization", format!("Bearer {}", self.token))
471            .send()
472            .await?;
473
474        let status = response.status();
475        if !status.is_success() {
476            let body = response.text().await?;
477            return Err(ClientError::from_response(status.as_u16(), &body));
478        }
479
480        let result: ArmListResponse<Subscription> = response.json().await?;
481        // Only return enabled subscriptions
482        Ok(result
483            .value
484            .into_iter()
485            .filter(|s| s.state == "Enabled")
486            .collect())
487    }
488
489    /// List Azure AI Search services in a subscription
490    pub async fn list_search_services(
491        &self,
492        subscription_id: &str,
493    ) -> Result<Vec<SearchService>, ClientError> {
494        let url = format!(
495            "{}/subscriptions/{}/providers/Microsoft.Search/searchServices?api-version=2023-11-01",
496            ARM_BASE_URL, subscription_id
497        );
498        debug!("Listing search services: {}", url);
499
500        let response = self
501            .http
502            .get(&url)
503            .header("Authorization", format!("Bearer {}", self.token))
504            .send()
505            .await?;
506
507        let status = response.status();
508        if !status.is_success() {
509            let body = response.text().await?;
510            return Err(ClientError::from_response(status.as_u16(), &body));
511        }
512
513        let result: ArmListResponse<SearchService> = response.json().await?;
514        Ok(result.value)
515    }
516
517    /// Find the resource group of a search service by scanning the subscription.
518    ///
519    /// Returns the resource group name extracted from the service's ARM resource ID.
520    pub async fn find_resource_group(
521        &self,
522        subscription_id: &str,
523        service_name: &str,
524    ) -> Result<String, ClientError> {
525        let services = self.list_search_services(subscription_id).await?;
526
527        for svc in &services {
528            if svc.name.eq_ignore_ascii_case(service_name) {
529                // Parse resource group from ARM ID:
530                // /subscriptions/{sub}/resourceGroups/{rg}/providers/...
531                return parse_resource_group(&svc.id).ok_or_else(|| ClientError::Api {
532                    status: 0,
533                    message: format!("Could not parse resource group from ARM ID: {}", svc.id),
534                });
535            }
536        }
537
538        Err(ClientError::NotFound {
539            kind: "Search service".to_string(),
540            name: service_name.to_string(),
541        })
542    }
543
544    /// List Azure AI Services accounts in a subscription (filtered to kind=AIServices)
545    pub async fn list_ai_services_accounts(
546        &self,
547        subscription_id: &str,
548    ) -> Result<Vec<AiServicesAccount>, ClientError> {
549        let url = format!(
550            "{}/subscriptions/{}/providers/Microsoft.CognitiveServices/accounts?api-version=2024-10-01",
551            ARM_BASE_URL, subscription_id
552        );
553        debug!("Listing AI Services accounts: {}", url);
554
555        let response = self
556            .http
557            .get(&url)
558            .header("Authorization", format!("Bearer {}", self.token))
559            .send()
560            .await?;
561
562        let status = response.status();
563        if !status.is_success() {
564            let body = response.text().await?;
565            return Err(ClientError::from_response(status.as_u16(), &body));
566        }
567
568        let result: ArmListResponse<AiServicesAccount> = response.json().await?;
569        Ok(result
570            .value
571            .into_iter()
572            .filter(|a| a.kind.eq_ignore_ascii_case("AIServices"))
573            .collect())
574    }
575
576    /// Find the ARM resource id of ANY Microsoft.CognitiveServices account
577    /// by name — regardless of kind (AIServices, CognitiveServices, OpenAI,
578    /// ...). Unlike [`Self::list_ai_services_accounts`] (which serves
579    /// Foundry discovery and filters to kind AIServices), this covers e.g.
580    /// the plain CognitiveServices accounts skillsets use for enrichment.
581    pub async fn find_cognitive_account_id(&self, name: &str) -> Result<String, ClientError> {
582        for sub in self.list_subscriptions().await? {
583            let url = format!(
584                "{}/subscriptions/{}/providers/Microsoft.CognitiveServices/accounts?api-version=2024-10-01",
585                ARM_BASE_URL, sub.subscription_id
586            );
587            let response = self
588                .http
589                .get(&url)
590                .header("Authorization", format!("Bearer {}", self.token))
591                .send()
592                .await?;
593            let status = response.status();
594            if !status.is_success() {
595                continue; // no access to this subscription — keep looking
596            }
597            let result: ArmListResponse<AiServicesAccount> = response.json().await?;
598            for acct in result.value {
599                if acct.name.eq_ignore_ascii_case(name) && !acct.id.is_empty() {
600                    return Ok(acct.id);
601                }
602            }
603        }
604        Err(ClientError::NotFound {
605            kind: "Cognitive Services account".to_string(),
606            name: name.to_string(),
607        })
608    }
609
610    /// List Microsoft Foundry projects under a specific AI Services account.
611    ///
612    /// Projects are sub-resources at:
613    /// `Microsoft.CognitiveServices/accounts/{accountName}/projects`
614    ///
615    /// The `account_id` should be the full ARM resource ID of the account,
616    /// from which we extract the resource group.
617    pub async fn list_foundry_projects(
618        &self,
619        account: &AiServicesAccount,
620        subscription_id: &str,
621    ) -> Result<Vec<FoundryProject>, ClientError> {
622        let resource_group = parse_resource_group(&account.id).ok_or_else(|| ClientError::Api {
623            status: 0,
624            message: format!("Could not parse resource group from ARM ID: {}", account.id),
625        })?;
626
627        let url = format!(
628            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.CognitiveServices/accounts/{}/projects?api-version=2025-06-01",
629            ARM_BASE_URL, subscription_id, resource_group, account.name
630        );
631        debug!("Listing Foundry projects: {}", url);
632
633        let response = self
634            .http
635            .get(&url)
636            .header("Authorization", format!("Bearer {}", self.token))
637            .send()
638            .await?;
639
640        let status = response.status();
641        if !status.is_success() {
642            let body = response.text().await?;
643            return Err(ClientError::from_response(status.as_u16(), &body));
644        }
645
646        let result: ArmListResponse<FoundryProject> = response.json().await?;
647        Ok(result.value)
648    }
649
650    /// List storage accounts in a resource group.
651    pub async fn list_storage_accounts(
652        &self,
653        subscription_id: &str,
654        resource_group: &str,
655    ) -> Result<Vec<StorageAccount>, ClientError> {
656        let url = format!(
657            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.Storage/storageAccounts?api-version=2023-05-01",
658            ARM_BASE_URL, subscription_id, resource_group
659        );
660        debug!("Listing storage accounts: {}", url);
661
662        let response = self
663            .http
664            .get(&url)
665            .header("Authorization", format!("Bearer {}", self.token))
666            .send()
667            .await?;
668
669        let status = response.status();
670        if !status.is_success() {
671            let body = response.text().await?;
672            return Err(ClientError::from_response(status.as_u16(), &body));
673        }
674
675        let result: ArmListResponse<StorageAccount> = response.json().await?;
676        Ok(result.value)
677    }
678
679    /// List storage accounts across a whole subscription.
680    pub async fn list_storage_accounts_subscription(
681        &self,
682        subscription_id: &str,
683    ) -> Result<Vec<StorageAccount>, ClientError> {
684        let url = format!(
685            "{}/subscriptions/{}/providers/Microsoft.Storage/storageAccounts?api-version=2023-05-01",
686            ARM_BASE_URL, subscription_id
687        );
688        debug!("Listing storage accounts (subscription-wide): {}", url);
689
690        let response = self
691            .http
692            .get(&url)
693            .header("Authorization", format!("Bearer {}", self.token))
694            .send()
695            .await?;
696
697        let status = response.status();
698        if !status.is_success() {
699            let body = response.text().await?;
700            return Err(ClientError::from_response(status.as_u16(), &body));
701        }
702
703        let result: ArmListResponse<StorageAccount> = response.json().await?;
704        Ok(result.value)
705    }
706
707    /// Whether a storage account contains a blob container with this name
708    /// (checked via ARM with the caller's CLI token — no data-plane access
709    /// or account keys involved).
710    pub async fn storage_account_has_container(
711        &self,
712        account_id: &str,
713        container: &str,
714    ) -> Result<bool, ClientError> {
715        let url = format!(
716            "{ARM_BASE_URL}{account_id}/blobServices/default/containers/{container}?api-version=2023-05-01"
717        );
718        debug!("Checking container: {}", url);
719
720        let response = self
721            .http
722            .get(&url)
723            .header("Authorization", format!("Bearer {}", self.token))
724            .send()
725            .await?;
726
727        match response.status().as_u16() {
728            200 => Ok(true),
729            404 => Ok(false),
730            status => {
731                let body = response.text().await?;
732                Err(ClientError::from_response(status, &body))
733            }
734        }
735    }
736
737    /// Find every storage account (across all visible subscriptions) that
738    /// holds a blob container named `container`. Used to auto-construct
739    /// identity-based `ResourceId=` data-source connections — the user is
740    /// already logged in via Azure CLI, so rigg discovers instead of asking.
741    /// Accounts that fail the container check (e.g. insufficient RBAC on an
742    /// unrelated subscription) are skipped, not fatal.
743    pub async fn find_storage_accounts_with_container(
744        &self,
745        container: &str,
746    ) -> Result<Vec<StorageAccount>, ClientError> {
747        let mut matches = Vec::new();
748        for sub in self.list_subscriptions().await? {
749            let accounts = match self
750                .list_storage_accounts_subscription(&sub.subscription_id)
751                .await
752            {
753                Ok(a) => a,
754                Err(e) => {
755                    debug!(
756                        "skipping subscription {} ({}): {e}",
757                        sub.display_name, sub.subscription_id
758                    );
759                    continue;
760                }
761            };
762            for account in accounts {
763                if account.id.is_empty() {
764                    continue;
765                }
766                match self
767                    .storage_account_has_container(&account.id, container)
768                    .await
769                {
770                    Ok(true) => matches.push(account),
771                    Ok(false) => {}
772                    Err(e) => debug!("skipping account {} ({e})", account.name),
773                }
774            }
775        }
776        Ok(matches)
777    }
778
779    /// Get the primary access key for a storage account.
780    pub async fn get_storage_account_key(
781        &self,
782        subscription_id: &str,
783        resource_group: &str,
784        account_name: &str,
785    ) -> Result<String, ClientError> {
786        let url = format!(
787            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.Storage/storageAccounts/{}/listKeys?api-version=2023-05-01",
788            ARM_BASE_URL, subscription_id, resource_group, account_name
789        );
790        debug!("Getting storage account keys: {}", url);
791
792        let response = self
793            .http
794            .post(&url)
795            .header("Authorization", format!("Bearer {}", self.token))
796            .header("Content-Length", "0")
797            .send()
798            .await?;
799
800        let status = response.status();
801        if !status.is_success() {
802            let body = response.text().await?;
803            return Err(ClientError::from_response(status.as_u16(), &body));
804        }
805
806        let key_list: StorageKeyList = response.json().await?;
807        key_list
808            .keys
809            .into_iter()
810            .next()
811            .map(|k| k.value)
812            .ok_or_else(|| ClientError::Api {
813                status: 0,
814                message: "No keys found for storage account".to_string(),
815            })
816    }
817
818    /// List model deployments for an AI Services account.
819    pub async fn list_model_deployments(
820        &self,
821        account: &AiServicesAccount,
822        subscription_id: &str,
823    ) -> Result<Vec<ModelDeployment>, ClientError> {
824        let resource_group = parse_resource_group(&account.id).ok_or_else(|| ClientError::Api {
825            status: 0,
826            message: format!("Could not parse resource group from ARM ID: {}", account.id),
827        })?;
828
829        let url = format!(
830            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.CognitiveServices/accounts/{}/deployments?api-version=2024-10-01",
831            ARM_BASE_URL, subscription_id, resource_group, account.name
832        );
833        debug!("Listing model deployments: {}", url);
834
835        let response = self
836            .http
837            .get(&url)
838            .header("Authorization", format!("Bearer {}", self.token))
839            .send()
840            .await?;
841
842        let status = response.status();
843        if !status.is_success() {
844            let body = response.text().await?;
845            return Err(ClientError::from_response(status.as_u16(), &body));
846        }
847
848        let result: ArmListResponse<ModelDeployment> = response.json().await?;
849        Ok(result.value)
850    }
851
852    /// Create a model deployment on an AI Services account.
853    pub async fn create_model_deployment(
854        &self,
855        account: &AiServicesAccount,
856        subscription_id: &str,
857        deployment_name: &str,
858        model_name: &str,
859        model_version: &str,
860    ) -> Result<(), ClientError> {
861        let resource_group = parse_resource_group(&account.id).ok_or_else(|| ClientError::Api {
862            status: 0,
863            message: format!("Could not parse resource group from ARM ID: {}", account.id),
864        })?;
865
866        let url = format!(
867            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.CognitiveServices/accounts/{}/deployments/{}?api-version=2024-10-01",
868            ARM_BASE_URL, subscription_id, resource_group, account.name, deployment_name
869        );
870        debug!("Creating model deployment: {}", url);
871
872        let body = serde_json::json!({
873            "sku": {
874                "name": "GlobalStandard",
875                "capacity": 1
876            },
877            "properties": {
878                "model": {
879                    "format": "OpenAI",
880                    "name": model_name,
881                    "version": model_version
882                }
883            }
884        });
885
886        let response = self
887            .http
888            .put(&url)
889            .header("Authorization", format!("Bearer {}", self.token))
890            .json(&body)
891            .send()
892            .await?;
893
894        let status = response.status();
895        if !status.is_success() {
896            let body = response.text().await?;
897            return Err(ClientError::from_response(status.as_u16(), &body));
898        }
899
900        Ok(())
901    }
902
903    /// Build a full connection string for a storage account.
904    pub async fn get_storage_connection_string(
905        &self,
906        subscription_id: &str,
907        resource_group: &str,
908        account_name: &str,
909    ) -> Result<String, ClientError> {
910        let key = self
911            .get_storage_account_key(subscription_id, resource_group, account_name)
912            .await?;
913
914        Ok(format!(
915            "DefaultEndpointsProtocol=https;AccountName={};AccountKey={};EndpointSuffix=core.windows.net",
916            account_name, key
917        ))
918    }
919}
920
921/// Parse resource group from an ARM resource ID.
922///
923/// ARM IDs look like: `/subscriptions/{sub}/resourceGroups/{rg}/providers/...`
924fn parse_resource_group(arm_id: &str) -> Option<String> {
925    let parts: Vec<&str> = arm_id.split('/').collect();
926    for (i, part) in parts.iter().enumerate() {
927        if part.eq_ignore_ascii_case("resourceGroups")
928            || part.eq_ignore_ascii_case("resourcegroups")
929        {
930            return parts.get(i + 1).map(|s| s.to_string());
931        }
932    }
933    None
934}
935
936#[cfg(test)]
937mod tests {
938    use super::*;
939
940    #[test]
941    fn test_parse_resource_group() {
942        let id = "/subscriptions/abc-123/resourceGroups/my-rg/providers/Microsoft.Search/searchServices/my-svc";
943        assert_eq!(parse_resource_group(id), Some("my-rg".to_string()));
944    }
945
946    #[test]
947    fn test_parse_resource_group_case_insensitive() {
948        let id = "/subscriptions/abc/resourcegroups/MyRG/providers/Something";
949        assert_eq!(parse_resource_group(id), Some("MyRG".to_string()));
950    }
951
952    #[test]
953    fn test_parse_resource_group_missing() {
954        let id = "/subscriptions/abc/providers/Something";
955        assert_eq!(parse_resource_group(id), None);
956    }
957
958    #[test]
959    fn test_ai_services_account_display() {
960        let account = AiServicesAccount {
961            name: "my-ai-service".to_string(),
962            location: "eastus".to_string(),
963            kind: "AIServices".to_string(),
964            id: String::new(),
965            properties: AiServicesAccountProperties::default(),
966        };
967        assert_eq!(format!("{}", account), "my-ai-service (eastus)");
968    }
969
970    #[test]
971    fn test_agents_endpoint_from_arm_endpoint() {
972        let account = AiServicesAccount {
973            name: "irma-prod-foundry".to_string(),
974            location: "swedencentral".to_string(),
975            kind: "AIServices".to_string(),
976            id: String::new(),
977            properties: AiServicesAccountProperties {
978                endpoint: Some("https://custom-subdomain.cognitiveservices.azure.com/".to_string()),
979            },
980        };
981        assert_eq!(
982            account.agents_endpoint(),
983            "https://custom-subdomain.services.ai.azure.com"
984        );
985    }
986
987    #[test]
988    fn test_agents_endpoint_fallback_to_name() {
989        let account = AiServicesAccount {
990            name: "irma-prod-foundry".to_string(),
991            location: "swedencentral".to_string(),
992            kind: "AIServices".to_string(),
993            id: String::new(),
994            properties: AiServicesAccountProperties::default(),
995        };
996        assert_eq!(
997            account.agents_endpoint(),
998            "https://irma-prod-foundry.services.ai.azure.com"
999        );
1000    }
1001
1002    #[test]
1003    fn test_extract_subdomain() {
1004        assert_eq!(
1005            extract_subdomain("https://my-svc.cognitiveservices.azure.com/"),
1006            Some("my-svc")
1007        );
1008        assert_eq!(
1009            extract_subdomain("https://custom.services.ai.azure.com"),
1010            Some("custom")
1011        );
1012        assert_eq!(extract_subdomain("not-a-url"), None);
1013    }
1014
1015    #[test]
1016    fn test_foundry_project_display_with_display_name() {
1017        let project = FoundryProject {
1018            name: "my-account/my-project".to_string(),
1019            location: "westus2".to_string(),
1020            id: String::new(),
1021            properties: FoundryProjectProperties {
1022                display_name: "my-project".to_string(),
1023            },
1024        };
1025        assert_eq!(format!("{}", project), "my-project (westus2)");
1026        assert_eq!(project.display_name(), "my-project");
1027    }
1028
1029    #[test]
1030    fn test_model_deployment_display() {
1031        let deployment = ModelDeployment {
1032            name: "gpt-4o-mini".to_string(),
1033            properties: ModelDeploymentProperties {
1034                model: ModelDeploymentModel {
1035                    name: "gpt-4o-mini".to_string(),
1036                    version: "2024-07-18".to_string(),
1037                },
1038            },
1039            sku: ModelDeploymentSku {
1040                name: "GlobalStandard".to_string(),
1041                capacity: 1,
1042            },
1043        };
1044        assert_eq!(
1045            format!("{}", deployment),
1046            "gpt-4o-mini (gpt-4o-mini, GlobalStandard)"
1047        );
1048    }
1049
1050    #[test]
1051    fn test_foundry_project_display_name_fallback() {
1052        let project = FoundryProject {
1053            name: "my-account/proj-default".to_string(),
1054            location: "swedencentral".to_string(),
1055            id: String::new(),
1056            properties: FoundryProjectProperties::default(),
1057        };
1058        assert_eq!(project.display_name(), "proj-default");
1059        assert_eq!(format!("{}", project), "proj-default (swedencentral)");
1060    }
1061}
1062
1063#[cfg(test)]
1064mod identity_tests {
1065    use super::*;
1066
1067    #[test]
1068    fn deterministic_uuid_stable_and_shaped() {
1069        let a = deterministic_uuid("scope|principal|role");
1070        let b = deterministic_uuid("scope|principal|role");
1071        let c = deterministic_uuid("scope|principal|other-role");
1072        assert_eq!(a, b);
1073        assert_ne!(a, c);
1074        assert_eq!(a.len(), 36);
1075        assert_eq!(a.chars().filter(|ch| *ch == '-').count(), 4);
1076    }
1077
1078    #[test]
1079    fn resource_identity_principal_ids() {
1080        let id = ResourceIdentity {
1081            kind: "SystemAssigned, UserAssigned".into(),
1082            principal_id: Some("sys".into()),
1083            user_assigned: vec![("id1".into(), "ua1".into())],
1084        };
1085        assert_eq!(id.principal_ids(), vec!["sys", "ua1"]);
1086    }
1087}
rigg_client/arm.rs

rigg_client/
arm.rs
