Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add SSL/TLS support for Postgres connections #104

Merged
merged 9 commits into from
Jul 26, 2021
Merged

Add SSL/TLS support for Postgres connections #104

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
868ed6c
inital commit
wseaton Jul 24, 2021
ac0de20
add docs
wseaton Jul 24, 2021
679c00c
fix TLS for partition range query; clean up arg mangling
wseaton Jul 25, 2021
10cc636
clean up some unneeded unwraps; rustfmt
wseaton Jul 25, 2021
66fe50d
1. add tests for tls in python 2. support both NoTls and Tls 3. opens…
wangxiaoying Jul 25, 2021
ba6c2d8
fix rust test
wangxiaoying Jul 25, 2021
b67b807
skip tls in ci for now
wangxiaoying Jul 25, 2021
9363007
1. clean up transports 2. add more tests, passed on aws rds
wangxiaoying Jul 26, 2021
9e73e86
mild refactor
dovahcrow Jul 26, 2021
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
27 changes: 27 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Justfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ build-debug:
cargo build

test:
cargo test --features src_postgres --features src_sqlite --features src_mysql --features src_csv --features dst_arrow --features dst_memory --features dst_polars -- --nocapture
cargo test --features all -- --nocapture

bootstrap-python:
cp README.md connectorx-python/README.md
Expand Down
1 change: 1 addition & 0 deletions connectorx-python/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ sqlparser = "0.9"
thiserror = "1"
url = "2"
uuid = "0.8"
postgres-native-tls = {version = "0.5"}

[build-dependencies]
built = {version = "0.5", features = ["chrono"]}
Expand Down
115 changes: 112 additions & 3 deletions connectorx-python/connectorx/tests/test_postgres.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,18 @@ def postgres_url() -> str:
return conn


@pytest.fixture(scope="module") # type: ignore
def postgres_url_tls() -> str:
conn = os.environ["POSTGRES_URL_TLS"]
return conn


@pytest.fixture(scope="module") # type: ignore
def postgres_rootcert() -> str:
cert = os.environ["POSTGRES_ROOTCERT"]
return cert


@pytest.mark.xfail
def test_on_non_select(postgres_url: str) -> None:
query = "CREATE TABLE non_select(id INTEGER NOT NULL)"
Expand All @@ -36,7 +48,8 @@ def test_partition_on_aggregation(postgres_url: str) -> None:
query = (
"SELECT test_bool, SUM(test_int) AS test_int FROM test_table GROUP BY test_bool"
)
df = read_sql(postgres_url, query, partition_on="test_int", partition_num=2)
df = read_sql(postgres_url, query,
partition_on="test_int", partition_num=2)
expected = pd.DataFrame(
index=range(3),
data={
Expand Down Expand Up @@ -74,7 +87,8 @@ def test_partition_on_aggregation2(postgres_url: str) -> None:

def test_udf(postgres_url: str) -> None:
query = "select increment(test_int) as test_int from test_table ORDER BY test_int"
df = read_sql(postgres_url, query, partition_on="test_int", partition_num=2)
df = read_sql(postgres_url, query,
partition_on="test_int", partition_num=2)
expected = pd.DataFrame(
index=range(6),
data={
Expand Down Expand Up @@ -484,7 +498,8 @@ def test_empty_result(postgres_url: str) -> None:

def test_empty_result_on_some_partition(postgres_url: str) -> None:
query = "SELECT * FROM test_table where test_int < 1"
df = read_sql(postgres_url, query, partition_on="test_int", partition_num=3)
df = read_sql(postgres_url, query,
partition_on="test_int", partition_num=3)
expected = pd.DataFrame(
data={
"test_int": pd.Series([0], dtype="Int64"),
Expand All @@ -495,3 +510,97 @@ def test_empty_result_on_some_partition(postgres_url: str) -> None:
}
)
assert_frame_equal(df, expected, check_names=True)


@pytest.mark.skipif(os.environ.get("TEST_COVER", "main") not in ("all", "pgtls"), reason="Do not test TLS unless TEST_COVER=all/pgtls")
def test_read_sql_tls(postgres_url_tls: str) -> None:
query = "SELECT * FROM test_table"
df = read_sql(
f"{postgres_url_tls}?sslmode=require",
query,
partition_on="test_int",
partition_range=(0, 2000),
partition_num=3,
)
expected = pd.DataFrame(
index=range(6),
data={
"test_int": pd.Series([1, 2, 0, 3, 4, 1314], dtype="Int64"),
"test_nullint": pd.Series([3, None, 5, 7, 9, 2], dtype="Int64"),
"test_str": pd.Series(
["str1", "str2", "a", "b", "c", None], dtype="object"
),
"test_float": pd.Series([None, 2.2, 3.1, 3, 7.8, -10], dtype="float64"),
"test_bool": pd.Series(
[True, False, None, False, None, True], dtype="boolean"
),
},
)
assert_frame_equal(df, expected, check_names=True)


@pytest.mark.skipif(os.environ.get("TEST_COVER", "main") not in ("all", "pgtls"), reason="Do not test TLS unless TEST_COVER=all/pgtls")
def test_read_sql_tls_with_cert(postgres_url_tls: str, postgres_rootcert: str) -> None:
query = "SELECT * FROM test_table"
df = read_sql(
f"{postgres_url_tls}?sslmode=require&sslrootcert={postgres_rootcert}",
query,
partition_on="test_int",
partition_range=(0, 2000),
partition_num=3,
)
expected = pd.DataFrame(
index=range(6),
data={
"test_int": pd.Series([1, 2, 0, 3, 4, 1314], dtype="Int64"),
"test_nullint": pd.Series([3, None, 5, 7, 9, 2], dtype="Int64"),
"test_str": pd.Series(
["str1", "str2", "a", "b", "c", None], dtype="object"
),
"test_float": pd.Series([None, 2.2, 3.1, 3, 7.8, -10], dtype="float64"),
"test_bool": pd.Series(
[True, False, None, False, None, True], dtype="boolean"
),
},
)
assert_frame_equal(df, expected, check_names=True)


@pytest.mark.skipif(os.environ.get("TEST_COVER", "main") not in ("all", "pgtls"), reason="Do not test TLS unless TEST_COVER=all/pgtls")
def test_read_sql_tls_disable(postgres_url_tls: str) -> None:
query = "SELECT * FROM test_table"
df = read_sql(
f"{postgres_url_tls}?sslmode=disable",
query,
partition_on="test_int",
partition_range=(0, 2000),
partition_num=3,
)
expected = pd.DataFrame(
index=range(6),
data={
"test_int": pd.Series([1, 2, 0, 3, 4, 1314], dtype="Int64"),
"test_nullint": pd.Series([3, None, 5, 7, 9, 2], dtype="Int64"),
"test_str": pd.Series(
["str1", "str2", "a", "b", "c", None], dtype="object"
),
"test_float": pd.Series([None, 2.2, 3.1, 3, 7.8, -10], dtype="float64"),
"test_bool": pd.Series(
[True, False, None, False, None, True], dtype="boolean"
),
},
)
assert_frame_equal(df, expected, check_names=True)


@pytest.mark.skipif(os.environ.get("TEST_COVER", "main") not in ("all", "pgtls"), reason="Do not test TLS unless TEST_COVER=all/pgtls")
@pytest.mark.xfail
def test_read_sql_tls_fail(postgres_url_tls: str) -> None:
query = "SELECT * FROM test_table"
df = read_sql(
f"{postgres_url_tls}?sslmode=require&sslrootcert=fake.cert",
query,
partition_on="test_int",
partition_range=(0, 2000),
partition_num=3,
)
Loading