mirror of
https://github.com/neondatabase/neon.git
synced 2026-01-15 17:32:56 +00:00
30680d1f32
## Problem https://github.com/neondatabase/neon/pull/9596 did not update tests because that would've broken the compat tests. ## Summary of Changes Use infra scope where possible.
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
import jwt
|
|
|
|
from fixtures.common_types import TenantId
|
|
|
|
|
|
@dataclass
|
|
class AuthKeys:
|
|
priv: str
|
|
|
|
def generate_token(self, *, scope: TokenScope, **token_data: Any) -> str:
|
|
token_data = {key: str(val) for key, val in token_data.items()}
|
|
token = jwt.encode({"scope": scope, **token_data}, self.priv, algorithm="EdDSA")
|
|
# cast(Any, self.priv)
|
|
|
|
# jwt.encode can return 'bytes' or 'str', depending on Python version or type
|
|
# hinting or something (not sure what). If it returned 'bytes', convert it to 'str'
|
|
# explicitly.
|
|
if isinstance(token, bytes):
|
|
token = token.decode()
|
|
|
|
return token
|
|
|
|
def generate_pageserver_token(self) -> str:
|
|
return self.generate_token(scope=TokenScope.PAGE_SERVER_API)
|
|
|
|
def generate_safekeeper_token(self) -> str:
|
|
return self.generate_token(scope=TokenScope.SAFEKEEPER_DATA)
|
|
|
|
# generate token giving access to only one tenant
|
|
def generate_tenant_token(self, tenant_id: TenantId) -> str:
|
|
return self.generate_token(scope=TokenScope.TENANT, tenant_id=str(tenant_id))
|
|
|
|
|
|
# TODO: Replace with `StrEnum` when we upgrade to python 3.11
|
|
class TokenScope(str, Enum):
|
|
ADMIN = "admin"
|
|
PAGE_SERVER_API = "pageserverapi"
|
|
GENERATIONS_API = "generations_api"
|
|
SAFEKEEPER_DATA = "safekeeperdata"
|
|
TENANT = "tenant"
|
|
SCRUBBER = "scrubber"
|
|
INFRA = "infra"
|