Check connection and DB schema existence of role service on start

LayerManager · Dec 20, 2023 · e1e60e6 · e1e60e6
1 parent c4cf313
commit e1e60e6
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 5 deletions.
diff --git a/src/geoserver/__init__.py b/src/geoserver/__init__.py
@@ -14,6 +14,8 @@
 GS_REST_WMS_SETTINGS = None
 GS_REST_WFS_SETTINGS = None
 GS_REST_TIMEOUT = None
+GS_JDBC_ROLE_SERVICE_URI = None
+GS_JDBC_ROLE_SERVICE_SCHEMA = None
 
 
 def ensure_data_dir(data_dir, data_dir_initial, normalized_raster_data_dir):
@@ -25,13 +27,14 @@ def ensure_data_dir(data_dir, data_dir_initial, normalized_raster_data_dir):
     os.chmod(norm_data_path, mode=0o777)
 
 
-def set_settings(gs_url, role_service, user_group_service, timeout):
+def set_settings(gs_url, role_service, jdbc_role_service_uri, jdbc_role_service_schema, user_group_service, timeout):
     role_service = role_service or 'default'
     user_group_service = user_group_service or 'default'
     # pylint: disable=global-statement
     global GS_REST, GS_REST_STYLES, GS_REST_WORKSPACES, GS_REST_SETTINGS, \
         GS_REST_SECURITY_ACL_LAYERS, GS_REST_ROLES, GS_REST_USERS, GS_REST_USER, \
-        GS_REST_WMS_SETTINGS, GS_REST_WFS_SETTINGS, GS_REST_TIMEOUT
+        GS_REST_WMS_SETTINGS, GS_REST_WFS_SETTINGS, GS_REST_TIMEOUT, GS_JDBC_ROLE_SERVICE_URI, \
+        GS_JDBC_ROLE_SERVICE_SCHEMA
 
     GS_REST = urljoin(gs_url, 'rest/')
     GS_REST_STYLES = urljoin(GS_REST, 'styles/')
@@ -43,5 +46,7 @@ def set_settings(gs_url, role_service, user_group_service, timeout):
     GS_REST_USER = urljoin(GS_REST, f'security/usergroup/service/{user_group_service}/user/')
     GS_REST_WMS_SETTINGS = urljoin(GS_REST, f'services/wms/settings/')
     GS_REST_WFS_SETTINGS = urljoin(GS_REST, f'services/wfs/settings/')
+    GS_JDBC_ROLE_SERVICE_URI = jdbc_role_service_uri
+    GS_JDBC_ROLE_SERVICE_SCHEMA = jdbc_role_service_schema
 
     GS_REST_TIMEOUT = timeout
diff --git a/src/geoserver/role_service.py b/src/geoserver/role_service.py
@@ -4,9 +4,11 @@
 import sys
 from urllib.parse import urlparse
 from xml.sax.saxutils import escape
+import psycopg2
 
+from db import util as db_util
 from requests_util import url_util
-from . import authn
+from . import authn, GS_JDBC_ROLE_SERVICE_URI, GS_JDBC_ROLE_SERVICE_SCHEMA
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
@@ -59,3 +61,20 @@ def set_primary_role_service(data_dir, role_service_name):
     element.text = role_service_name
     security_path = os.path.join(data_dir, 'security/config.xml')
     security_xml.write(security_path)
+
+
+def check_jdbc_role_service(role_service_db_uri=None, role_service_schema=None):
+    role_service_db_uri = role_service_db_uri or GS_JDBC_ROLE_SERVICE_URI
+    role_service_schema = role_service_schema or GS_JDBC_ROLE_SERVICE_SCHEMA
+
+    try:
+        db_util.get_connection_pool(db_uri_str=role_service_db_uri, encapsulate_exception=False)
+    except psycopg2.OperationalError as exc:
+        secret_conn_dict = url_util.redact_uri(role_service_db_uri)
+        raise Exception(f"Failed to connect to role service database {secret_conn_dict}") from exc
+
+    try:
+        db_util.run_query(f"select name, parent from {role_service_schema}.roles limit 0")
+    except psycopg2.OperationalError as exc:
+        secret_conn_dict = url_util.redact_uri(role_service_db_uri)
+        raise Exception(f"Error querying role service database {secret_conn_dict}") from exc
diff --git a/src/layman_settings.py b/src/layman_settings.py
@@ -125,7 +125,8 @@ class EnumWfsWmsStatus(Enum):
 LAYMAN_GS_PATH = '/geoserver/'
 
 LAYMAN_GS_URL = f"http://{LAYMAN_GS_HOST}:{LAYMAN_GS_PORT}{LAYMAN_GS_PATH}"
-geoserver.set_settings(LAYMAN_GS_URL, LAYMAN_GS_ROLE_SERVICE, LAYMAN_GS_USER_GROUP_SERVICE, DEFAULT_CONNECTION_TIMEOUT, )
+geoserver.set_settings(LAYMAN_GS_URL, LAYMAN_GS_ROLE_SERVICE, LAYMAN_ROLE_SERVICE_URI, LAYMAN_ROLE_SERVICE_SCHEMA,
+                       LAYMAN_GS_USER_GROUP_SERVICE, DEFAULT_CONNECTION_TIMEOUT, )
 geoserver.GS_AUTH = LAYMAN_GS_AUTH
 
 LAYMAN_GS_ROLE = os.environ['LAYMAN_GS_ROLE']

diff --git a/src/wait_for_deps.py b/src/wait_for_deps.py
@@ -5,7 +5,7 @@
 
 import geoserver
 import layman_settings as settings
-
+from geoserver.role_service import check_jdbc_role_service
 
 ATTEMPT_INTERVAL = 2
 MAX_ATTEMPTS = 60
@@ -48,6 +48,11 @@ def main():
             attempt += 1
     print()
 
+    # Check PostgreSQL role service and stops immediately in case of any error
+    print(f"Checking PostgreSQL role service")
+    check_jdbc_role_service()
+    print()
+
     # QGIS Server
     wait_for_msg = f"QGIS Server, url={settings.LAYMAN_QGIS_URL}"
     print(f"Waiting for {wait_for_msg}")