Fix failover logic (#170)

* Fix failover partner connection rotation

Should address issue: #167

src/pytds/__init__.py

@@ -349,6 +349,7 @@ def attempt(attempt_timeout: float) -> Connection:
                         tds_socket=tds_socket,
                     )
         host, port, instance = login.servers[0]
+        login.servers.rotate(1)
         return _connect(
             login=login,
             host=host,
@@ -368,7 +369,7 @@ def attempt(attempt_timeout: float) -> Connection:
     def ex_handler(ex: Exception) -> None:
         if isinstance(ex, LoginError):
             raise ex
-        elif isinstance(ex, BrokenPipeError):
+        elif isinstance(ex, BrokenPipeError) or isinstance(ex, socket.timeout):
             # Allow to retry when BrokenPipeError is received
             pass
         elif isinstance(ex, OperationalError):
@@ -416,17 +417,14 @@ def _connect(
     """
     Establish physical connection and login.
     """
-    try:
-        login.server_name = host
-        login.instance_name = instance
-        resolved_port = instance_browser_client.resolve_instance_port(
-            server=host, port=port, instance=instance, timeout=timeout
-        )
-        if not sock:
-            logger.info("Opening socket to %s:%d", host, resolved_port)
-            sock = socket.create_connection((host, resolved_port), timeout)
-    except Exception as e:
-        raise LoginError(f"Cannot connect to server '{host}': {e}") from e
+    login.server_name = host
+    login.instance_name = instance
+    resolved_port = instance_browser_client.resolve_instance_port(
+        server=host, port=port, instance=instance, timeout=timeout
+    )
+    if not sock:
+        logger.info("Opening socket to %s:%d", host, resolved_port)
+        sock = socket.create_connection((host, resolved_port), timeout)
     try:
         sock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
 

tests/all_test.py

@@ -210,21 +210,6 @@ def test_connection_timeout_no_mars():
             assert cur.fetchall() == [(1,)]
 
 
-@unittest.skipUnless(LIVE_TEST, "requires HOST variable to be set")
-def test_connection_no_mars_no_pooling():
-    kwargs = settings.CONNECT_KWARGS.copy()
-    kwargs.update(
-        {
-            "use_mars": False,
-            "pooling": False,
-        }
-    )
-    with connect(**kwargs) as conn:
-        with conn.cursor() as cur:
-            cur.execute("select 1")
-            assert cur.fetchall() == [(1,)]
-
-
 @unittest.skipUnless(LIVE_TEST, "requires HOST variable to be set")
 def test_row_strategies():
     kwargs = settings.CONNECT_KWARGS.copy()
@@ -324,9 +309,9 @@ def test_new_datetime(self):
 @unittest.skipUnless(LIVE_TEST, "requires HOST variable to be set")
 class BadConnection(unittest.TestCase):
     def test_invalid_parameters(self):
-        with self.assertRaises(Error):
+        with self.assertRaises(socket.gaierror):
             with connect(
-                server=settings.HOST + "bad",
+                server="badhost",
                 database="master",
                 user="baduser",
                 password=settings.PASSWORD,

tests/connection_test.py

@@ -0,0 +1,42 @@
+import pytest
+import pytds
+import unittest
+import settings
+from pytds import (
+    connect,
+    Error
+)
+
+LIVE_TEST = getattr(settings, "LIVE_TEST", True)
+
+
+
+@unittest.skipUnless(LIVE_TEST, "requires HOST variable to be set")
+def test_connection_no_mars_no_pooling():
+    kwargs = settings.CONNECT_KWARGS.copy()
+    kwargs.update(
+        {
+            "use_mars": False,
+            "pooling": False,
+        }
+    )
+    with connect(**kwargs) as conn:
+        with conn.cursor() as cur:
+            cur.execute("select 1")
+            assert cur.fetchall() == [(1,)]
+
+
+@unittest.skipUnless(LIVE_TEST, "requires HOST variable to be set")
+def test_failover_partner():
+    kwargs = settings.CONNECT_KWARGS.copy()
+    kwargs.update(
+        {
+            "server": "192.168.1.1\\sqlexpress-doesnotexist",
+            "failover_partner": settings.CONNECT_KWARGS["server"],
+            "pooling": False,
+        }
+    )
+    with connect(**kwargs) as conn:
+        with conn.cursor() as cur:
+            cur.execute("select 1")
+            assert cur.fetchall() == [(1,)]