fix a bug when partner_id is None in customer_reolver

unified get_partner_by_name to return object instead of dict

gso/graphql_api/resolvers/customer.py

@@ -11,18 +11,12 @@ from gso.services.partners import get_partner_by_id, get_partner_by_name
 
 async def resolve_customer(root: CustomerType) -> CustomerType:
     """Resolve the customer field for a subscription or process."""
-    if root.customer_id:
-        partner = get_partner_by_id(root.customer_id)
-        return CustomerType(
-            customer_id=partner.partner_id,
-            fullname=partner.name,
-            shortcode=partner.email,
-        )
-    partner = get_partner_by_name("GEANT")
+    partner = get_partner_by_id(root.customer_id) if root.customer_id else get_partner_by_name("GEANT")
+
     return CustomerType(
-        customer_id=partner["partner_id"],
-        fullname=partner["name"],
-        shortcode=partner["email"],
+        customer_id=partner.partner_id,
+        fullname=partner.name,
+        shortcode=partner.email,
     )
 
 

gso/services/partners.py

@@ -7,7 +7,6 @@ from uuid import uuid4
 from orchestrator.db import db
 from pydantic import AfterValidator, BaseModel, ConfigDict, EmailStr, Field
 from sqlalchemy import func
-from sqlalchemy.exc import NoResultFound
 
 from gso.db.models import PartnerTable
 
@@ -59,20 +58,18 @@ class PartnerNotFoundError(Exception):
     """Exception raised when a partner is not found."""
 
 
-def get_all_partners() -> list[dict]:
+def get_all_partners() -> list[PartnerTable]:
     """Fetch all partners from the database and serialize them to JSON."""
-    partners = PartnerTable.query.all()
-    return [partner.__json__() for partner in partners]
+    return PartnerTable.query.all()
 
 
-def get_partner_by_name(name: str) -> dict[str, Any]:
+def get_partner_by_name(name: str) -> PartnerTable:
     """Try to get a partner by their name."""
-    try:
-        partner = db.session.query(PartnerTable).filter(PartnerTable.name == name).one()
-        return partner.__json__()
-    except NoResultFound as e:
-        msg = f"partner {name} not found"
-        raise PartnerNotFoundError(msg) from e
+    partner = db.session.query(PartnerTable).filter_by(name=name).first()
+    if not partner:
+        raise PartnerNotFoundError
+
+    return partner
 
 
 def get_partner_by_id(partner_id: str) -> PartnerTable:

gso/utils/helpers.py

@@ -277,7 +277,7 @@ def active_edge_port_selector(*, partner_id: UUIDstr | None = None) -> Choice:
 
 def partner_choice() -> Choice:
     """Return a Choice object containing a list of available partners."""
-    partners = {partner["partner_id"]: partner["name"] for partner in get_all_partners()}
+    partners = {partner.partner_id: partner.name for partner in get_all_partners()}
 
     return Choice("Select a partner", zip(partners.values(), partners.items(), strict=True))  # type: ignore[arg-type]
 

gso/workflows/edge_port/create_imported_edge_port.py

@@ -27,7 +27,7 @@ from gso.utils.types.interfaces import LAGMember, PhysicalPortCapacity
 @step("Create subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription object."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = get_product_id_by_name(ProductName.IMPORTED_EDGE_PORT)
     subscription = ImportedEdgePortInactive.from_product_id(product_id, partner_id)
 

gso/workflows/iptrunk/create_imported_iptrunk.py

@@ -59,7 +59,7 @@ def initial_input_form_generator() -> FormGenerator:
 @step("Create a new subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription in the service database."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = subscriptions.get_product_id_by_name(ProductName.IMPORTED_IP_TRUNK)
     subscription = ImportedIptrunkInactive.from_product_id(product_id, partner_id)
 

gso/workflows/iptrunk/create_iptrunk.py

@@ -215,7 +215,7 @@ def initial_input_form_generator(product_name: str) -> FormGenerator:
 @step("Create subscription")
 def create_subscription(product: UUIDstr, partner: str) -> State:
     """Create a new subscription object in the database."""
-    subscription = IptrunkInactive.from_product_id(product, get_partner_by_name(partner)["partner_id"])
+    subscription = IptrunkInactive.from_product_id(product, get_partner_by_name(partner).partner_id)
 
     return {
         "subscription": subscription,

gso/workflows/l2_circuit/create_imported_layer_2_circuit.py

@@ -83,7 +83,7 @@ def initial_input_form_generator() -> FormGenerator:
 @step("Create subscription")
 def create_subscription(partner: str, service_type: Layer2CircuitServiceType) -> State:
     """Create a new subscription object."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = get_product_id_by_name(ProductName(service_type))
     subscription = ImportedLayer2CircuitInactive.from_product_id(product_id, partner_id)
 

gso/workflows/l2_circuit/create_layer_2_circuit.py

@@ -29,7 +29,7 @@ from gso.utils.types.virtual_identifiers import VLAN_ID
 
 def initial_input_generator(product_name: str) -> FormGenerator:
     """Gather input from the operator about a new Layer 2 Circuit subscription."""
-    geant_partner_id = get_partner_by_name("GEANT")["partner_id"]
+    geant_partner_id = get_partner_by_name("GEANT").partner_id
 
     class CreateLayer2CircuitServicePage(FormPage):
         model_config = ConfigDict(title=f"{product_name}")

gso/workflows/l3_core_service/create_imported_l3_core_service.py

@@ -80,7 +80,7 @@ def initial_input_form_generator() -> FormGenerator:
 @step("Create subscription")
 def create_subscription(partner: str, service_type: L3CoreServiceType) -> dict:
     """Create a new subscription object in the database."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     match service_type:
         case L3CoreServiceType.GEANT_IP:
             product_id = get_product_id_by_name(ProductName.IMPORTED_GEANT_IP)

gso/workflows/lan_switch_interconnect/create_imported_lan_switch_interconnect.py

@@ -44,7 +44,7 @@ def _initial_input_form_generator() -> FormGenerator:
 @step("Create subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription object."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = get_product_id_by_name(ProductName.IMPORTED_LAN_SWITCH_INTERCONNECT)
     subscription = ImportedLanSwitchInterconnectInactive.from_product_id(product_id, partner_id)
 

gso/workflows/lan_switch_interconnect/create_lan_switch_interconnect.py

@@ -125,7 +125,7 @@ def _initial_input_form(product_name: str) -> FormGenerator:
 @step("Create subscription")
 def create_subscription(product: UUIDstr, partner: str) -> State:
     """Create a new subscription object in the database."""
-    subscription = LanSwitchInterconnectInactive.from_product_id(product, get_partner_by_name(partner)["partner_id"])
+    subscription = LanSwitchInterconnectInactive.from_product_id(product, get_partner_by_name(partner).partner_id)
 
     return {"subscription": subscription, "subscription_id": subscription.subscription_id}
 

gso/workflows/office_router/create_imported_office_router.py

@@ -20,7 +20,7 @@ from gso.utils.types.ip_address import IPv4AddressType, IPv6AddressType, PortNum
 @step("Create subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription object."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = subscriptions.get_product_id_by_name(ProductName.IMPORTED_OFFICE_ROUTER)
     subscription = ImportedOfficeRouterInactive.from_product_id(product_id, partner_id)
 

gso/workflows/opengear/create_imported_opengear.py

@@ -18,7 +18,7 @@ from gso.utils.types.ip_address import IPv4AddressType
 @step("Create subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription object."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = get_product_id_by_name(ProductName.IMPORTED_OPENGEAR)
     subscription = ImportedOpengearInactive.from_product_id(product_id, partner_id)
 

gso/workflows/router/create_imported_router.py

@@ -21,7 +21,7 @@ from gso.utils.types.ip_address import IPv4AddressType, IPv6AddressType, PortNum
 @step("Create subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription object."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = get_product_id_by_name(ProductName.IMPORTED_ROUTER)
     subscription = ImportedRouterInactive.from_product_id(product_id, partner_id)
 

gso/workflows/router/create_router.py

@@ -117,7 +117,7 @@ def initial_input_form_generator(product_name: str) -> FormGenerator:
 @step("Create subscription")
 def create_subscription(product: UUIDstr, partner: str) -> State:
     """Create a new subscription object."""
-    subscription = RouterInactive.from_product_id(product, get_partner_by_name(partner)["partner_id"])
+    subscription = RouterInactive.from_product_id(product, get_partner_by_name(partner).partner_id)
 
     return {
         "subscription": subscription,

gso/workflows/site/create_imported_site.py

@@ -22,7 +22,7 @@ from gso.utils.types.ip_address import IPAddress
 @step("Create subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription object in the service database."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id: UUID = subscriptions.get_product_id_by_name(ProductName.IMPORTED_SITE)
     subscription = ImportedSiteInactive.from_product_id(product_id, partner_id)
 

gso/workflows/site/create_site.py

@@ -52,7 +52,7 @@ def initial_input_form_generator(product_name: str) -> FormGenerator:
 @step("Create subscription")
 def create_subscription(product: UUIDstr, partner: str) -> State:
     """Create a new subscription object in the service database."""
-    subscription = site.SiteInactive.from_product_id(product, get_partner_by_name(partner)["partner_id"])
+    subscription = site.SiteInactive.from_product_id(product, get_partner_by_name(partner).partner_id)
 
     return {
         "subscription": subscription,

gso/workflows/super_pop_switch/create_imported_super_pop_switch.py

@@ -21,7 +21,7 @@ from gso.utils.types.ip_address import IPv4AddressType, PortNumber
 @step("Create subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription object."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = subscriptions.get_product_id_by_name(ProductName.IMPORTED_SUPER_POP_SWITCH)
     subscription = ImportedSuperPopSwitchInactive.from_product_id(product_id, partner_id)
 

gso/workflows/switch/create_imported_switch.py

@@ -32,7 +32,7 @@ def _initial_input_form_generator() -> FormGenerator:
 @step("Create subscription")
 def create_subscription(partner: str) -> State:
     """Create a new subscription object."""
-    partner_id = get_partner_by_name(partner)["partner_id"]
+    partner_id = get_partner_by_name(partner).partner_id
     product_id = get_product_id_by_name(ProductName.IMPORTED_SWITCH)
     subscription = ImportedSwitchInactive.from_product_id(product_id, partner_id)
 

gso/workflows/switch/create_switch.py

@@ -67,7 +67,7 @@ def _initial_input_form_generator(product_name: str) -> FormGenerator:
 @step("Create subscription")
 def create_subscription(product: UUIDstr, partner: str) -> State:
     """Create a new subscription object."""
-    subscription = SwitchInactive.from_product_id(product, get_partner_by_name(partner)["partner_id"])
+    subscription = SwitchInactive.from_product_id(product, get_partner_by_name(partner).partner_id)
 
     return {"subscription": subscription, "subscription_id": subscription.subscription_id}