Skip to content

Commit

Permalink
inline ssl certificate
Browse files Browse the repository at this point in the history
  • Loading branch information
Rohit Chatterjee committed Dec 17, 2024
1 parent 46f0429 commit 5422818
Show file tree
Hide file tree
Showing 2 changed files with 149 additions and 7 deletions.
43 changes: 39 additions & 4 deletions ddpui/datainsights/warehouse/postgres.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import tempfile
from urllib.parse import quote

from sqlalchemy.engine import create_engine
Expand All @@ -18,11 +19,45 @@ def __init__(self, creds: dict):
"""
creds["encoded_username"] = quote(creds["username"].strip())
creds["encoded_password"] = quote(creds["password"].strip())
connection_string = (
"postgresql://{encoded_username}:{encoded_password}@{host}/{database}".format(**creds)
)

self.engine = create_engine(connection_string, pool_size=5, pool_timeout=30)
connection_args = {
"host": creds["host"],
"port": creds["port"],
"dbname": creds["database"],
"user": creds["encoded_username"],
"password": creds["encoded_password"],
}

connection_string = "postgresql+psycopg2://"

if "ssl_mode" in creds:
creds["sslmode"] = creds["ssl_mode"]

Check warning on line 34 in ddpui/datainsights/warehouse/postgres.py

View check run for this annotation

Codecov / codecov/patch

ddpui/datainsights/warehouse/postgres.py#L34 

Added line #L34 was not covered by tests

if "sslrootcert" in creds:
connection_args["sslrootcert"] = creds["sslrootcert"]

if "sslmode" in creds and isinstance(creds["sslmode"], str):
connection_args["sslmode"] = creds["sslmode"]

if "sslmode" in creds and isinstance(creds["sslmode"], bool):
connection_args["sslmode"] = "require" if creds["sslmode"] else "disable"

if (
"sslmode" in creds
and isinstance(creds["sslmode"], dict)
and "ca_certificate" in creds["sslmode"]
):
# connect_params['sslcert'] needs a file path but
# creds['sslmode']['ca_certificate']
# is a string (i.e. the actual certificate). so we write
# it to disk and pass the file path
with tempfile.NamedTemporaryFile(delete=False) as fp:
fp.write(creds["sslmode"]["ca_certificate"].encode())
connection_args["sslrootcert"] = fp.name

Check warning on line 56 in ddpui/datainsights/warehouse/postgres.py

View check run for this annotation

Codecov / codecov/patch

ddpui/datainsights/warehouse/postgres.py#L54-L56 

Added lines #L54 - L56 were not covered by tests

self.engine = create_engine(
connection_string, connect_args=connection_args, pool_size=5, pool_timeout=30
)
self.inspect_obj: Inspector = inspect(
self.engine
) # this will be used to fetch metadata of the database
Expand Down
113 changes: 110 additions & 3 deletions ddpui/tests/core/datainsights/factories/test_warehouse_factory.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,14 +45,121 @@ def test_warehouse_factory():
WarehouseFactory.connect({}, "some-no-supported-warehouse-type")


def test_connect_args_1():
"""tests creation on connect_args parameter to create_engine"""
with patch("ddpui.datainsights.warehouse.postgres.inspect"):
with patch("ddpui.datainsights.warehouse.postgres.create_engine") as mock_create_engine:
PostgresClient(
{
"username": "user name",
"password": "pass word",
"host": "host",
"port": 1234,
"database": "db",
"sslrootcert": "sslrootcert",
"sslmode": "require",
}
)
mock_create_engine.assert_called_with(
"postgresql+psycopg2://",
connect_args={
"host": "host",
"port": 1234,
"dbname": "db",
"user": "user%20name",
"password": "pass%20word",
"sslrootcert": "sslrootcert",
"sslmode": "require",
},
pool_size=5,
pool_timeout=30,
)


def test_connect_args_2():
"""tests creation on connect_args parameter to create_engine"""
with patch("ddpui.datainsights.warehouse.postgres.inspect"):
with patch("ddpui.datainsights.warehouse.postgres.create_engine") as mock_create_engine:
PostgresClient(
{
"username": "user name",
"password": "pass word",
"host": "host",
"port": 1234,
"database": "db",
"sslrootcert": "sslrootcert",
"sslmode": True,
}
)
mock_create_engine.assert_called_with(
"postgresql+psycopg2://",
connect_args={
"host": "host",
"port": 1234,
"dbname": "db",
"user": "user%20name",
"password": "pass%20word",
"sslrootcert": "sslrootcert",
"sslmode": "require",
},
pool_size=5,
pool_timeout=30,
)


def test_connect_args_3():
"""tests creation on connect_args parameter to create_engine"""
with patch("ddpui.datainsights.warehouse.postgres.inspect"):
with patch("ddpui.datainsights.warehouse.postgres.create_engine") as mock_create_engine:
PostgresClient(
{
"username": "user name",
"password": "pass word",
"host": "host",
"port": 1234,
"database": "db",
"sslrootcert": "sslrootcert",
"sslmode": False,
}
)
mock_create_engine.assert_called_with(
"postgresql+psycopg2://",
connect_args={
"host": "host",
"port": 1234,
"dbname": "db",
"user": "user%20name",
"password": "pass%20word",
"sslrootcert": "sslrootcert",
"sslmode": "disable",
},
pool_size=5,
pool_timeout=30,
)


def test_url_encoding():
"""tests url encoding of username and password"""

with patch("ddpui.datainsights.warehouse.postgres.inspect"):
with patch("ddpui.datainsights.warehouse.postgres.create_engine") as mock_create_engine:
PostgresClient(
{"username": "user name", "password": "pass word", "host": "host", "database": "db"}
{
"username": "user name",
"password": "pass word",
"host": "host",
"port": 1234,
"database": "db",
}
)
mock_create_engine.assert_called_with(
"postgresql://user%20name:pass%20word@host/db", pool_size=5, pool_timeout=30
"postgresql+psycopg2://",
connect_args={
"host": "host",
"port": 1234,
"dbname": "db",
"user": "user%20name",
"password": "pass%20word",
},
pool_size=5,
pool_timeout=30,
)

0 comments on commit 5422818

Please sign in to comment.