Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions awswrangler/_databases.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Databases Utilities."""

import logging
import ssl
from typing import Any, Dict, Generator, Iterator, List, NamedTuple, Optional, Tuple, Union, cast

import boto3
Expand All @@ -22,6 +23,7 @@ class ConnectionAttributes(NamedTuple):
host: str
port: int
database: str
ssl_context: Optional[ssl.SSLContext]


def _get_dbname(cluster_id: str, boto3_session: Optional[boto3.Session] = None) -> str:
Expand All @@ -41,13 +43,28 @@ def _get_connection_attributes_from_catalog(
else:
database_sep = "/"
port, database = details["JDBC_CONNECTION_URL"].split(":")[3].split(database_sep)
ssl_context: Optional[ssl.SSLContext] = None
if details.get("JDBC_ENFORCE_SSL") == "true":
ssl_cert_path: Optional[str] = details.get("CUSTOM_JDBC_CERT")
ssl_cadata: Optional[str] = None
if ssl_cert_path:
bucket_name, key_path = _utils.parse_path(ssl_cert_path)
client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session)
try:
ssl_cadata = client_s3.get_object(Bucket=bucket_name, Key=key_path)["Body"].read().decode("utf-8")
except client_s3.exception.NoSuchKey:
raise exceptions.NoFilesFound( # pylint: disable=raise-missing-from
f"No CA certificate found at {ssl_cert_path}."
)
ssl_context = ssl.create_default_context(cadata=ssl_cadata)
return ConnectionAttributes(
kind=details["JDBC_CONNECTION_URL"].split(":")[1].lower(),
user=details["USERNAME"],
password=details["PASSWORD"],
host=details["JDBC_CONNECTION_URL"].split(":")[2].replace("/", ""),
port=int(port),
database=dbname if dbname is not None else database,
ssl_context=ssl_context,
)


Expand All @@ -71,6 +88,7 @@ def _get_connection_attributes_from_secrets_manager(
host=secret_value["host"],
port=secret_value["port"],
database=_dbname,
ssl_context=None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess there is really no way to specify/pass the ca data via secrets manager... Do you think we should clarify the docs in that regard?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Doesn't hurt to clarify, of course.

)


Expand Down
8 changes: 7 additions & 1 deletion awswrangler/mysql.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,10 +78,15 @@ def connect(
write_timeout: Optional[int] = None,
connect_timeout: int = 10,
) -> pymysql.connections.Connection:
"""Return a pymysql connection from a Glue Catalog Connection.
"""Return a pymysql connection from a Glue Catalog Connection or Secrets Manager.

https://pymysql.readthedocs.io

Note
----
It is only possible to configure SSL using Glue Catalog Connection. More at:
https://docs.aws.amazon.com/glue/latest/dg/connection-defining.html

Parameters
----------
connection : str
Expand Down Expand Up @@ -136,6 +141,7 @@ def connect(
password=attrs.password,
port=attrs.port,
host=attrs.host,
ssl=attrs.ssl_context,
read_timeout=read_timeout,
write_timeout=write_timeout,
connect_timeout=connect_timeout,
Expand Down
26 changes: 26 additions & 0 deletions cloudformation/databases.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -384,6 +384,32 @@ Resources:
PASSWORD:
Ref: DatabasesPassword
Name: aws-data-wrangler-mysql
MysqlGlueConnectionSSL:
Type: AWS::Glue::Connection
Properties:
CatalogId:
Ref: AWS::AccountId
ConnectionInput:
Description: Connect to Aurora (MySQL) SSL enabled.
ConnectionType: JDBC
PhysicalConnectionRequirements:
AvailabilityZone:
Fn::Select:
- 0
- Fn::GetAZs: ''
SecurityGroupIdList:
- Ref: DatabaseSecurityGroup
SubnetId:
Fn::ImportValue: aws-data-wrangler-base-PrivateSubnet
ConnectionProperties:
JDBC_CONNECTION_URL:
Fn::Sub: jdbc:mysql://${AuroraInstanceMysql.Endpoint.Address}:${AuroraInstanceMysql.Endpoint.Port}/test
JDBC_ENFORCE_SSL: true
CUSTOM_JDBC_CERT: s3://rds-downloads/rds-combined-ca-bundle.pem
USERNAME: test
PASSWORD:
Ref: DatabasesPassword
Name: aws-data-wrangler-mysql-ssl
SqlServerGlueConnection:
Type: AWS::Glue::Connection
Properties:
Expand Down
17 changes: 15 additions & 2 deletions tests/test_mysql.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,16 @@ def mysql_con():
con.close()


def test_connection():
wr.mysql.connect("aws-data-wrangler-mysql", connect_timeout=10).close()
@pytest.fixture(scope="function")
def mysql_con_ssl():
con = wr.mysql.connect("aws-data-wrangler-mysql-ssl")
yield con
con.close()


@pytest.mark.parametrize("connection", ["aws-data-wrangler-mysql", "aws-data-wrangler-mysql-ssl"])
def test_connection(connection):
wr.mysql.connect(connection, connect_timeout=10).close()


def test_read_sql_query_simple(databases_parameters):
Expand All @@ -42,6 +50,11 @@ def test_to_sql_simple(mysql_table, mysql_con):
wr.mysql.to_sql(df, mysql_con, mysql_table, "test", "overwrite", True)


def test_to_sql_simple_ssl(mysql_table, mysql_con_ssl):
df = pd.DataFrame({"c0": [1, 2, 3], "c1": ["foo", "boo", "bar"]})
wr.mysql.to_sql(df, mysql_con_ssl, mysql_table, "test", "overwrite", True)


def test_sql_types(mysql_table, mysql_con):
table = mysql_table
df = get_df()
Expand Down