[Databricks] API를 이용한 Token정보 가져오기 (Python)

 

안녕하세요. 오랜만에 Databricks 관련 글을 올립니다. 오랜만에 쓰는 글인데 심화적인 글은 아니고 간단한 API를 이용한 내용입니다. Databricks에서는 AWS나 GCP 같은 플랫폼에서와 같이 Key 형태의 Token을 제공합니다. Token을 이용해서 여러 가지 Databricks의 모니터링을 할 때 Token을 이용해서 코드에 삽입하여 사용합니다. 

 

제가 만든 수많은 모니터링 시스템에서 필요한 것이 Databricks의 Token입니다. 모니터링을 하려면 AWS에서는 Access key 가 필요하고 GCP에서는 key file이 필요하듯이 여기서도 Token이라는 것이 존재 합니다.

 

Authenticate with Databricks personal access tokens (legacy) | Databricks on AWS

Authenticate with Databricks personal access tokens (legacy) | Databricks on AWS

 

하지만 이 Token에도 관리가 필요합니다. 특히 최근에 많은 이슈가 되고 있는 보안을 위해서 Token을 주기적으로 교체하거나 수명주기를 관리하여 Token을 관리해야 합니다. 

 

Token 확인 방법

Databricks Token의 경우 위의 스크린숏과 같이 우측 상단의 워크스페이스명을 누르면 나오는 곳의 가장 아래에 "계정 관리"부분을 클릭하면 나오는 페이지로 접속 후 아래와 같이 확인 가능 합니다.

 

 

여기서 Token을 생성하는 방법이나 각각의 워크스페이스별로 보는 방법은 다루지 않겠습니다. (글 목적과 맞지 않으므로)

---

Token 정보 가져오기

Databricks의 Token 정보의 경우 다음의 API를 통해서 가져올 수 있습니다. 

List all tokens | Token management API | REST API reference | Databricks on AWS

Databricks REST API reference

 

Token의 정보를 가져올 수 있으며 결과값 예시는 다음과 같습니다. 

{
  "token_infos": [
    {
      "comment": "This is for the ABC department automation scripts.",
      "created_by_id": 202480738464078,
      "created_by_username": "jsmith@example.com",
      "creation_time": 1580265020299,
      "expiry_time": 1580265020299,
      "last_used_day": 1710374400000,
      "owner_id": 202480738464078,
      "token_id": "5684c955822ac792a51ae2aeb80190f13457bab3e2e2934c133a08b38454816c",
      "workspace_id": 5750939285643220
    }
  ]
}

 

다음의 코드로 Token의 정보를 가져올수 있습니다. info_list에 보시면 list형태로 넣은 이유는 list에 여러 개의 워크스페이스를 넣어서 한 번에 정보를 전부 가져오기 위함입니다. (각 워크스페이스마다 Token을 생성하여 각 워크스페이스마다 API를 실행해야지 모든 Token정보를 획득 가능함)

import requests
from datetime import datetime,timedelta

info_list = [
             {'account':'Please enter your account.','domain':'Please enter your domain.','token':'Please enter your token.'}
             ]

for info_row in info_list:
    account = info_row.get('account')
    domain = info_row.get('domain')
    token = info_row.get('token')

    print(f'Collect the corresponding {account}')

    url = f"https://{domain}/api/2.0/token-management/tokens"
    headers = {"Authorization": f"Bearer {token}" }
    response = requests.get(url=url, headers=headers)
    result = response.json()
    token_list = result.get('token_infos')

    if token_list == []:
        pass

    else:
        for token in token_list:
            print(token)

 

위의 코드를 이용해서 Token 정보를 가져오고 저희 같은 경우 다음과 같이 Teams 메신저를 통해서 알림을 받고 있습니다. 이렇게 Token의 만료기간을 알려서 사용자가 갑자기 Token이 안 되는 상황을 막고 사전에 이를 방지하기 위함입니다.