[MISC] Sync postgresql lib

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -35,7 +35,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 53
+LIBPATCH = 52
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -700,15 +700,22 @@ def list_valid_privileges_and_roles(self) -> Tuple[Set[str], Set[str]]:
                 "superuser",
             }, {role[0] for role in cursor.fetchall() if role[0]}
 
-    def set_up_database(self) -> None:
+    def set_up_database(self, temp_location: Optional[str] = None) -> None:
         """Set up postgres database with the right permissions."""
         connection = None
+        cursor = None
         try:
-            with self._connect_to_database() as connection, connection.cursor() as cursor:
-                cursor.execute("SELECT TRUE FROM pg_roles WHERE rolname='admin';")
-                if cursor.fetchone() is not None:
-                    return
+            connection = self._connect_to_database()
+            cursor = connection.cursor()
 
+            if temp_location is not None:
+                cursor.execute("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';")
+                if cursor.fetchone() is None:
+                    cursor.execute(f"CREATE TABLESPACE temp LOCATION '{temp_location}';")
+                    cursor.execute("GRANT CREATE ON TABLESPACE temp TO public;")
+
+            cursor.execute("SELECT TRUE FROM pg_roles WHERE rolname='admin';")
+            if cursor.fetchone() is None:
                 # Allow access to the postgres database only to the system users.
                 cursor.execute("REVOKE ALL PRIVILEGES ON DATABASE postgres FROM PUBLIC;")
                 cursor.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC;")
@@ -727,6 +734,8 @@ def set_up_database(self) -> None:
             logger.error(f"Failed to set up databases: {e}")
             raise PostgreSQLDatabasesSetupError() from e
         finally:
+            if cursor is not None:
+                cursor.close()
             if connection is not None:
                 connection.close()