From 2b3993a07be5de5f9a817da461fc66dc3247474a Mon Sep 17 00:00:00 2001
From: Karel van Klink <karel.vanklink@geant.org>
Date: Wed, 27 Nov 2024 11:51:40 +0100
Subject: [PATCH] Add prefix limit as an attribute to BGP session

---
 gso/cli/imports.py                            |  3 +-
 ...fff041f9_add_prefix_limit_to_bgpsession.py | 41 +++++++++++++++++++
 gso/products/product_blocks/bgp_session.py    |  6 ++-
 .../create_imported_l3_core_service.py        |  3 +-
 .../l3_core_service/create_l3_core_service.py |  8 ++--
 .../l3_core_service/modify_l3_core_service.py |  8 ++--
 test/fixtures/l3_core_service_fixtures.py     |  3 ++
 test/services/test_infoblox.py                |  2 +-
 .../test_modify_l3_core_service.py            |  2 +
 9 files changed, 66 insertions(+), 10 deletions(-)
 create mode 100644 gso/migrations/versions/2024-11-27_543afff041f9_add_prefix_limit_to_bgpsession.py

diff --git a/gso/cli/imports.py b/gso/cli/imports.py
index 85eba733..bca9ccd4 100644
--- a/gso/cli/imports.py
+++ b/gso/cli/imports.py
@@ -13,7 +13,7 @@ import yaml
 from orchestrator.db import db
 from orchestrator.services.processes import start_process
 from orchestrator.types import SubscriptionLifecycle, UUIDstr
-from pydantic import BaseModel, ValidationError, field_validator, model_validator
+from pydantic import BaseModel, NonNegativeInt, ValidationError, field_validator, model_validator
 from sqlalchemy.exc import SQLAlchemyError
 
 from gso.db.models import PartnerTable
@@ -254,6 +254,7 @@ class L3CoreServiceImportModel(BaseModel):
         families: list[IPFamily]
         is_multi_hop: bool
         rtbh_enabled: bool  # whether Remote Triggered Blackhole is enabled
+        prefix_limit: NonNegativeInt | None = None
 
     class BFDSettingsModel(BaseModel):
         """BFD Settings model."""
diff --git a/gso/migrations/versions/2024-11-27_543afff041f9_add_prefix_limit_to_bgpsession.py b/gso/migrations/versions/2024-11-27_543afff041f9_add_prefix_limit_to_bgpsession.py
new file mode 100644
index 00000000..a239b880
--- /dev/null
+++ b/gso/migrations/versions/2024-11-27_543afff041f9_add_prefix_limit_to_bgpsession.py
@@ -0,0 +1,41 @@
+"""Add prefix limit to BGPSession.
+
+Revision ID: 543afff041f9
+Revises: 2746f861a765
+Create Date: 2024-11-27 10:34:29.855749
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = '543afff041f9'
+down_revision = '2746f861a765'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    conn.execute(sa.text("""
+INSERT INTO resource_types (resource_type, description) VALUES ('prefix_limit', 'Prefix limit for a BGP session') RETURNING resource_types.resource_type_id
+    """))
+    conn.execute(sa.text("""
+INSERT INTO product_block_resource_types (product_block_id, resource_type_id) VALUES ((SELECT product_blocks.product_block_id FROM product_blocks WHERE product_blocks.name IN ('BGPSession')), (SELECT resource_types.resource_type_id FROM resource_types WHERE resource_types.resource_type IN ('prefix_limit')))
+    """))
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    conn.execute(sa.text("""
+DELETE FROM product_block_resource_types WHERE product_block_resource_types.product_block_id IN (SELECT product_blocks.product_block_id FROM product_blocks WHERE product_blocks.name IN ('BGPSession')) AND product_block_resource_types.resource_type_id = (SELECT resource_types.resource_type_id FROM resource_types WHERE resource_types.resource_type IN ('prefix_limit'))
+    """))
+    conn.execute(sa.text("""
+DELETE FROM subscription_instance_values USING product_block_resource_types WHERE subscription_instance_values.subscription_instance_id IN (SELECT subscription_instances.subscription_instance_id FROM subscription_instances WHERE subscription_instances.subscription_instance_id IN (SELECT product_blocks.product_block_id FROM product_blocks WHERE product_blocks.name IN ('BGPSession'))) AND product_block_resource_types.resource_type_id = (SELECT resource_types.resource_type_id FROM resource_types WHERE resource_types.resource_type IN ('prefix_limit'))
+    """))
+    conn.execute(sa.text("""
+DELETE FROM subscription_instance_values WHERE subscription_instance_values.resource_type_id IN (SELECT resource_types.resource_type_id FROM resource_types WHERE resource_types.resource_type IN ('prefix_limit'))
+    """))
+    conn.execute(sa.text("""
+DELETE FROM resource_types WHERE resource_types.resource_type IN ('prefix_limit')
+    """))
diff --git a/gso/products/product_blocks/bgp_session.py b/gso/products/product_blocks/bgp_session.py
index 2e3563b6..2a57eb69 100644
--- a/gso/products/product_blocks/bgp_session.py
+++ b/gso/products/product_blocks/bgp_session.py
@@ -3,7 +3,7 @@
 import strawberry
 from orchestrator.domain.base import ProductBlockModel
 from orchestrator.types import SubscriptionLifecycle
-from pydantic import Field
+from pydantic import Field, NonNegativeInt
 from pydantic_forms.types import strEnum
 
 from gso.utils.types.ip_address import IPAddress
@@ -41,6 +41,7 @@ class BGPSessionInactive(ProductBlockModel, lifecycle=[SubscriptionLifecycle.INI
     rtbh_enabled: bool = False
     bfd_enabled: bool = False
     ip_type: IPTypes | None = None
+    prefix_limit: NonNegativeInt | None = None
 
 
 class BGPSessionProvisioning(BGPSessionInactive, lifecycle=[SubscriptionLifecycle.PROVISIONING]):
@@ -57,6 +58,7 @@ class BGPSessionProvisioning(BGPSessionInactive, lifecycle=[SubscriptionLifecycl
     rtbh_enabled: bool
     bfd_enabled: bool
     ip_type: IPTypes
+    prefix_limit: NonNegativeInt | None
 
 
 class BGPSession(BGPSessionProvisioning, lifecycle=[SubscriptionLifecycle.ACTIVE]):
@@ -84,3 +86,5 @@ class BGPSession(BGPSessionProvisioning, lifecycle=[SubscriptionLifecycle.ACTIVE
     bfd_enabled: bool
     #: The IP type of the session.
     ip_type: IPTypes
+    #: A prefix limit, if required
+    prefix_limit: NonNegativeInt | None
diff --git a/gso/workflows/l3_core_service/create_imported_l3_core_service.py b/gso/workflows/l3_core_service/create_imported_l3_core_service.py
index f63fcec0..0ff4c30f 100644
--- a/gso/workflows/l3_core_service/create_imported_l3_core_service.py
+++ b/gso/workflows/l3_core_service/create_imported_l3_core_service.py
@@ -9,7 +9,7 @@ from orchestrator.types import FormGenerator, SubscriptionLifecycle
 from orchestrator.utils.errors import ProcessFailureError
 from orchestrator.workflow import StepList, begin, done, step
 from orchestrator.workflows.steps import resync, set_status, store_process_subscription
-from pydantic import BaseModel
+from pydantic import BaseModel, NonNegativeInt
 from pydantic_forms.types import UUIDstr
 
 from gso.products import ProductName
@@ -45,6 +45,7 @@ def initial_input_form_generator() -> FormGenerator:
         families: list[IPFamily]
         is_multi_hop: bool
         rtbh_enabled: bool
+        prefix_limit: NonNegativeInt | None = None
 
     class ServiceBindingPort(BaseModel):
         edge_port: UUIDstr
diff --git a/gso/workflows/l3_core_service/create_l3_core_service.py b/gso/workflows/l3_core_service/create_l3_core_service.py
index 67723553..f23a2fcf 100644
--- a/gso/workflows/l3_core_service/create_l3_core_service.py
+++ b/gso/workflows/l3_core_service/create_l3_core_service.py
@@ -10,7 +10,7 @@ from orchestrator.types import FormGenerator, State, SubscriptionLifecycle, UUID
 from orchestrator.workflow import StepList, begin, done, step, workflow
 from orchestrator.workflows.steps import resync, set_status, store_process_subscription
 from orchestrator.workflows.utils import wrap_create_initial_input_form
-from pydantic import BaseModel, ConfigDict, Field, computed_field
+from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt, computed_field
 from pydantic_forms.validators import Divider
 
 from gso.products.product_blocks.bgp_session import BGPSession, IPFamily, IPTypes
@@ -67,6 +67,7 @@ def initial_input_form_generator(product_name: str) -> FormGenerator:
         has_custom_policies: bool = False
         bfd_enabled: bool = False
         multipath_enabled: bool = False
+        prefix_limit: NonNegativeInt | None = None
         is_passive: bool = False
         add_v4_multicast: bool = Field(default=False, exclude=True)
         send_default_route: bool = False
@@ -87,6 +88,7 @@ def initial_input_form_generator(product_name: str) -> FormGenerator:
         has_custom_policies: bool = False
         bfd_enabled: bool = False
         multipath_enabled: bool = False
+        prefix_limit: NonNegativeInt | None = None
         is_passive: bool = False
         add_v6_multicast: bool = Field(default=False, exclude=True)
         send_default_route: bool = False
@@ -118,12 +120,12 @@ def initial_input_form_generator(product_name: str) -> FormGenerator:
         v4_label: Label = Field("IPV4 SBP interface params", exclude=True)
         ipv4_address: IPv4AddressType
         ipv4_mask: IPV4Netmask
-        v4_bfd_settings: BFDSettingsForm = BFDSettingsForm(bfd_enabled=False)
+        v4_bfd_settings: BFDSettingsForm
         divider2: Divider = Field(None, exclude=True)
         v6_label: Label = Field("IPV6 SBP interface params", exclude=True)
         ipv6_address: IPv6AddressType
         ipv6_mask: IPV6Netmask
-        v6_bfd_settings: BFDSettingsForm = BFDSettingsForm(bfd_enabled=False)
+        v6_bfd_settings: BFDSettingsForm
         divider3: Divider = Field(None, exclude=True)
         v4_bgp_peer: IPv4BGPPeer
         v6_bgp_peer: IPv6BGPPeer
diff --git a/gso/workflows/l3_core_service/modify_l3_core_service.py b/gso/workflows/l3_core_service/modify_l3_core_service.py
index f165424b..3073b6de 100644
--- a/gso/workflows/l3_core_service/modify_l3_core_service.py
+++ b/gso/workflows/l3_core_service/modify_l3_core_service.py
@@ -10,7 +10,7 @@ from orchestrator.types import FormGenerator, UUIDstr
 from orchestrator.workflow import StepList
 from orchestrator.workflows.steps import resync, store_process_subscription, unsync
 from orchestrator.workflows.utils import wrap_modify_initial_input_form
-from pydantic import AfterValidator, BaseModel, ConfigDict, Field, computed_field
+from pydantic import AfterValidator, BaseModel, ConfigDict, Field, NonNegativeInt, computed_field
 from pydantic_forms.types import State
 from pydantic_forms.validators import Divider, Label
 
@@ -65,10 +65,11 @@ def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
 
     class IPv4BGPPeer(BaseModel):
         peer_address: IPv4AddressType
-        authentication_key: str | None
+        authentication_key: str | None = None
         has_custom_policies: bool = False
         bfd_enabled: bool = False
         multipath_enabled: bool = False
+        prefix_limit: NonNegativeInt | None = None
         is_passive: bool = False
         add_v4_multicast: bool = Field(default=False, exclude=True)
         send_default_route: bool = False
@@ -85,10 +86,11 @@ def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
 
     class IPv6BGPPeer(BaseModel):
         peer_address: IPv6AddressType
-        authentication_key: str | None
+        authentication_key: str | None = None
         has_custom_policies: bool = False
         bfd_enabled: bool = False
         multipath_enabled: bool = False
+        prefix_limit: NonNegativeInt | None = None
         is_passive: bool = False
         add_v6_multicast: bool = Field(default=False, exclude=True)
         send_default_route: bool = False
diff --git a/test/fixtures/l3_core_service_fixtures.py b/test/fixtures/l3_core_service_fixtures.py
index b37d71f7..da94298e 100644
--- a/test/fixtures/l3_core_service_fixtures.py
+++ b/test/fixtures/l3_core_service_fixtures.py
@@ -5,6 +5,7 @@ import pytest
 from orchestrator.db import db
 from orchestrator.domain import SubscriptionModel
 from orchestrator.types import SubscriptionLifecycle, UUIDstr
+from pydantic import NonNegativeInt
 
 from gso.products import ProductName
 from gso.products.product_blocks.bgp_session import BGPSession, IPFamily, IPTypes
@@ -51,6 +52,7 @@ def bgp_session_subscription_factory(faker):
         is_multi_hop: bool = False,
         has_custom_policies: bool = False,
         multipath_enabled: bool | None = True,
+        prefix_limit: NonNegativeInt | None = None,
         send_default_route: bool | None = True,
         is_passive: bool | None = False,
         rtbh_enabled: bool | None = False,
@@ -65,6 +67,7 @@ def bgp_session_subscription_factory(faker):
             has_custom_policies=has_custom_policies,
             authentication_key=authentication_key or faker.password(),
             multipath_enabled=multipath_enabled,
+            prefix_limit=prefix_limit,
             send_default_route=send_default_route,
             is_multi_hop=is_multi_hop,
             rtbh_enabled=rtbh_enabled,
diff --git a/test/services/test_infoblox.py b/test/services/test_infoblox.py
index 93d526db..3a733260 100644
--- a/test/services/test_infoblox.py
+++ b/test/services/test_infoblox.py
@@ -62,7 +62,7 @@ def _set_up_host_responses():
     responses.add(
         method=responses.GET,
         url="https://10.0.0.1/wapi/v2.12/record%3Ahost?name=test.lo.geant.net&_return_fields=extattrs%2Cipv4addrs%2Cnam"
-        "e%2Cview%2Ciases",
+        "e%2Cview%2Caliases",
         json=[],
     )
 
diff --git a/test/workflows/l3_core_service/test_modify_l3_core_service.py b/test/workflows/l3_core_service/test_modify_l3_core_service.py
index 2b84f59b..67ff8880 100644
--- a/test/workflows/l3_core_service/test_modify_l3_core_service.py
+++ b/test/workflows/l3_core_service/test_modify_l3_core_service.py
@@ -92,6 +92,7 @@ def test_modify_l3_core_service_add_new_edge_port_success(
                 "authentication_key": faker.password(),
                 "peer_address": faker.ipv4(),
                 "bfd_enabled": False,
+                "prefix_limit": 1000,
             },
             "v6_bgp_peer": {
                 "authentication_key": faker.password(),
@@ -151,6 +152,7 @@ def sbp_input_form_data(faker):
                 "is_passive": True,
                 "peer_address": faker.ipv6(),
                 "add_v6_multicast": True,
+                "prefix_limit": 3000,
             },
         }
 
-- 
GitLab