feat: integrate unity catalog with datafusion #1338

roeap · 2023-05-06T16:59:15Z

Description

This is a follow-up from #1331, where we added support for the unity catalog and does a couple of things:

Adds a port of the client implementation from object_store to communicate with unity catalog.
refactors and extends the current client to accept configuration much in the way we do for the storage backends
adds additional (so far only Azure) credentials to authorize requests
adds an integration with datafusion's SchemaProvider and CatalogProvider
add unit tests for various components.

At first glance porting the whole client implementation may seem excessive, but my hope is that with this we can eventually replace the rusoto dependency from the glue catalog, as well as implement the LogStore locking mechanisms suggested by @nkarpov while keeping everything SDK-free :).

@nohajc - unfortunately I do not have access to a unity catalog instance right now. Would it be possible for you to validate, that the requests are still working?

Related Issue(s)

No issues

nohajc · 2023-05-06T17:45:49Z

@nohajc - unfortunately I do not have access to a unity catalog instance right now. Would it be possible for you to validate, that the requests are still working?

Sure, I will gladly test it.

nohajc · 2023-05-06T20:44:39Z

rust/src/data_catalog/unity/credential.rs

+    pub const AZURE_PUBLIC_CLOUD: &str = "https://login.microsoftonline.com";
+}
+
+/// Trait for providing authrization tokens for catalog requests


Typo in authrization

nohajc · 2023-05-06T23:40:49Z

rust/src/data_catalog/unity/mod.rs

+                    .map(|c| c.as_ref().to_string())
+                    .unwrap_or("main".into()),


You should be able to do this instead:

Suggested change

.map(|c| c.as_ref().to_string())

.unwrap_or("main".into()),

.as_ref()

.map(|c| c.as_ref())

.unwrap_or("main"),

nohajc · 2023-05-07T01:45:45Z

rust/src/data_catalog/unity/models.rs

+    Success(Table),
+    /// Error response
+    Error(ErrorResponse),
+}


I tested it and unfortunately ~~Error(ErrorResponse) is not equivalent to Error { error_code: String, message: String }~~, so right now it fails to parse (the same for Success(Table).

~~So far I haven't found any other way other than defining the structs inline, like it was before.~~

~~I thought the following might work but it didn't.~~

#[derive(Deserialize)] #[serde(untagged)] pub enum GetTableResponse { /// Successful response Success { /// Table information #[serde(flatten)] table: Table, }, /// Error response Error { /// Error information #[serde(flatten)] err: ErrorResponse, }, }

hmm .. if we have to we'll do it, but this makes the types a bit annoying to work with.

The example responses from the docs parse fine in the tests though. would it be possible to provide a response body that doe snot parse? Admittedly, the error case I just inferred from your code :).

Ok, now I'm even more confused. From the test, it looks like it should work. I don't know what serde is doing but it didn't parse the real response before I changed the definition back.

I'll try to write a test case to reproduce the problem.

So, the error cause was different than I thought (see the comment in struct Table). At least the definitions don't have to be inlined.

nohajc · 2023-05-07T01:50:10Z

rust/src/data_catalog/unity/mod.rs

+            }
+            "access_token" | "unity_access_token" | "databricks_access_token" => {
+                Ok(UnityCatalogConfigKey::AccessToken)
+            }


Seems like this should also match AZURE_TENANT_ID, AZURE_CLIENT_ID and AZURE_CLIENT_SECRET, right?
Btw, I wouldn't use the unity_ alias for Azure tenant and client because it's not Databrick/Unity specific.

I have been going back and forth on this and hope to experiment a bit here.

The main nuisance I encountered was that in scenarios where you are juggling many connections sometime we pick up things from the environment that aren't intended.

Generally speaking though yes, these are quite common and users would likely expect this to be picked up.

nohajc · 2023-05-07T09:21:22Z

rust/src/data_catalog/unity/models.rs

+    pub updated_by: String,
+
+    /// List of schemes whose objects can be referenced without qualification.
+    pub sql_path: String,


This is the culprit. My Unity Catalog returns response without any sql_path.

Suggested change

pub sql_path: String,

#[serde(default)]

pub sql_path: String,

I fixed it like this but looking at the documentation, it seems like most of the response fields could be optional (even though it doesn't make much sense in some cases).

nohajc · 2023-05-07T09:30:13Z

rust/src/data_catalog/unity/models.rs

+        "created_by": "string",
+        "name": "table_name",
+        "updated_by": "string",
+        "sql_path": "string",


We could add a test case with omitted sql_path but who knows which other fields can be missing from the response under different circumstances...

I guess we can just make a bunch of these fields optional. the ones we actually need are quite limited and should always be available - like name etc..

WE now just derive default for Schema and Table

I rechecked with the latest changes and deriving Default for the struct does not change the deserializer behavior.
The #[serde(default)] macro still needs to be used either per field or on the whole struct making all the fields optional with empty default. An alternative would be to use Option<String> where it makes sense.

nohajc · 2023-05-07T09:40:58Z

@roeap So far I only verified it works with Databricks personal access token but if I figure out which AAD roles are needed, I could check the OAuth flows (+ azure cli) as well.

roeap · 2023-05-07T20:38:52Z

So far I only verified it works with Databricks personal access token but if I figure out which AAD roles are needed, I could check the OAuth flows (+ azure cli) as well.

That would be awesome. Haven't worked with this yet. but from the docs i gather you don't really need specific roles, but rather "Just" need the AAD object user to be available / assigned in the unity catalog.

jordandakota · 2023-05-10T01:13:15Z

That would be awesome. Haven't worked with this yet. but from the docs i gather you don't really need specific roles, but rather "Just" need the AAD object user to be available / assigned in the unity catalog.

This is correct. If you have an AAD bearer token with scope set to the azure databricks application you can do everything a databricks personal access token can do. The scope being -> 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d

rust/src/data_catalog/unity/mod.rs

nohajc · 2023-05-13T14:01:16Z

rust/src/data_catalog/unity/mod.rs

+                Box::new(ClientSecretOAuthProvider::new(
+                    client_id,
+                    client_secret,
+                    authority_id,


Using Azure service principal account works (specifying Azure tenant ID as authority_id and account appId with password as client_id and client_secret respectively).

nohajc · 2023-05-13T14:04:36Z

That would be awesome. Haven't worked with this yet. but from the docs i gather you don't really need specific roles, but rather "Just" need the AAD object user to be available / assigned in the unity catalog.

Successfully tested service principal (client ID / client secret) and Azure CLI credentials. I don't have an opportunity to try with managed identity or workload identity right now...

roeap · 2023-05-13T16:19:07Z

Thanks for testing this @nohajc!

I don't have an opportunity to try with managed identity or workload identity right now...

I guess we know now, that we are using the right scopes, the rest is just how we get the AAD token, that we have validated elsewhere, so we should be good 🤞.

nohajc · 2023-05-13T16:19:41Z

(I still need to use #[serde(default)] on the sql_path response field.)

roeap · 2023-05-13T16:28:32Z

(I still need to use #[serde(default)] on the sql_path response field.)

on it :) - applying some fixes and adding datafusion tests.

Co-authored-by: nohajc <[email protected]>

ion-elgreco · 2023-08-15T13:33:49Z

@roeap Is this ready to get merged?

…he workspace

github-actions bot added binding/python Issues for the Python package binding/rust Issues for the Rust crate rust labels May 6, 2023

roeap force-pushed the unity-catalog branch from 29f2b9d to 086d9e9 Compare May 6, 2023 18:08

nohajc reviewed May 6, 2023

View reviewed changes

nohajc reviewed May 7, 2023

View reviewed changes

nohajc reviewed May 13, 2023

View reviewed changes

rust/src/data_catalog/unity/mod.rs Show resolved Hide resolved

nohajc reviewed May 13, 2023

View reviewed changes

roeap force-pushed the unity-catalog branch from 7003fcd to 047132f Compare May 13, 2023 15:59

roeap mentioned this pull request May 13, 2023

ci: disable full debug symbol generation #1364

Merged

roeap and others added 8 commits May 13, 2023 19:57

feat: extend unit catalog support

8649548

chore: draft datafusion integration

ac9bec1

fix: allow passing catalog options from python

ddc1098

chore: clippy

e7178f4

feat: add more azure credentials

c57d8db

fix: add defaults for return types

ab59f9a

fix: simpler defaults

05642da

Update rust/src/data_catalog/unity/mod.rs

a1ea5aa

Co-authored-by: nohajc <[email protected]>

Merge branch 'main' into unity-catalog

1020898

github-actions bot removed the documentation Improvements or additions to documentation label Jun 2, 2023

roeap added 2 commits June 2, 2023 08:35

fix: remove artifact

80b25ca

fix: errors after merge with main

83dcb6d

rtyler added 3 commits September 18, 2023 23:11

Merge remote-tracking branch 'upstream/main' into unity-catalog

719266e

Correct some merge related errors with redundant package names from t…

b2ba0d7

…he workspace

Address some latent clippy failures after merging main

b1fab03

rtyler marked this pull request as ready for review September 19, 2023 06:46

rtyler requested review from wjones127, fvaleye, rtyler, houqp, xianwill and mosyp as code owners September 19, 2023 06:46

Correct the incorrect documentation for Backoff

9119ee6

rtyler added this to the Rust v0.16 milestone Sep 19, 2023

rtyler approved these changes Sep 19, 2023

View reviewed changes

rtyler enabled auto-merge (rebase) September 19, 2023 15:19

Merge branch 'main' into unity-catalog

fff05c7

auto-merge was automatically disabled September 19, 2023 15:34
Rebase failed

rtyler enabled auto-merge (rebase) September 19, 2023 20:33

Merge branch 'main' into unity-catalog

51db313

auto-merge was automatically disabled September 19, 2023 20:51
Rebase failed

rtyler added 2 commits September 19, 2023 16:30

Merge remote-tracking branch 'upstream/main' into unity-catalog

645abb9

Merge remote-tracking branch 'roeap/unity-catalog' into unity-catalog

9b3a0e5

rtyler enabled auto-merge (rebase) September 19, 2023 23:33

auto-merge was automatically disabled September 19, 2023 23:48
Rebase failed

rtyler merged commit 6746dd4 into delta-io:main Sep 20, 2023
21 checks passed

roeap deleted the unity-catalog branch May 26, 2024 17:09

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: integrate unity catalog with datafusion #1338

feat: integrate unity catalog with datafusion #1338

roeap commented May 6, 2023

nohajc commented May 6, 2023

nohajc May 6, 2023

nohajc May 6, 2023

nohajc May 7, 2023 •

edited

Loading

roeap May 7, 2023

nohajc May 7, 2023

nohajc May 7, 2023

nohajc May 7, 2023

roeap May 7, 2023

nohajc May 7, 2023 •

edited

Loading

nohajc May 7, 2023

roeap May 7, 2023

roeap May 7, 2023

nohajc May 7, 2023

nohajc commented May 7, 2023 •

edited

Loading

roeap commented May 7, 2023

jordandakota commented May 10, 2023

nohajc May 13, 2023

nohajc commented May 13, 2023

roeap commented May 13, 2023

nohajc commented May 13, 2023

roeap commented May 13, 2023

ion-elgreco commented Aug 15, 2023

feat: integrate unity catalog with datafusion #1338

feat: integrate unity catalog with datafusion #1338

Conversation

roeap commented May 6, 2023

Description

Related Issue(s)

nohajc commented May 6, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nohajc May 7, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nohajc May 7, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nohajc commented May 7, 2023 • edited Loading

roeap commented May 7, 2023

jordandakota commented May 10, 2023

Choose a reason for hiding this comment

nohajc commented May 13, 2023

roeap commented May 13, 2023

nohajc commented May 13, 2023

roeap commented May 13, 2023

ion-elgreco commented Aug 15, 2023

nohajc May 7, 2023 •

edited

Loading

nohajc May 7, 2023 •

edited

Loading

nohajc commented May 7, 2023 •

edited

Loading