Skip to content

Commit 18ed86c

Browse files
SmritiAgrawal04Smriti Agrawal
SmritiAgrawal04
and
Smriti Agrawal
authored
Whitelisting Onelake API & Workspace PL FQDNs (#1) (#711)
* Whitelisting Onelake API & Workspace PL FQDNs * Addressing comments to whitelist api-onelake fqdns and add UTs * Reverting the mistakenly modified UT * Adding validation for xy in WS-PL URL & case insensitive regex * Making regex an optional dependency * Eliminating the use of regex for ws-pl fqdn matching * Adding a separte method 'azure_test_workspace_private_link' for ws-pl tests * Organising private link UTs * Adding integration test for WSPL & rebasing with unpstream * Whitelisting Onelake API & Workspace PL FQDNs * Addressing comments to whitelist api-onelake fqdns and add UTs * Reverting the mistakenly modified UT * Adding validation for xy in WS-PL URL & case insensitive regex * Making regex an optional dependency * Eliminating the use of regex for ws-pl fqdn matching * Adding a separte method 'azure_test_workspace_private_link' for ws-pl tests * Organising private link UTs * Fixing cargo fmt * Fixing Clippy * Fixing Clippy * Fixing Clippy * Modified Integration test to pass workspace & artifact id within URL only * Fixing Clippy * Adding WSPL support for ABFSS scheme as well * Adding comments for readability * Adding comments for readability --------- Co-authored-by: Smriti Agrawal <smagrawal@microsoft.com>
1 parent af634d0 commit 18ed86c

2 files changed

Lines changed: 163 additions & 0 deletions

File tree

‎src/azure/builder.rs‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -666,10 +666,26 @@ impl MicrosoftAzureBuilder {
666666
self.container_name = Some(validate(host)?);
667667
} else {
668668
match host.split_once('.') {
669+
// Workspace-level Private Link detection
670+
// "{workspaceid}.z??.(onelake|dfs|blob).fabric.microsoft.com"
671+
Some((workspaceid, rest))
672+
if rest.starts_with('z') && rest.ends_with("fabric.microsoft.com") =>
673+
{
674+
// Account name for WS-PL is two labels: "{workspaceid}.z{xy}"
675+
let (zone, _) = rest.split_once('.').unwrap_or((rest, ""));
676+
677+
self.account_name = Some(format!("{workspaceid}.{zone}"));
678+
self.endpoint = Some(format!("https://{}", host));
679+
680+
self.container_name = Some(validate(parsed.username())?);
681+
self.use_fabric_endpoint = true.into();
682+
}
683+
669684
Some((a, "dfs.core.windows.net")) | Some((a, "blob.core.windows.net")) => {
670685
self.account_name = Some(validate(a)?);
671686
self.container_name = Some(validate(parsed.username())?);
672687
}
688+
673689
Some((a, "dfs.fabric.microsoft.com"))
674690
| Some((a, "blob.fabric.microsoft.com")) => {
675691
self.account_name = Some(validate(a)?);
@@ -681,6 +697,30 @@ impl MicrosoftAzureBuilder {
681697
}
682698
}
683699
"https" => match host.split_once('.') {
700+
// Workspace-level Private Link detection
701+
// "{workspaceid}.z??.(onelake|dfs|blob).fabric.microsoft.com"
702+
Some((workspaceid, rest))
703+
if rest.starts_with('z') && rest.ends_with("fabric.microsoft.com") =>
704+
{
705+
// rest looks like: "z28.dfs.fabric.microsoft.com" / "z28.blob.fabric.microsoft.com" / etc.
706+
// Account name for WS-PL is two labels: "{workspaceid}.z{xy}"
707+
let (zone, _) = rest.split_once('.').unwrap_or((rest, ""));
708+
709+
self.account_name = Some(format!("{workspaceid}.{zone}"));
710+
self.endpoint = Some(format!("https://{}", host));
711+
712+
// Attempt to infer the container name from the URL
713+
let container = parsed.path_segments().unwrap().next().expect(
714+
"iterator always contains at least one string (which may be empty)",
715+
);
716+
717+
if !container.is_empty() {
718+
self.container_name = Some(validate(container)?);
719+
}
720+
721+
self.use_fabric_endpoint = true.into();
722+
}
723+
684724
Some((a, "dfs.core.windows.net")) | Some((a, "blob.core.windows.net")) => {
685725
self.account_name = Some(validate(a)?);
686726
let container = parsed.path_segments().unwrap().next().expect(
@@ -1204,6 +1244,17 @@ mod tests {
12041244
assert_eq!(builder.container_name.as_deref(), Some("container"));
12051245
assert!(builder.use_fabric_endpoint.get().unwrap());
12061246

1247+
let mut builder = MicrosoftAzureBuilder::new();
1248+
builder
1249+
.parse_url("https://onelake.dfs.fabric.microsoft.com/c047b3e3-4e89-407a-98d7-cf9949ae92a3/9f1a2b3c-4d5e-6f70-8a9b-c0d1e2f3a456.lakehouse/Files/tables/sales/data.parquet")
1250+
.unwrap();
1251+
assert_eq!(builder.account_name, Some("onelake".to_string()));
1252+
assert_eq!(
1253+
builder.container_name.as_deref(),
1254+
Some("c047b3e3-4e89-407a-98d7-cf9949ae92a3")
1255+
);
1256+
assert!(builder.use_fabric_endpoint.get().unwrap());
1257+
12071258
let mut builder = MicrosoftAzureBuilder::new();
12081259
builder
12091260
.parse_url("https://account.blob.fabric.microsoft.com/")
@@ -1235,6 +1286,77 @@ mod tests {
12351286
}
12361287
}
12371288

1289+
#[test]
1290+
fn azure_test_workspace_private_link() {
1291+
let test_cases: Vec<(&str, &str, Option<&str>)> = vec![
1292+
(
1293+
"https://Ab000000000000000000000000000000.zAb.dfs.fabric.microsoft.com/",
1294+
"ab000000000000000000000000000000.zab",
1295+
None,
1296+
),
1297+
(
1298+
"https://ab000000000000000000000000000000.zab.dfs.fabric.microsoft.com/",
1299+
"ab000000000000000000000000000000.zab",
1300+
None,
1301+
),
1302+
(
1303+
"https://c047b3e34e89407a98d7cf9949ae92a3.zc0.blob.fabric.microsoft.com/c047b3e3-4e89-407a-98d7-cf9949ae92a3/9f1a2b3c-4d5e-6f70-8a9b-c0d1e2f3a456/file",
1304+
"c047b3e34e89407a98d7cf9949ae92a3.zc0",
1305+
Some("c047b3e3-4e89-407a-98d7-cf9949ae92a3"),
1306+
),
1307+
(
1308+
"https://c047b3e34e89407a98d7cf9949ae92a3.zc0.dfs.fabric.microsoft.com/c047b3e3-4e89-407a-98d7-cf9949ae92a3/9f1a2b3c-4d5e-6f70-8a9b-c0d1e2f3a456/file",
1309+
"c047b3e34e89407a98d7cf9949ae92a3.zc0",
1310+
Some("c047b3e3-4e89-407a-98d7-cf9949ae92a3"),
1311+
),
1312+
(
1313+
"https://c047b3e34e89407a98d7cf9949ae92a3.zc0.onelake.fabric.microsoft.com/c047b3e3-4e89-407a-98d7-cf9949ae92a3/9f1a2b3c-4d5e-6f70-8a9b-c0d1e2f3a456/file",
1314+
"c047b3e34e89407a98d7cf9949ae92a3.zc0",
1315+
Some("c047b3e3-4e89-407a-98d7-cf9949ae92a3"),
1316+
),
1317+
(
1318+
"https://c047b3e34e89407a98d7cf9949ae92a3.zc0.w.api.fabric.microsoft.com/c047b3e3-4e89-407a-98d7-cf9949ae92a3/9f1a2b3c-4d5e-6f70-8a9b-c0d1e2f3a456/file",
1319+
"c047b3e34e89407a98d7cf9949ae92a3.zc0",
1320+
Some("c047b3e3-4e89-407a-98d7-cf9949ae92a3"),
1321+
),
1322+
(
1323+
"https://c047b3e34e89407a98d7cf9949ae92a3.zc0.c.api.fabric.microsoft.com/c047b3e3-4e89-407a-98d7-cf9949ae92a3/9f1a2b3c-4d5e-6f70-8a9b-c0d1e2f3a456/file",
1324+
"c047b3e34e89407a98d7cf9949ae92a3.zc0",
1325+
Some("c047b3e3-4e89-407a-98d7-cf9949ae92a3"),
1326+
),
1327+
(
1328+
"abfss://c047b3e34e89407a98d7cf9949ae92a3@c047b3e34e89407a98d7cf9949ae92a3.zc0.dfs.fabric.microsoft.com/9f1a2b3c-4d5e-6f70-8a9b-c0d1e2f3a456/file",
1329+
"c047b3e34e89407a98d7cf9949ae92a3.zc0",
1330+
Some("c047b3e34e89407a98d7cf9949ae92a3"),
1331+
),
1332+
(
1333+
"abfss://c047b3e34e89407a98d7cf9949ae92a3@c047b3e34e89407a98d7cf9949ae92a3.zc0.blob.fabric.microsoft.com/9f1a2b3c-4d5e-6f70-8a9b-c0d1e2f3a456/file",
1334+
"c047b3e34e89407a98d7cf9949ae92a3.zc0",
1335+
Some("c047b3e34e89407a98d7cf9949ae92a3"),
1336+
),
1337+
];
1338+
1339+
for (url, expected_account, expected_container) in &test_cases {
1340+
let mut builder = MicrosoftAzureBuilder::new();
1341+
builder.parse_url(url).unwrap();
1342+
1343+
assert_eq!(
1344+
builder.account_name.as_deref(),
1345+
Some(*expected_account),
1346+
"account mismatch for URL: {url}"
1347+
);
1348+
assert_eq!(
1349+
builder.container_name.as_deref(),
1350+
*expected_container,
1351+
"container mismatch for URL: {url}"
1352+
);
1353+
assert!(
1354+
builder.use_fabric_endpoint.get().unwrap(),
1355+
"use_fabric_endpoint not set for URL: {url}"
1356+
);
1357+
}
1358+
}
1359+
12381360
#[test]
12391361
fn azure_test_config_from_map() {
12401362
let azure_client_id = "object_store:fake_access_key_id";

‎src/azure/mod.rs‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,6 +378,47 @@ mod tests {
378378
}
379379
}
380380

381+
#[ignore = "Used for manual testing against a real Workspace Private Link Endpoint."]
382+
#[tokio::test]
383+
async fn azure_onelake_wspl_test() {
384+
maybe_skip_integration!();
385+
386+
let url =
387+
std::env::var("AZURE_ONELAKE_URL").expect("Set AZURE_ONELAKE_URL to a WS-PL FQDN");
388+
let parsed = url::Url::parse(&url).unwrap();
389+
390+
let path = match parsed.scheme() {
391+
"abfss" | "abfs" => {
392+
// abfss://<container>@<host>/<path...>
393+
// container is in username, entire path is the object path
394+
let segments: Vec<&str> = parsed.path_segments().unwrap().collect();
395+
Path::from(segments.join("/"))
396+
}
397+
_ => {
398+
// https://<host>/<container>/<path...>
399+
// first segment is container, rest is the object path
400+
let segments: Vec<&str> = parsed.path_segments().unwrap().collect();
401+
Path::from(segments[1..].join("/"))
402+
}
403+
};
404+
405+
let store = MicrosoftAzureBuilder::new()
406+
.with_url(&url)
407+
.with_bearer_token_authorization(
408+
std::env::var("AZURE_STORAGE_TOKEN").expect("Set AZURE_STORAGE_TOKEN"),
409+
)
410+
.build()
411+
.unwrap();
412+
413+
let data = Bytes::from("Hello OneLake WSPL");
414+
415+
store.put(&path, data.clone().into()).await.unwrap();
416+
let result = store.get(&path).await.unwrap();
417+
let loaded = result.bytes().await.unwrap();
418+
assert_eq!(data, loaded);
419+
store.delete(&path).await.unwrap();
420+
}
421+
381422
#[ignore = "Used for manual testing against a real storage account."]
382423
#[tokio::test]
383424
async fn test_user_delegation_key() {

0 commit comments

Comments
 (0)