Add token refresh for UC Lineage ingestion using system tables

[Anshul759]

If using system tables to query lineage, the sql connection fails.
Reason - sql connection requires PAT to connect to databricks and run queries. As the token is valid for only 60mins, and some of the ingestions run for long duration, the sql connection expires and the process fails with authentication error.
Fix - To use azure credentials to generate tokens. Have updated the code to generate PAT using azure credentials. The validity of PAT is 60 mins here too, but we can refresh the PAT before expiring so that the process continues without any halt

[kyungsoo-datahub]

1. Could you add some unit tests?

2. For the future improvement, it would be nice if we can separate out the token provider part. For example:

class SqlConnectionTokenProvider(Protocol):                                                                             
    def get_token(self) -> str: ...                                                                                     
                                                                                                                        
class StaticTokenProvider:                                                                                              
    def __init__(self, token: str):                                                                                     
        self._token = token                                                                                             
                                                                                                                        
    def get_token(self) -> str:                                                                                         
        return self._token                                                                                              
                                                                                                                        
class AzureTokenProvider:
        ...

[kyungsoo-datahub]

Do we want to refresh within 10 mins before expiration? Line 414 says that it would be within 5 mins.

[Anshul759]

Have used authenticate function, so removed this part of the code

[kyungsoo-datahub]

It would be good to be debug level logging.

[Anshul759]

Done

[kyungsoo-datahub]

Based on the brief investigation, recreating client is redundant. authenticate() will refresh token if it's expired. Please, check this.

authenticate() -> _header_factory() (config.py) -> OAuthCredentialsProvider.call() (credentials_provider.py) -> refreshed_headers() (credentials_provider.py) -> inner.token() (credentials_provider.py) -> token() (oath.py) -> _blocking_token() (oath.py) -> refresh() (oath.py)

Please, check token() and _blcoking_token() functions: https://github.com/databricks/databricks-sdk-py/blob/main/databricks/sdk/oauth.py#L280-L333

[Anshul759]

Thanks for the suggestion. authenticate function does refresh the token when expired.
For Testing, ran pipeline for around 3 hours without any error -

[kyungsoo-datahub]

I'm wondering if it's safe with multi threads.

[Anshul759]

@kyungsoo-datahub

1. Have added unit test cases
2. Seperated the token management in a different class
3. Used inbuild authenticate function to refresh token instead of using custom logic