diff --git a/src/firebolt/service/V1/base.py b/src/firebolt/service/V1/base.py index 24b6caea66c..2039f320df6 100644 --- a/src/firebolt/service/V1/base.py +++ b/src/firebolt/service/V1/base.py @@ -1,11 +1,13 @@ -from typing import Optional +from typing import TYPE_CHECKING, Optional from firebolt.client import ClientV1 as Client -from firebolt.service.manager import ResourceManager + +if TYPE_CHECKING: + from firebolt.service.manager import ResourceManager class BaseService: - def __init__(self, resource_manager: ResourceManager): + def __init__(self, resource_manager: "ResourceManager"): self.resource_manager = resource_manager @property diff --git a/src/firebolt/service/V1/region.py b/src/firebolt/service/V1/region.py index 710894b8969..f688e26fa3e 100644 --- a/src/firebolt/service/V1/region.py +++ b/src/firebolt/service/V1/region.py @@ -1,14 +1,16 @@ -from typing import Dict, List +from typing import TYPE_CHECKING, Dict, List from firebolt.model.V1.region import Region, RegionKey -from firebolt.service.manager import ResourceManager from firebolt.service.V1.base import BaseService from firebolt.utils.urls import REGIONS_URL from firebolt.utils.util import cached_property +if TYPE_CHECKING: + from firebolt.service.manager import ResourceManager + class RegionService(BaseService): - def __init__(self, resource_manager: ResourceManager): + def __init__(self, resource_manager: "ResourceManager"): """ Service to manage AWS regions (us-east-1, etc) diff --git a/src/firebolt/service/manager.py b/src/firebolt/service/manager.py index 0b66c1d86bf..435b072ef50 100644 --- a/src/firebolt/service/manager.py +++ b/src/firebolt/service/manager.py @@ -14,7 +14,14 @@ ) from firebolt.common import Settings from firebolt.db import connect +from firebolt.service.V1.binding import BindingService +from firebolt.service.V1.database import DatabaseService as DatabaseServiceV1 +from firebolt.service.V1.engine import EngineService as EngineServiceV1 from firebolt.service.V1.provider import get_provider_id +from firebolt.service.V1.region import RegionService +from firebolt.service.V2.database import DatabaseService as DatabaseServiceV2 +from firebolt.service.V2.engine import EngineService as EngineServiceV2 +from firebolt.service.V2.instance_type import InstanceTypeService from firebolt.utils.util import fix_url_schema DEFAULT_TIMEOUT_SECONDS: int = 60 * 2 @@ -127,34 +134,25 @@ def __init__( def _init_services_v2(self) -> None: # avoid circular import - from firebolt.service.V2.database import DatabaseService - from firebolt.service.V2.engine import EngineService - from firebolt.service.V2.instance_type import InstanceTypeService # Cloud Platform Resources (AWS) self.instance_types = InstanceTypeService(resource_manager=self) # Firebolt Resources - self.databases = DatabaseService(resource_manager=self) - self.engines = EngineService(resource_manager=self) + self.databases = DatabaseServiceV2(resource_manager=self) + self.engines = EngineServiceV2(resource_manager=self) # Not applicable to V2 self.provider_id = None def _init_services_v1(self) -> None: - # avoid circular import - from firebolt.service.V1.binding import BindingService - from firebolt.service.V1.database import DatabaseService - from firebolt.service.V1.engine import EngineService - from firebolt.service.V1.region import RegionService - # Cloud Platform Resources (AWS) self.regions = RegionService(resource_manager=self) # type: ignore # Firebolt Resources self.bindings = BindingService(resource_manager=self) # type: ignore - self.engines = EngineService(resource_manager=self) # type: ignore - self.databases = DatabaseService(resource_manager=self) # type: ignore + self.engines = EngineServiceV1(resource_manager=self) # type: ignore + self.databases = DatabaseServiceV1(resource_manager=self) # type: ignore self.provider_id = get_provider_id(client=self._client) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index f1766cbb085..6a7f9d43ecf 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -1,5 +1,7 @@ from logging import getLogger from os import environ +from time import time +from typing import Optional from pytest import fixture, mark @@ -142,3 +144,20 @@ def engine_url() -> str: @fixture(scope="session") def stopped_engine_url() -> str: return must_env(STOPPED_ENGINE_URL_ENV) + + +@fixture(scope="function") +def minimal_time(): + limit: Optional[float] = None + + def setter(value): + nonlocal limit + limit = value + + start = time() + yield setter + end = time() + if limit is not None: + assert ( + end - start >= limit + ), f"Test took {end - start} seconds, less than {limit} seconds" diff --git a/tests/integration/dbapi/async/V1/test_queries_async.py b/tests/integration/dbapi/async/V1/test_queries_async.py index a5129c438b8..26dba4b90dd 100644 --- a/tests/integration/dbapi/async/V1/test_queries_async.py +++ b/tests/integration/dbapi/async/V1/test_queries_async.py @@ -1,6 +1,6 @@ from datetime import date, datetime from decimal import Decimal -from typing import Any, List +from typing import Any, Callable, List from pytest import fixture, mark, raises @@ -170,14 +170,19 @@ async def test_select( @mark.slow -@mark.timeout(timeout=600) +@mark.timeout(timeout=1000) async def test_long_query( connection: Connection, + minimal_time: Callable[[float], None], ) -> None: """AWS ALB TCP timeout set to 350; make sure we handle the keepalive correctly.""" + + # Fail test if it takes less than 350 seconds + minimal_time(350) + with connection.cursor() as c: await c.execute( - "SELECT checksum(*) FROM GENERATE_SERIES(1, 200000000000)", # approx 6m runtime + "SELECT checksum(*) FROM GENERATE_SERIES(1, 250000000000)", # approx 6m runtime ) data = await c.fetchall() assert len(data) == 1, "Invalid data size returned by fetchall" diff --git a/tests/integration/dbapi/async/V2/test_queries_async.py b/tests/integration/dbapi/async/V2/test_queries_async.py index 102d5ebd429..24a85373836 100644 --- a/tests/integration/dbapi/async/V2/test_queries_async.py +++ b/tests/integration/dbapi/async/V2/test_queries_async.py @@ -1,7 +1,7 @@ from datetime import date, datetime from decimal import Decimal from os import environ -from typing import List +from typing import Callable, List from pytest import fixture, mark, raises @@ -104,8 +104,12 @@ async def test_select( @mark.timeout(timeout=550) async def test_long_query( connection: Connection, + minimal_time: Callable[[float], None], ) -> None: """AWS ALB TCP timeout set to 350; make sure we handle the keepalive correctly.""" + + minimal_time(350) + with connection.cursor() as c: await c.execute( "SELECT checksum(*) FROM GENERATE_SERIES(1, 200000000000)", # approx 6m runtime diff --git a/tests/integration/dbapi/sync/V1/test_queries.py b/tests/integration/dbapi/sync/V1/test_queries.py index 780563f4c55..22be1209346 100644 --- a/tests/integration/dbapi/sync/V1/test_queries.py +++ b/tests/integration/dbapi/sync/V1/test_queries.py @@ -1,7 +1,7 @@ from datetime import date, datetime from decimal import Decimal from threading import Thread -from typing import Any, List +from typing import Any, Callable, List from pytest import fixture, mark, raises @@ -121,14 +121,18 @@ def test_select( @mark.slow -@mark.timeout(timeout=600) +@mark.timeout(timeout=1000) def test_long_query( - connection: Connection, + connection: Connection, minimal_time: Callable[[float], None] ) -> None: """AWS ALB TCP timeout set to 350, make sure we handle the keepalive correctly.""" + + # Fail test if it takes less than 350 seconds + minimal_time(350) + with connection.cursor() as c: c.execute( - "SELECT checksum(*) FROM GENERATE_SERIES(1, 200000000000)", # approx 6m runtime + "SELECT checksum(*) FROM GENERATE_SERIES(1, 250000000000)", # approx 6m runtime ) data = c.fetchall() assert len(data) == 1, "Invalid data size returned by fetchall" diff --git a/tests/integration/dbapi/sync/V2/test_queries.py b/tests/integration/dbapi/sync/V2/test_queries.py index 5bfdfc6acc6..68325205efc 100644 --- a/tests/integration/dbapi/sync/V2/test_queries.py +++ b/tests/integration/dbapi/sync/V2/test_queries.py @@ -2,7 +2,7 @@ from decimal import Decimal from os import environ from threading import Thread -from typing import Any, List +from typing import Any, Callable, List from pytest import fixture, mark, raises @@ -110,8 +110,12 @@ def test_select( @mark.timeout(timeout=550) def test_long_query( connection: Connection, + minimal_time: Callable[[float], None], ) -> None: - """AWS ALB TCP timeout set to 350, make sure we handle the keepalive correctly.""" + """AWS ALB TCP timeout set to 350; make sure we handle the keepalive correctly.""" + + minimal_time(350) + with connection.cursor() as c: c.execute( "SELECT checksum(*) FROM GENERATE_SERIES(1, 200000000000)", # approx 6m runtime