feat: svc-rabbitmq support policy (#2735)

xogas · web-flow · commit b3cd9f51962c · 2026-04-24T11:08:23.000+08:00
diff --git a/svc-rabbitmq/pyproject.toml b/svc-rabbitmq/pyproject.toml
@@ -54,3 +54,8 @@ dependencies = [
     'urllib3 (==1.26.20)',
     'pyjwt (==2.4.0)'
 ]
+
+
+[tool.pytest.ini_options]
+DJANGO_SETTINGS_MODULE = "svc_rabbitmq.settings"
+python_files = ["test_*.py"]
diff --git a/svc-rabbitmq/svc_rabbitmq/settings/__init__.py b/svc-rabbitmq/svc_rabbitmq/settings/__init__.py
@@ -35,15 +35,15 @@ class PatchFeatures:
 
     @cached_property
     def minimum_database_version(self):
-        if self.connection.mysql_is_mariadb:  # noqa
+        if self.connection.mysql_is_mariadb:
             return (10, 4)
         else:
             return (5, 7)
 
 
 # Django 4.2+ 不再官方支持 Mysql 5.7，但目前 Django 仅是对 5.7 做了软性的不兼容改动，
 # 在没有使用 8.0 特异的功能时，对 5.7 版本的使用无影响，为兼容存量的 Mysql 5.7 DB 做此 Patch
-DatabaseFeatures.minimum_database_version = PatchFeatures.minimum_database_version  # noqa
+DatabaseFeatures.minimum_database_version = PatchFeatures.minimum_database_version
 
 pymysql.install_as_MySQLdb()
 # Patch version info to force pass Django client check
@@ -176,7 +176,7 @@ def minimum_database_version(self):
     "disable_existing_loggers": True,
     "formatters": {
         "verbose": {
-            "format": "%(name)s %(levelname)s [%(asctime)s] %(name)s(ln:%(lineno)d): %(message)s",  # noqa
+            "format": "%(name)s %(levelname)s [%(asctime)s] %(name)s(ln:%(lineno)d): %(message)s",
             "datefmt": "%Y-%m-%d %H:%M:%S",
         },
         "simple": {"format": "%(levelname)s %(message)s"},
@@ -254,6 +254,7 @@ def minimum_database_version(self):
 RABBITMQ_DEFAULT_DEAD_LETTER_QUEUE = env.str("RABBITMQ_DEFAULT_DEAD_LETTER_QUEUE", "dlx.queue")
 RABBITMQ_DEFAULT_DEAD_LETTER_QUEUE_DURABLE = env.bool("RABBITMQ_DEFAULT_DEAD_LETTER_QUEUE_DURABLE", False)
 RABBITMQ_MANAGEMENT_API_CACHE_SECONDS = env.int("RABBITMQ_MANAGEMENT_API_CACHE_SECONDS", 3)
+RABBITMQ_HA_POLICY_ENABLED = env.bool("RABBITMQ_HA_POLICY_ENABLED", True)
 
 MODEL_TAG_CLUSTER = "cluster"
 
diff --git a/svc-rabbitmq/tests/__init__.py b/svc-rabbitmq/tests/__init__.py
@@ -0,0 +1,18 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making
+蓝鲸智云 - PaaS 平台 (BlueKing - PaaS System) available.
+Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except
+in compliance with the License. You may obtain a copy of the License at
+
+    http://opensource.org/licenses/MIT
+
+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.
+
+We undertake not to change the open source license (MIT license) applicable
+to the current version of the project delivered to anyone in the future.
+"""
diff --git a/svc-rabbitmq/tests/conftest.py b/svc-rabbitmq/tests/conftest.py
@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making
+蓝鲸智云 - PaaS 平台 (BlueKing - PaaS System) available.
+Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except
+in compliance with the License. You may obtain a copy of the License at
+
+    http://opensource.org/licenses/MIT
+
+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.
+
+We undertake not to change the open source license (MIT license) applicable
+to the current version of the project delivered to anyone in the future.
+"""
+
+from unittest.mock import MagicMock, create_autospec
+
+import pytest
+from vendor.client import ManagementClient, VirtualHostHandler
+from vendor.clusters import Cluster
+
+
+@pytest.fixture()
+def mock_client() -> MagicMock:
+    """Fixture to provide a mocked ManagementClient."""
+
+    client: ManagementClient = create_autospec(ManagementClient, instance=True)
+    # Ensure nested handlers are plain MagicMocks for easy assertion
+    client.virtual_host = MagicMock(spec=VirtualHostHandler)
+    client.user_policy = MagicMock()
+    client.limit_policy = MagicMock()
+    client.exchange = MagicMock()
+    client.queue = MagicMock()
+    client.user = MagicMock()
+
+    return client
+
+
+def make_cluster(version: str = "3.9.0", **overrides) -> Cluster:
+    """Build a Cluster pydantic model for testing."""
+
+    defaults = {
+        "host": "127.0.0.1",
+        "port": 5672,
+        "management_api": "http://127.0.0.1:15672",
+        "admin": "admin",
+        "password": "password",
+        "version": version,
+        "tls": {},
+    }
+    defaults.update(overrides)
+
+    return Cluster(**defaults)
diff --git a/svc-rabbitmq/tests/test_provider.py b/svc-rabbitmq/tests/test_provider.py
@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making
+蓝鲸智云 - PaaS 平台 (BlueKing - PaaS System) available.
+Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except
+in compliance with the License. You may obtain a copy of the License at
+
+    http://opensource.org/licenses/MIT
+
+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.
+
+We undertake not to change the open source license (MIT license) applicable
+to the current version of the project delivered to anyone in the future.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from vendor.provider import Provider
+
+from .conftest import make_cluster
+
+
+class TestProviderPickCluster:
+    """pick_cluster: fallback from flat fields vs. weighted clusters list."""
+
+    def test_from_flat_fields(self):
+        provider = Provider(
+            host="10.0.0.1",
+            port=5672,
+            management_api="http://10.0.0.1:15672",
+            admin="admin",
+            password="secret",
+            version="3.9.0",
+        )
+        cluster = provider.pick_cluster()
+        assert cluster.host == "10.0.0.1"
+        assert cluster.version == "3.9.0"
+
+    def test_raises_on_invalid_flat_fields(self):
+        provider = Provider(host="10.0.0.1")
+        with pytest.raises(ValueError, match="cluster 配置不正确"):
+            provider.pick_cluster()
+
+    def test_from_clusters_list(self):
+        provider = Provider(
+            clusters=[
+                {
+                    "weight": 100,
+                    "values": {
+                        "host": "10.0.0.2",
+                        "port": 5672,
+                        "management_api": "http://10.0.0.2:15672",
+                        "admin": "admin",
+                        "password": "secret",
+                        "version": "3.8.0",
+                    },
+                }
+            ]
+        )
+        cluster = provider.pick_cluster()
+        assert cluster.host == "10.0.0.2"
+
+
+class TestProviderResolveClusterPK:
+    """resolve_cluster_pk: match by management_api or fallback to host/port."""
+
+    def _db_cluster(self, pk, management_api, host="127.0.0.1", port=5672):
+        obj = MagicMock()
+        obj.pk = pk
+        obj.management_api = management_api
+        obj.host = host
+        obj.port = port
+        return obj
+
+    @patch("vendor.provider.ClusterModel.objects")
+    def test_match_by_management_api(self, mock_objects):
+        db_cluster = self._db_cluster(pk=42, management_api="http://10.0.0.1:15672/")
+        mock_objects.filter.return_value = [db_cluster]
+
+        cluster = make_cluster(management_api="http://10.0.0.1:15672")
+        assert Provider().resolve_cluster_pk(cluster) == 42
+
+    @patch("vendor.provider.ClusterModel.objects")
+    def test_returns_none_when_no_match(self, mock_objects):
+        db_cluster = self._db_cluster(pk=1, management_api="http://other:15672")
+        mock_objects.filter.return_value = [db_cluster]
+
+        cluster = make_cluster(management_api="http://10.0.0.1:15672")
+        assert Provider().resolve_cluster_pk(cluster) is None
diff --git a/svc-rabbitmq/tests/test_provider_plugins.py b/svc-rabbitmq/tests/test_provider_plugins.py
@@ -0,0 +1,127 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making
+蓝鲸智云 - PaaS 平台 (BlueKing - PaaS System) available.
+Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except
+in compliance with the License. You may obtain a copy of the License at
+
+    http://opensource.org/licenses/MIT
+
+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.
+
+We undertake not to change the open source license (MIT license) applicable
+to the current version of the project delivered to anyone in the future.
+"""
+
+import pytest
+from django.test import override_settings
+from vendor.provider import DeadLetterRoutingProviderPlugin, HAProviderPlugin
+
+from .conftest import make_cluster
+
+
+@pytest.fixture()
+def cluster():
+    """Cluster with version 3.9.0 (>= 3.8, supports quorum)."""
+    return make_cluster("3.9.0")
+
+
+@pytest.fixture()
+def old_cluster():
+    """Cluster with version 3.7.0 (< 3.8, classic mirror only)."""
+    return make_cluster("3.7.0")
+
+
+class TestHAProviderPlugin:
+    """HAProviderPlugin: version-based branching between quorum and classic mirror."""
+
+    def test_quorum_path_for_version_gte_3_8(self, cluster, mock_client):
+        """For RabbitMQ >= 3.8, HAProviderPlugin should use quorum queues."""
+        plugin = HAProviderPlugin(context={}, client=mock_client, cluster=cluster, virtual_host="vh")
+        plugin.on_create()
+
+        mock_client.virtual_host.create.assert_called_once_with("vh", default_queue_type="quorum")
+        mock_client.user_policy.create.assert_not_called()
+
+    def test_classic_mirror_path_for_version_lt_3_8(self, old_cluster, mock_client):
+        """For RabbitMQ < 3.8, HAProviderPlugin should use classic mirrored queues."""
+        plugin = HAProviderPlugin(context={}, client=mock_client, cluster=old_cluster, virtual_host="vh")
+        plugin.on_create()
+
+        mock_client.virtual_host.create.assert_not_called()
+        mock_client.user_policy.create.assert_called_once()
+        _, args, _ = mock_client.user_policy.create.mock_calls[0]
+        assert args[0] == "vh"
+        assert args[1] == HAProviderPlugin.HA_POLICY_NAME
+        assert args[2]["definition"]["ha-mode"] == "all"
+
+    @override_settings(RABBITMQ_HA_POLICY_ENABLED=False)
+    def test_noop_when_disabled(self, cluster, mock_client):
+        """Setting disabled → no calls at all."""
+        plugin = HAProviderPlugin(context={}, client=mock_client, cluster=cluster, virtual_host="vh")
+        plugin.on_create()
+
+        mock_client.virtual_host.create.assert_not_called()
+        mock_client.user_policy.create.assert_not_called()
+
+
+class TestDeadLetterRoutingProviderPlugin:
+    """DeadLetterRoutingProviderPlugin: dead letter exchange/queue setup with quorum awareness."""
+
+    DLX_SETTINGS = {
+        "RABBITMQ_DEFAULT_DEAD_LETTER_ROUTING_KEY": "",
+        "RABBITMQ_DEFAULT_DEAD_LETTER_EXCHANGE": "dlx.exchange",
+        "RABBITMQ_DEFAULT_DEAD_LETTER_EXCHANGE_TYPE": "direct",
+        "RABBITMQ_DEFAULT_DEAD_LETTER_EXCHANGE_DURABLE": False,
+        "RABBITMQ_DEFAULT_DEAD_LETTER_QUEUE": "dlx.queue",
+        "RABBITMQ_DEFAULT_DEAD_LETTER_QUEUE_DURABLE": False,
+    }
+
+    @override_settings(**DLX_SETTINGS)
+    def test_quorum_queue_for_version_gte_3_8(self, cluster, mock_client):
+        """DLQ declared as quorum type with durable=True for >= 3.8."""
+        ctx = {}
+        plugin = DeadLetterRoutingProviderPlugin(context=ctx, client=mock_client, cluster=cluster, virtual_host="vh")
+        plugin.on_create()
+
+        _, kwargs = mock_client.queue.declare.call_args
+        assert kwargs["arguments"] == {"x-queue-type": "quorum"}
+        assert kwargs["durable"] is True
+
+    override_settings(**DLX_SETTINGS)
+
+    def test_classic_queue_for_version_lt_3_8(self, old_cluster, mock_client):
+        """DLQ declared without quorum arguments for < 3.8."""
+        plugin = DeadLetterRoutingProviderPlugin(
+            context={}, client=mock_client, cluster=old_cluster, virtual_host="vh"
+        )
+        plugin.on_create()
+
+        _, kwargs = mock_client.queue.declare.call_args
+        assert kwargs["arguments"] == {}
+        assert kwargs["durable"] is False
+
+    @override_settings(**DLX_SETTINGS)
+    def test_skipped_for_version_lt_2_8(self, mock_client):
+        """Versions below 2.8 → on_create is a no-op."""
+        c = make_cluster("2.7.0")
+        plugin = DeadLetterRoutingProviderPlugin(context={}, client=mock_client, cluster=c, virtual_host="vh")
+        plugin.on_create()
+
+        mock_client.exchange.declare.assert_not_called()
+        mock_client.queue.declare.assert_not_called()
+
+    @override_settings(**DLX_SETTINGS)
+    def test_context_populated(self, cluster, mock_client):
+        """After on_create, context should contain dlx-* keys."""
+        ctx = {}
+        plugin = DeadLetterRoutingProviderPlugin(context=ctx, client=mock_client, cluster=cluster, virtual_host="vh")
+        plugin.on_create()
+
+        assert ctx["dlx-exchange"] == "dlx.exchange"
+        assert ctx["dlx-queue"] == "dlx.queue"
+        assert ctx["dlx-routing-key"] == ""
diff --git a/svc-rabbitmq/vendor/client.py b/svc-rabbitmq/vendor/client.py
@@ -81,6 +81,26 @@ def partial(self, path: "str", method: "str", headers=None) -> "typing.Callable"
         return partial(self.request, path=path, method=method, headers=headers)
 
 
+class VirtualHostHandler(VirtualHost):
+    """Extended VirtualHost handler that supports default_queue_type parameter."""
+
+    path = "vhosts/"
+
+    def create(self, name: str, default_queue_type: str | None = None):
+        """Create a virtual host.
+
+        :param name: The name of the virtual host to create.
+        :param default_queue_type: Default queue type for the vhost (Optional, e.g. 'quorum', 'classic', 'stream')
+        """
+        payload = {}
+        if default_queue_type:
+            payload["default_queue_type"] = default_queue_type
+        self.http_client.put(
+            urljoin(self.path, quote(name)),
+            payload=payload,
+        )
+
+
 class DefinitionsHandler(ManagementHandler):
     path = "definitions/"
 
@@ -228,7 +248,7 @@ def __init__(self, api_url, username, password, timeout=10, verify=None, cert=No
         self.health_checks = HealthChecks(self.http_client)
         self.queue = Queue(self.http_client)
         self.user = User(self.http_client)
-        self.virtual_host = VirtualHost(self.http_client)
+        self.virtual_host = VirtualHostHandler(self.http_client)
         self.definitions = DefinitionsHandler(self.http_client)
         self.user_policy = UserPolicyHandler(self.http_client)
         self.limit_policy = LimitPolicyHandler(self.http_client)
diff --git a/svc-rabbitmq/vendor/provider.py b/svc-rabbitmq/vendor/provider.py

Original file line number	Diff line number	Diff line change
`@@ -54,3 +54,8 @@ dependencies = [`
`54`	`54`	`'urllib3 (==1.26.20)',`
`55`	`55`	`'pyjwt (==2.4.0)'`
`56`	`56`	`]`
	`57`	`+`
	`58`	`+`
	`59`	`+[tool.pytest.ini_options]`
	`60`	`+DJANGO_SETTINGS_MODULE = "svc_rabbitmq.settings"`
	`61`	`+python_files = ["test_*.py"]`