Skip to content
Merged
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
10 changes: 8 additions & 2 deletions testcontainers/postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,18 @@ class PostgresContainer(DbContainer):
POSTGRES_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "test")
POSTGRES_DB = os.environ.get("POSTGRES_DB", "test")

def __init__(self, image="postgres:latest", port=5432, user=None, password=None, dbname=None):
def __init__(self,
image="postgres:latest",
port=5432, user=None,
password=None,
dbname=None,
driver="psycopg2"):
super(PostgresContainer, self).__init__(image=image)
self.POSTGRES_USER = user or self.POSTGRES_USER
self.POSTGRES_PASSWORD = password or self.POSTGRES_PASSWORD
self.POSTGRES_DB = dbname or self.POSTGRES_DB
self.port_to_expose = port
self.driver = driver

self.with_exposed_ports(self.port_to_expose)

Expand All @@ -47,7 +53,7 @@ def _configure(self):
self.with_env("POSTGRES_DB", self.POSTGRES_DB)

def get_connection_url(self, host=None):
return super()._create_connection_url(dialect="postgresql+psycopg2",
return super()._create_connection_url(dialect="postgresql+{}".format(self.driver),
username=self.POSTGRES_USER,
password=self.POSTGRES_PASSWORD,
db_name=self.POSTGRES_DB,
Expand Down