mirror of
https://github.com/neondatabase/neon.git
synced 2025-12-22 21:59:59 +00:00
30a7dd630c
## Problem `TYPE_CHECKING` is used inconsistently across Python tests. ## Summary of changes - Update `ruff`: 0.7.0 -> 0.11.2 - Enable TC (flake8-type-checking): https://docs.astral.sh/ruff/rules/#flake8-type-checking-tc - (auto)fix all new issues
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import StrEnum
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import jwt
|
|
|
|
if TYPE_CHECKING:
|
|
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))
|
|
|
|
|
|
class TokenScope(StrEnum):
|
|
ADMIN = "admin"
|
|
PAGE_SERVER_API = "pageserverapi"
|
|
GENERATIONS_API = "generations_api"
|
|
SAFEKEEPER_DATA = "safekeeperdata"
|
|
TENANT = "tenant"
|
|
SCRUBBER = "scrubber"
|
|
INFRA = "infra"
|