Skip to content

Commit 1da1724

Browse files
Tranquility2Tranquility2
committed
Refactor generic sql
1 parent 03dd5e0 commit 1da1724

File tree

3 files changed

+86
-84
lines changed

3 files changed

+86
-84
lines changed

‎modules/generic/testcontainers/generic/sql.py‎

Lines changed: 3 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -4,88 +4,10 @@
44

55
from testcontainers.core.container import DockerContainer
66
from testcontainers.core.exceptions import ContainerStartException
7-
from testcontainers.core.utils import raise_for_deprecated_parameter
8-
from testcontainers.core.waiting_utils import WaitStrategy, WaitStrategyTarget
97

10-
logger = logging.getLogger(__name__)
11-
12-
ADDITIONAL_TRANSIENT_ERRORS = []
13-
try:
14-
from sqlalchemy.exc import DBAPIError
15-
16-
ADDITIONAL_TRANSIENT_ERRORS.append(DBAPIError)
17-
except ImportError:
18-
logger.debug("SQLAlchemy not available, skipping DBAPIError handling")
19-
SQL_TRANSIENT_EXCEPTIONS = (TimeoutError, ConnectionError, *ADDITIONAL_TRANSIENT_ERRORS)
20-
21-
22-
class ConnectWaitStrategy(WaitStrategy):
23-
"""
24-
Wait strategy that tests database connectivity until it succeeds or times out.
25-
26-
This strategy performs database connection testing using SQLAlchemy directly,
27-
handling transient connection errors and providing appropriate retry logic
28-
for database connectivity testing.
29-
"""
8+
from .sql_utils import SqlConnectWaitStrategy
309

31-
def __init__(self, transient_exceptions: Optional[tuple] = None):
32-
super().__init__()
33-
self.transient_exceptions = transient_exceptions or (TimeoutError, ConnectionError)
34-
35-
def wait_until_ready(self, container: WaitStrategyTarget) -> None:
36-
"""
37-
Test database connectivity with retry logic until it succeeds or times out.
38-
39-
Args:
40-
container: The SQL container that must have get_connection_url method
41-
42-
Raises:
43-
TimeoutError: If connection fails after timeout
44-
AttributeError: If container doesn't have get_connection_url method
45-
ImportError: If SQLAlchemy is not installed
46-
Exception: Any non-transient errors from connection attempts
47-
"""
48-
import time
49-
50-
if not hasattr(container, "get_connection_url"):
51-
raise AttributeError(f"Container {container} must have a get_connection_url method")
52-
53-
try:
54-
import sqlalchemy
55-
except ImportError as e:
56-
logger.error("SQLAlchemy is required for database connectivity testing")
57-
raise ImportError("SQLAlchemy is required for database containers") from e
58-
59-
start_time = time.time()
60-
61-
while True:
62-
if time.time() - start_time > self._startup_timeout:
63-
raise TimeoutError(
64-
f"Database connection failed after {self._startup_timeout}s timeout. "
65-
f"Hint: Check if the container is ready and the database is accessible."
66-
)
67-
68-
try:
69-
connection_url = container.get_connection_url()
70-
engine = sqlalchemy.create_engine(connection_url)
71-
72-
try:
73-
with engine.connect():
74-
logger.info("Database connection test successful")
75-
return
76-
except Exception as e:
77-
logger.debug(f"Database connection attempt failed: {e}")
78-
raise
79-
finally:
80-
engine.dispose()
81-
82-
except self.transient_exceptions as e:
83-
logger.debug(f"Connection attempt failed: {e}, retrying in {self._poll_interval}s...")
84-
except Exception as e:
85-
logger.error(f"Connection failed with non-transient error: {e}")
86-
raise
87-
88-
time.sleep(self._poll_interval)
10+
logger = logging.getLogger(__name__)
8911

9012

9113
class SqlContainer(DockerContainer):
@@ -128,8 +50,6 @@ def _create_connection_url(
12850
ValueError: If unexpected arguments are provided or required parameters are missing
12951
ContainerStartException: If container is not started
13052
"""
131-
if raise_for_deprecated_parameter(kwargs, "db_name", "dbname"):
132-
raise ValueError(f"Unexpected arguments: {','.join(kwargs)}")
13353

13454
if self._container is None:
13555
raise ContainerStartException("Container has not been started")
@@ -179,7 +99,7 @@ def start(self) -> "SqlContainer":
17999

180100
try:
181101
self._configure()
182-
self.waiting_for(ConnectWaitStrategy(SQL_TRANSIENT_EXCEPTIONS))
102+
self.waiting_for(SqlConnectWaitStrategy())
183103
super().start()
184104
self._transfer_seed()
185105
logger.info("Database container started successfully")

‎modules/generic/testcontainers/generic/sql_utils.py‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
import logging
2+
3+
from testcontainers.core.waiting_utils import WaitStrategy, WaitStrategyTarget
4+
5+
logger = logging.getLogger(__name__)
6+
7+
ADDITIONAL_TRANSIENT_ERRORS = []
8+
try:
9+
from sqlalchemy.exc import DBAPIError
10+
11+
ADDITIONAL_TRANSIENT_ERRORS.append(DBAPIError)
12+
except ImportError:
13+
logger.debug("SQLAlchemy not available, skipping DBAPIError handling")
14+
15+
16+
class SqlConnectWaitStrategy(WaitStrategy):
17+
"""
18+
Wait strategy that tests database connectivity until it succeeds or times out.
19+
20+
This strategy performs database connection testing using SQLAlchemy directly,
21+
handling transient connection errors and providing appropriate retry logic
22+
for database connectivity testing.
23+
"""
24+
25+
def __init__(self):
26+
super().__init__()
27+
self.transient_exceptions = (TimeoutError, ConnectionError, *ADDITIONAL_TRANSIENT_ERRORS)
28+
29+
def wait_until_ready(self, container: WaitStrategyTarget) -> None:
30+
"""
31+
Test database connectivity with retry logic until it succeeds or times out.
32+
33+
Args:
34+
container: The SQL container that must have get_connection_url method
35+
36+
Raises:
37+
TimeoutError: If connection fails after timeout
38+
AttributeError: If container doesn't have get_connection_url method
39+
ImportError: If SQLAlchemy is not installed
40+
Exception: Any non-transient errors from connection attempts
41+
"""
42+
import time
43+
44+
if not hasattr(container, "get_connection_url"):
45+
raise AttributeError(f"Container {container} must have a get_connection_url method")
46+
47+
try:
48+
import sqlalchemy
49+
except ImportError as e:
50+
logger.error("SQLAlchemy is required for database connectivity testing")
51+
raise ImportError("SQLAlchemy is required for database containers") from e
52+
53+
start_time = time.time()
54+
55+
while True:
56+
if time.time() - start_time > self._startup_timeout:
57+
raise TimeoutError(
58+
f"Database connection failed after {self._startup_timeout}s timeout. "
59+
f"Hint: Check if the container is ready and the database is accessible."
60+
)
61+
62+
try:
63+
connection_url = container.get_connection_url()
64+
engine = sqlalchemy.create_engine(connection_url)
65+
66+
try:
67+
with engine.connect():
68+
logger.info("Database connection test successful")
69+
return
70+
except Exception as e:
71+
logger.debug(f"Database connection attempt failed: {e}")
72+
raise
73+
finally:
74+
engine.dispose()
75+
76+
except self.transient_exceptions as e:
77+
logger.debug(f"Connection attempt failed: {e}, retrying in {self._poll_interval}s...")
78+
except Exception as e:
79+
logger.error(f"Connection failed with non-transient error: {e}")
80+
raise
81+
82+
time.sleep(self._poll_interval)

‎modules/generic/tests/test_sql.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ def test_container_inheritance(self):
151151
assert hasattr(container, "start")
152152

153153
def test_additional_transient_errors_list(self):
154-
from testcontainers.generic.sql import ADDITIONAL_TRANSIENT_ERRORS
154+
from testcontainers.generic.sql_utils import ADDITIONAL_TRANSIENT_ERRORS
155155

156156
assert isinstance(ADDITIONAL_TRANSIENT_ERRORS, list)
157157
# List may be empty if SQLAlchemy not available, or contain DBAPIError if it is

0 commit comments

Comments
 (0)