Skip to content

Commit

Permalink
bump postgresql.py lib to 28
Browse files Browse the repository at this point in the history
  • Loading branch information
lucasgameiroborges committed Jun 13, 2024
1 parent 3d0ddfb commit 17427df
Showing 1 changed file with 21 additions and 11 deletions.
32 changes: 21 additions & 11 deletions lib/charms/postgresql_k8s/v0/postgresql.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@

# Increment this PATCH version before using `charmcraft publish-lib` or reset
# to 0 if you are raising the major API version
LIBPATCH = 26
LIBPATCH = 28

INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE = "invalid role(s) for extra user roles"

Expand Down Expand Up @@ -111,20 +111,19 @@ def __init__(
self.system_users = system_users

def _connect_to_database(
self, database: str = None, connect_to_current_host: bool = False
self, database: str = None, database_host: str = None
) -> psycopg2.extensions.connection:
"""Creates a connection to the database.
Args:
database: database to connect to (defaults to the database
provided when the object for this class was created).
connect_to_current_host: whether to connect to the current host
instead of the primary host.
database_host: host to connect to instead of the primary host.
Returns:
psycopg2 connection object.
"""
host = self.current_host if connect_to_current_host else self.primary_host
host = database_host if database_host is not None else self.primary_host
connection = psycopg2.connect(
f"dbname='{database if database else self.database}' user='{self.user}' host='{host}'"
f"password='{self.password}' connect_timeout=1"
Expand Down Expand Up @@ -231,7 +230,10 @@ def create_user(
user_definition += f"WITH {'NOLOGIN' if user == 'admin' else 'LOGIN'}{' SUPERUSER' if admin else ''} ENCRYPTED PASSWORD '{password}'{'IN ROLE admin CREATEDB' if admin_role else ''}"
if privileges:
user_definition += f' {" ".join(privileges)}'
cursor.execute(sql.SQL("BEGIN;"))
cursor.execute(sql.SQL("SET LOCAL log_statement = 'none';"))
cursor.execute(sql.SQL(f"{user_definition};").format(sql.Identifier(user)))
cursor.execute(sql.SQL("COMMIT;"))

# Add extra user roles to the new user.
if roles:
Expand Down Expand Up @@ -388,7 +390,7 @@ def get_postgresql_text_search_configs(self) -> Set[str]:
Set of PostgreSQL text search configs.
"""
with self._connect_to_database(
connect_to_current_host=True
database_host=self.current_host
) as connection, connection.cursor() as cursor:
cursor.execute("SELECT CONCAT('pg_catalog.', cfgname) FROM pg_ts_config;")
text_search_configs = cursor.fetchall()
Expand All @@ -401,7 +403,7 @@ def get_postgresql_timezones(self) -> Set[str]:
Set of PostgreSQL timezones.
"""
with self._connect_to_database(
connect_to_current_host=True
database_host=self.current_host
) as connection, connection.cursor() as cursor:
cursor.execute("SELECT name FROM pg_timezone_names;")
timezones = cursor.fetchall()
Expand Down Expand Up @@ -434,7 +436,7 @@ def is_tls_enabled(self, check_current_host: bool = False) -> bool:
"""
try:
with self._connect_to_database(
connect_to_current_host=check_current_host
database_host=self.current_host if check_current_host else None
) as connection, connection.cursor() as cursor:
cursor.execute("SHOW ssl;")
return "on" in cursor.fetchone()[0]
Expand Down Expand Up @@ -502,24 +504,32 @@ def set_up_database(self) -> None:
if connection is not None:
connection.close()

def update_user_password(self, username: str, password: str) -> None:
def update_user_password(
self, username: str, password: str, database_host: str = None
) -> None:
"""Update a user password.
Args:
username: the user to update the password.
password: the new password for the user.
database_host: the host to connect to.
Raises:
PostgreSQLUpdateUserPasswordError if the password couldn't be changed.
"""
connection = None
try:
with self._connect_to_database() as connection, connection.cursor() as cursor:
with self._connect_to_database(
database_host=database_host
) as connection, connection.cursor() as cursor:
cursor.execute(sql.SQL("BEGIN;"))
cursor.execute(sql.SQL("SET LOCAL log_statement = 'none';"))
cursor.execute(
sql.SQL("ALTER USER {} WITH ENCRYPTED PASSWORD '" + password + "';").format(
sql.Identifier(username)
)
)
cursor.execute(sql.SQL("COMMIT;"))
except psycopg2.Error as e:
logger.error(f"Failed to update user password: {e}")
raise PostgreSQLUpdateUserPasswordError()
Expand Down Expand Up @@ -610,7 +620,7 @@ def validate_date_style(self, date_style: str) -> bool:
"""
try:
with self._connect_to_database(
connect_to_current_host=True
database_host=self.current_host
) as connection, connection.cursor() as cursor:
cursor.execute(
sql.SQL(
Expand Down

0 comments on commit 17427df

Please sign in to comment.