Modernize common/check_messages_to_submit

common/check_messages_to_submit

@@ -1,62 +1,57 @@
-#!/usr/bin/env python
-# Copyright European Organization for Nuclear Research (CERN) 2013
+#!/usr/bin/env python3
+# Copyright European Organization for Nuclear Research (CERN) since 2012
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
-# You may not use this file except in compliance with the License.
-# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
 #
-# Authors:
-# - Mario Lassnig, <[email protected]>, 2013-2014
-# - Thomas Beermann, <[email protected]>, 2019
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 """
 Probe to check the queues of messages to submit by Hermes to the broker
 """
-from __future__ import print_function
 
 import sys
+import traceback
+
+from sqlalchemy import func, select
 
-from prometheus_client import CollectorRegistry, Gauge, push_to_gateway
-from rucio.common.config import config_get
-from rucio.db.sqla.session import BASE, get_session
+from rucio.db.sqla import models
+from rucio.db.sqla.session import get_session
 
-from utils.common import probe_metrics
+from utils.common import PrometheusPusher
 
 # Exit statuses
 OK, WARNING, CRITICAL, UNKNOWN = 0, 1, 2, 3
 
-if BASE.metadata.schema:
-    schema = BASE.metadata.schema + '.'
-else:
-    schema = ''
-
-queue_sql = """SELECT COUNT(*) FROM {schema}messages""".format(schema=schema)
-
-PROM_SERVERS = config_get('monitor', 'prometheus_servers', raise_exception=False, default='')
-if PROM_SERVERS != '':
-    PROM_SERVERS = PROM_SERVERS.split(',')
-
 if __name__ == "__main__":
     try:
-        registry = CollectorRegistry()
         session = get_session()
-        result = session.execute(queue_sql).fetchall()
-        print('queues.messages %s' % result[0][0])
-        probe_metrics.gauge(name='queues.messages').set(result[0][0])
-        Gauge('hermes_queues_messages', '', registry=registry).set(result[0][0])
-
-        if len(PROM_SERVERS):
-            for server in PROM_SERVERS:
-                try:
-                    push_to_gateway(server.strip(), job='check_messages_to_submit', registry=registry)
-                except:
-                    continue
-
-        if result[0][0] > 100000:
+        statement = select(
+            func.count()
+        ).select_from(
+            models.Message
+        )
+        message_count = session.execute(statement).scalar_one()
+        print(f"queue.messages {message_count}")
+
+        with PrometheusPusher() as manager:
+            (manager.gauge(
+                    "messages_to_submit.queues.messages",
+                    documentation="Messages in queue, to submit")
+                .set(message_count))
+
+        if message_count > 100000:
             sys.exit(WARNING)
-        elif result[0][0] > 1000000:
+        elif message_count > 1000000:
             sys.exit(CRITICAL)
-
-    except Exception as e:
+    except Exception:
+        print(traceback.format_exc())
         sys.exit(UNKNOWN)
     sys.exit(OK)