Skip to content

Commit ccb1473

Browse files
committed
fix(connection, startup): support new startup message properties
1 parent dbe8cbd commit ccb1473

File tree

5 files changed

+78
-1
lines changed

5 files changed

+78
-1
lines changed

redshift_connector/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
PGVarchar,
3939
)
4040
from redshift_connector.redshift_property import RedshiftProperty
41+
from redshift_connector.utils import DriverInfo
4142

4243
from .version import __version__
4344

@@ -267,6 +268,7 @@ def connect(
267268
replication=info.replication,
268269
client_protocol_version=info.client_protocol_version,
269270
database_metadata_current_db_only=database_metadata_current_db_only,
271+
credentials_provider=info.credentials_provider,
270272
)
271273

272274

redshift_connector/core.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
FC_BINARY,
4848
NULL,
4949
NULL_BYTE,
50+
DriverInfo,
5051
array_check_dimensions,
5152
array_dim_lengths,
5253
array_find_first_element,
@@ -347,6 +348,16 @@ def _getError(self: "Connection", error):
347348
warn("DB-API extension connection.%s used" % error.__name__, stacklevel=3)
348349
return error
349350

351+
@property
352+
def client_os_version(self: "Connection") -> str:
353+
from platform import platform as CLIENT_PLATFORM
354+
355+
try:
356+
os_version: str = CLIENT_PLATFORM()
357+
except:
358+
os_version = "unknown"
359+
return os_version
360+
350361
def __init__(
351362
self: "Connection",
352363
user: str,
@@ -365,6 +376,7 @@ def __init__(
365376
replication: typing.Optional[str] = None,
366377
client_protocol_version: int = DEFAULT_PROTOCOL_VERSION,
367378
database_metadata_current_db_only: bool = True,
379+
credentials_provider: typing.Optional[str] = None,
368380
):
369381
"""
370382
Creates a :class:`Connection` to an Amazon Redshift cluster. For more information on establishing a connection to an Amazon Redshift cluster using `federated API access <https://aws.amazon.com/blogs/big-data/federated-api-access-to-amazon-redshift-using-an-amazon-redshift-connector-for-python/>`_ see our examples page.
@@ -401,6 +413,8 @@ def __init__(
401413
The requested server protocol version. The default value is 1 representing `EXTENDED_RESULT_METADATA`. If the requested server protocol cannot be satisfied, a warning will be displayed to the user.
402414
database_metadata_current_db_only : bool
403415
Is `datashare <https://docs.aws.amazon.com/redshift/latest/dg/datashare-overview.html>`_ disabled. Default value is True, implying datasharing will not be used.
416+
credentials_provider : Optional[str]
417+
The class-path of the IdP plugin used for authentication with Amazon Redshift.
404418
"""
405419
self.merge_socket_read = False
406420

@@ -432,8 +446,13 @@ def __init__(
432446
"application_name": application_name,
433447
"replication": replication,
434448
"client_protocol_version": str(self._client_protocol_version),
449+
"driver_version": DriverInfo.driver_full_name(),
450+
"os_version": self.client_os_version,
435451
}
436452

453+
if credentials_provider:
454+
init_params["plugin_name"] = credentials_provider
455+
437456
for k, v in tuple(init_params.items()):
438457
if isinstance(v, str):
439458
init_params[k] = v.encode("utf8")

redshift_connector/iam_helper.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,11 @@ def set_iam_properties(
146146
"Invalid connection property setting. It is not valid to provide both Credentials provider and "
147147
"AWS credentials or AWS profile"
148148
)
149+
elif not isinstance(credentials_provider, str):
150+
raise InterfaceError(
151+
"Invalid connection property setting. It is not valid to provide a non-string value to "
152+
"credentials_provider."
153+
)
149154
else:
150155
info.credentials_provider = credentials_provider
151156
elif profile is not None:
@@ -272,9 +277,11 @@ def set_iam_credentials(info: RedshiftProperty) -> None:
272277
except (AttributeError, ModuleNotFoundError):
273278
_logger.debug("Failed to load user defined plugin: {}".format(info.credentials_provider))
274279
try:
275-
klass = dynamic_plugin_import("redshift_connector.plugin.{}".format(info.credentials_provider))
280+
predefined_idp: str = "redshift_connector.plugin.{}".format(info.credentials_provider)
281+
klass = dynamic_plugin_import(predefined_idp)
276282
provider = klass() # type: ignore
277283
provider.add_parameter(info) # type: ignore
284+
info.credentials_provider = predefined_idp
278285
except (AttributeError, ModuleNotFoundError):
279286
_logger.debug(
280287
"Failed to load pre-defined IdP plugin from redshift_connector.plugin: {}".format(

redshift_connector/utils/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
array_has_null,
77
walk_array,
88
)
9+
from .driver_info import DriverInfo
910
from .type_utils import (
1011
FC_BINARY,
1112
NULL,
Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
class DriverInfo:
2+
"""
3+
No-op informative class containing Amazon Redshift Python driver specifications.
4+
"""
5+
6+
@staticmethod
7+
def version() -> str:
8+
"""
9+
The version of redshift_connector
10+
Returns
11+
-------
12+
The redshift_connector package version: str
13+
"""
14+
from redshift_connector import __version__ as DRIVER_VERSION
15+
16+
return str(DRIVER_VERSION)
17+
18+
@staticmethod
19+
def driver_name() -> str:
20+
"""
21+
The name of the Amazon Redshift Python driver, redshift_connector
22+
Returns
23+
-------
24+
The human readable name of the redshift_connector package: str
25+
"""
26+
return "Redshift Python Driver"
27+
28+
@staticmethod
29+
def driver_short_name() -> str:
30+
"""
31+
The shortened name of the Amazon Redshift Python driver, redshift_connector
32+
Returns
33+
-------
34+
The shortened human readable name of the Amazon Redshift Python driver: str
35+
"""
36+
return "RsPython"
37+
38+
@staticmethod
39+
def driver_full_name() -> str:
40+
"""
41+
The fully qualified name of the Amazon Redshift Python driver, redshift_connector
42+
Returns
43+
-------
44+
The fully qualified name of the Amazon Redshift Python driver: str
45+
"""
46+
return "{driver_name} {driver_version}".format(
47+
driver_name=DriverInfo.driver_name(), driver_version=DriverInfo.version()
48+
)

0 commit comments

Comments
 (0)