feat(Connection): add initializer parameter for server protocol version, default to extended metadata

Author: Brooke-white

redshift_connector/__init__.py

@@ -2,6 +2,7 @@
 import os
 import typing
 
+from redshift_connector.config import DEFAULT_PROTOCOL_VERSION
 from redshift_connector.core import BINARY, Connection, Cursor
 from redshift_connector.error import (
     ArrayContentNotHomogenousError,
@@ -110,6 +111,7 @@ def connect(
     db_groups: typing.List[str] = list(),
     force_lowercase: bool = False,
     allow_db_user_override: bool = False,
+    client_protocol_version: int = DEFAULT_PROTOCOL_VERSION,
     log_level: int = 0,
     log_path: str = log_path,
 ) -> Connection:
@@ -163,6 +165,7 @@ def connect(
         db_groups=db_groups,
         force_lowercase=force_lowercase,
         allow_db_user_override=allow_db_user_override,
+        client_protocol_version=client_protocol_version,
     )
 
     return Connection(
@@ -180,6 +183,7 @@ def connect(
         tcp_keepalive=info.tcp_keepalive,
         application_name=info.application_name,
         replication=info.replication,
+        client_protocol_version=info.client_protocol_version,
     )
 
 

redshift_connector/config.py

@@ -3,11 +3,32 @@
 from calendar import timegm
 from datetime import datetime as Datetime
 from datetime import timezone as Timezone
+from enum import IntEnum
 
 FC_TEXT: int = 0
 FC_BINARY: int = 1
 _client_encoding: str = "utf8"
 
+
+class ClientProtocolVersion(IntEnum):
+    BASE_SERVER = 0
+    EXTENDED_RESULT_METADATA = 1
+    BINARY = 2
+
+    @classmethod
+    def list(cls) -> typing.List[int]:
+        return list(map(lambda p: p.value, cls))  # type: ignore
+
+    @classmethod
+    def get_name(cls, i: int) -> str:
+        try:
+            return ClientProtocolVersion(i).name
+        except ValueError:
+            return str(i)
+
+
+DEFAULT_PROTOCOL_VERSION: int = ClientProtocolVersion.EXTENDED_RESULT_METADATA.value
+
 min_int2: int = -(2 ** 15)
 max_int2: int = 2 ** 15
 min_int4: int = -(2 ** 31)

redshift_connector/core.py

@@ -17,6 +17,8 @@
 from scramp import ScramClient  # type: ignore
 
 from redshift_connector.config import (
+    DEFAULT_PROTOCOL_VERSION,
+    ClientProtocolVersion,
     _client_encoding,
     max_int2,
     max_int4,
@@ -99,7 +101,6 @@
 
 __author__ = "Mathieu Fenniak"
 
-
 ZERO: Timedelta = Timedelta(0)
 BINARY: type = bytes
 
@@ -321,7 +322,6 @@ def create_message(code: bytes, data: bytes = b"") -> bytes:
 IDLE_IN_TRANSACTION: bytes = b"T"
 IDLE_IN_FAILED_TRANSACTION: bytes = b"E"
 
-
 arr_trans: typing.Mapping[int, typing.Optional[str]] = dict(zip(map(ord, "[] 'u"), ["{", "}", None, None, None]))
 
 
@@ -363,6 +363,7 @@ def __init__(
         tcp_keepalive: typing.Optional[bool] = True,
         application_name: typing.Optional[str] = None,
         replication: typing.Optional[str] = None,
+        client_protocol_version: int = DEFAULT_PROTOCOL_VERSION,
     ):
 
         self.merge_socket_read = False
@@ -382,6 +383,7 @@ def __init__(
         self.parameter_statuses: deque = deque(maxlen=100)
         self.max_prepared_statements: int = int(max_prepared_statements)
         self._run_cursor: Cursor = Cursor(self, paramstyle="named")
+        self._client_protocol_version: int = client_protocol_version
 
         if user is None:
             raise InterfaceError("The 'user' connection parameter cannot be None")
@@ -391,6 +393,7 @@ def __init__(
             "database": database,
             "application_name": application_name,
             "replication": replication,
+            "client_protocol_version": str(self._client_protocol_version),
         }
 
         for k, v in tuple(init_params.items()):
@@ -551,6 +554,15 @@ def __init__(
         if self.error is not None:
             raise self.error
 
+        # if we didn't receive a server_protocol_version from the server, default to
+        # using BASE_SERVER as the server is likely lacking this functionality due to
+        # being out of date
+        if (
+            self._client_protocol_version > ClientProtocolVersion.BASE_SERVER
+            and not (b"server_protocol_version", str(self._client_protocol_version).encode()) in self.parameter_statuses
+        ):
+            self._client_protocol_version = ClientProtocolVersion.BASE_SERVER
+
         self.in_transaction = False
 
     def handle_ERROR_RESPONSE(self: "Connection", data, ps):
@@ -845,24 +857,37 @@ def make_params(self: "Connection", values):
     # get the metadata of each row in database
     # and store these metadata into ps dictionary
     def handle_ROW_DESCRIPTION(self: "Connection", data, cursor: Cursor) -> None:
+        if cursor.ps is None:
+            raise InterfaceError("Cursor is missing prepared statement")
+        elif "row_desc" not in cursor.ps:
+            raise InterfaceError("Prepared Statement is missing row description")
+
         count: int = h_unpack(data)[0]
         idx = 2
         for i in range(count):
-            name = data[idx : data.find(NULL_BYTE, idx)]
-            idx += len(name) + 1
+            column_label = data[idx : data.find(NULL_BYTE, idx)]
+            idx += len(column_label) + 1
+
             field: typing.Dict = dict(
                 zip(
                     ("table_oid", "column_attrnum", "type_oid", "type_size", "type_modifier", "format"),
                     ihihih_unpack(data, idx),
                 )
             )
-            field["name"] = name
+            field["label"] = column_label
             idx += 18
 
-            if cursor.ps is None:
-                raise InterfaceError("Cursor is missing prepared statement")
-            elif "row_desc" not in cursor.ps:
-                raise InterfaceError("Prepared Statement is missing row description")
+            if self._client_protocol_version >= ClientProtocolVersion.EXTENDED_RESULT_METADATA:
+                for entry in ("schema_name", "table_name", "column_name", "catalog_name"):
+                    field[entry] = data[idx : data.find(NULL_BYTE, idx)]
+                    idx += len(field[entry]) + 1
+
+                temp: int = h_unpack(data, idx)[0]
+                field["nullable"] = temp & 0x1
+                field["autoincrement"] = (temp >> 4) & 0x1
+                field["read_only"] = (temp >> 8) & 0x1
+                field["searchable"] = (temp >> 12) & 0x1
+                idx += 2
 
             cursor.ps["row_desc"].append(field)
             field["pg8000_fc"], field["func"] = pg_types[field["type_oid"]]
@@ -1178,6 +1203,18 @@ def handle_PARAMETER_STATUS(self: "Connection", data: bytes, ps) -> None:
         if key == b"client_encoding":
             encoding = value.decode("ascii").lower()
             _client_encoding = pg_to_py_encodings.get(encoding, encoding)
+        elif key == b"server_protocol_version":
+            # when a mismatch occurs between the client's requested protocol version, and the server's response,
+            # warn the user and follow server
+            if self._client_protocol_version != int(value):
+                warn(
+                    "Server indicated {} transfer protocol will be used rather than protocol requested by client: {}".format(
+                        ClientProtocolVersion.get_name(int(value)),
+                        ClientProtocolVersion.get_name(self._client_protocol_version),
+                    ),
+                    stacklevel=3,
+                )
+                self._client_protocol_version = int(value)
         elif key == b"server_version":
             self._server_version: LooseVersion = LooseVersion(value.decode("ascii"))
             if self._server_version < LooseVersion("8.2.0"):

redshift_connector/cursor.py

@@ -112,7 +112,7 @@ def _getDescription(self: "Cursor") -> typing.Optional[typing.List[typing.Option
             return None
         columns: typing.List[typing.Optional[typing.Tuple]] = []
         for col in row_desc:
-            columns.append((col["name"], col["type_oid"], None, None, None, None, None))
+            columns.append((col["label"], col["type_oid"], None, None, None, None, None))
         return columns
 
     ##

redshift_connector/iam_helper.py

@@ -5,6 +5,7 @@
 import boto3  # type: ignore
 import botocore  # type: ignore
 
+from redshift_connector.config import ClientProtocolVersion
 from redshift_connector.credentials_holder import CredentialsHolder
 from redshift_connector.error import InterfaceError
 from redshift_connector.plugin import (
@@ -65,6 +66,7 @@ def set_iam_properties(
     db_groups: typing.List[str],
     force_lowercase: bool,
     allow_db_user_override: bool,
+    client_protocol_version: int,
 ) -> None:
     if info is None:
         raise InterfaceError("Invalid connection property setting. info must be specified")
@@ -108,6 +110,13 @@ def set_iam_properties(
     if password is None:
         raise InterfaceError("Invalid connection property setting. password must be specified")
 
+    if client_protocol_version not in ClientProtocolVersion.list():
+        raise InterfaceError(
+            "Invalid connection property setting. client_protocol_version must be in: {}".format(
+                ClientProtocolVersion.list()
+            )
+        )
+
     # basic driver parameters
     info.user_name = user
     info.host = host
@@ -121,6 +130,7 @@ def set_iam_properties(
     info.tcp_keepalive = tcp_keepalive
     info.application_name = application_name
     info.replication = replication
+    info.client_protocol_version = client_protocol_version
 
     # Idp parameters
     info.idp_host = idp_host

redshift_connector/redshift_property.py

@@ -1,6 +1,8 @@
 import typing
 from typing import TYPE_CHECKING
 
+from redshift_connector.config import DEFAULT_PROTOCOL_VERSION
+
 if TYPE_CHECKING:
     from redshift_connector.iam_helper import SSLMode
 
@@ -71,6 +73,8 @@ class RedshiftProperty:
     max_prepared_statements: int = 1000
     # Use this property to enable or disable TCP keepalives. The following values are possible:
     tcp_keepalive: bool = True
+    # client's requested transfer protocol version. See config.py for supported protocols
+    client_protocol_version: int = DEFAULT_PROTOCOL_VERSION
     # application name
     application_name: typing.Optional[str] = None
     # Used to run in streaming replication mode. If your server character encoding is not ascii or utf8,

test/integration/test_connection.py

@@ -7,6 +7,7 @@
 import pytest  # type: ignore
 
 import redshift_connector
+from redshift_connector.config import ClientProtocolVersion
 
 conf = configparser.ConfigParser()
 root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
@@ -181,3 +182,18 @@ def test_scram_sha_256(db_kwargs):
     # Should only raise an exception saying db doesn't exist
     with pytest.raises(redshift_connector.ProgrammingError, match="3D000"):
         redshift_connector.connect(**db_kwargs)
+
+
+@pytest.mark.parametrize("_input", ClientProtocolVersion.list()[:-1])
+def test_client_protocol_version_is_used(db_kwargs, _input):
+    db_kwargs["client_protocol_version"] = _input
+
+    with redshift_connector.connect(**db_kwargs) as conn:
+        assert conn._client_protocol_version == _input
+
+
+def test_client_protocol_version_invalid_warns_user(db_kwargs):
+    db_kwargs["client_protocol_version"] = max(ClientProtocolVersion.list()) + 1
+
+    with pytest.warns(UserWarning):
+        redshift_connector.Connection(**db_kwargs)