python-workflows/src/workflows/util/zocalo/configuration.py at 52fc2d5b1c2c61b0f2ba9ac4db5943c91684a39a · DiamondLightSource/python-workflows · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
from __future__ import annotations

from marshmallow import fields
from zocalo.configuration import PluginSchema

import workflows.transport
from workflows.transport.pika_transport import PikaTransport
from workflows.transport.stomp_transport import StompTransport


class OTEL:
    """A Zocalo configuration plugin to pre-populate OTELTracing config defaults"""

    class Schema(PluginSchema):
        host = fields.Str(required=True)
        port = fields.Int(required=True)
        timeout = fields.Int(required=False, load_default=10)

    # Store configuration for access by services
    config = {}

    @staticmethod
    def activate(configuration):
        # Build the full endpoint URL
        endpoint = f"https://{configuration['host']}:{configuration['port']}/v1/traces"
        OTEL.config["endpoint"] = endpoint
        OTEL.config["timeout"] = configuration.get("timeout", 10)
        return OTEL.config


class Stomp:
    """A Zocalo configuration plugin to pre-populate StompTransport config defaults"""

    class Schema(PluginSchema):
        host = fields.Str(required=True)
        port = fields.Int(required=True)
        username = fields.Str(required=True)
        password = fields.Str(required=True)
        prefix = fields.Str(required=True)

    @staticmethod
    def activate(configuration):
        for cfgoption, target in [
            ("host", "--stomp-host"),
            ("port", "--stomp-port"),
            ("username", "--stomp-user"),
            ("password", "--stomp-pass"),
            ("prefix", "--stomp-prfx"),
        ]:
            StompTransport.defaults[target] = configuration[cfgoption]


class Pika:
    """A Zocalo configuration plugin to pre-populate PikaTransport config defaults"""

    class Schema(PluginSchema):
        host = fields.Str(required=True)
        port = fields.Field(required=True)
        username = fields.Str(required=True)
        password = fields.Str(required=True)
        vhost = fields.Str(required=True)

    @staticmethod
    def activate(configuration):
        for cfgoption, target in [
            ("host", "--rabbit-host"),
            ("port", "--rabbit-port"),
            ("username", "--rabbit-user"),
            ("password", "--rabbit-pass"),
            ("vhost", "--rabbit-vhost"),
        ]:
            PikaTransport.defaults[target] = configuration[cfgoption]


class DefaultTransport:
    """A Zocalo configuration plugin to set a default transport class"""

    class Schema(PluginSchema):
        default = fields.Str(required=True)

    @staticmethod
    def activate(configuration):
        workflows.transport.default_transport = configuration["default"]