Improve some typing related to performance testing for LR

Signed-off-by: Tristan Partin <tristan@neon.tech>
This commit is contained in:
Tristan Partin
2024-11-04 15:52:01 -06:00
committed by GitHub
parent ee68bbf6f5
commit 34812a6aab
2 changed files with 25 additions and 13 deletions

View File

@@ -1,7 +1,7 @@
from __future__ import annotations
import time
from typing import TYPE_CHECKING, cast
from typing import TYPE_CHECKING, cast, final
import requests
@@ -261,17 +261,22 @@ class NeonAPI:
time.sleep(0.5)
@final
class NeonApiEndpoint:
def __init__(self, neon_api: NeonAPI, pg_version: PgVersion, project_id: Optional[str]):
self.neon_api = neon_api
self.project_id: str
self.endpoint_id: str
self.connstr: str
if project_id is None:
project = neon_api.create_project(pg_version)
neon_api.wait_for_operation_to_finish(project["project"]["id"])
neon_api.wait_for_operation_to_finish(cast("str", project["project"]["id"]))
self.project_id = project["project"]["id"]
self.endpoint_id = project["endpoints"][0]["id"]
self.connstr = project["connection_uris"][0]["connection_uri"]
self.pgbench_env = connection_parameters_to_env(
project["connection_uris"][0]["connection_parameters"]
cast("dict[str, str]", project["connection_uris"][0]["connection_parameters"])
)
self.is_new = True
else: