refactor: connect client once at server startup via lifespan
This commit is contained in:
@@ -1,5 +1,7 @@
|
|||||||
import os
|
import os
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
from sabatron_tryton_rpc_client.client import Client
|
from sabatron_tryton_rpc_client.client import Client
|
||||||
|
|
||||||
@@ -19,6 +21,7 @@ class TrytonSettings:
|
|||||||
port: int = field(
|
port: int = field(
|
||||||
default_factory=lambda: int(os.environ.get("TRYTON_PORT", "8000"))
|
default_factory=lambda: int(os.environ.get("TRYTON_PORT", "8000"))
|
||||||
)
|
)
|
||||||
|
_client: Optional[Client] = field(default=None, init=False, repr=False)
|
||||||
|
|
||||||
def to_dict(self) -> dict:
|
def to_dict(self) -> dict:
|
||||||
return {
|
return {
|
||||||
@@ -30,7 +33,28 @@ class TrytonSettings:
|
|||||||
}
|
}
|
||||||
|
|
||||||
def get_client(self) -> Client:
|
def get_client(self) -> Client:
|
||||||
return Client(**self.to_dict())
|
if self._client is None:
|
||||||
|
self._client = Client(**self.to_dict())
|
||||||
|
return self._client
|
||||||
|
|
||||||
|
def connect(self):
|
||||||
|
client = self.get_client()
|
||||||
|
client.connect()
|
||||||
|
|
||||||
|
def disconnect(self):
|
||||||
|
if self._client is not None:
|
||||||
|
close_method = getattr(self._client, "close", None) or getattr(
|
||||||
|
self._client, "disconnect", None
|
||||||
|
)
|
||||||
|
if close_method:
|
||||||
|
close_method()
|
||||||
|
self._client = None
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def lifespan(self):
|
||||||
|
self.connect()
|
||||||
|
yield
|
||||||
|
self.disconnect()
|
||||||
|
|
||||||
def __post_init__(self):
|
def __post_init__(self):
|
||||||
if not self.hostname:
|
if not self.hostname:
|
||||||
|
|||||||
@@ -2,15 +2,22 @@
|
|||||||
MCP Server for Naliia Module
|
MCP Server for Naliia Module
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from sabatron_tryton_rpc_client.client import Client
|
from contextlib import asynccontextmanager
|
||||||
from fastmcp import FastMCP
|
from fastmcp import FastMCP
|
||||||
from typing import List
|
from typing import List, AsyncIterator
|
||||||
import datetime
|
import datetime
|
||||||
|
|
||||||
from tryton_mcp.config import settings
|
from tryton_mcp.config import settings
|
||||||
|
|
||||||
|
|
||||||
mcp = FastMCP("Tryton MCP Server")
|
@asynccontextmanager
|
||||||
|
async def server_lifespan(server: FastMCP) -> AsyncIterator[None]:
|
||||||
|
settings.connect()
|
||||||
|
yield
|
||||||
|
settings.disconnect()
|
||||||
|
|
||||||
|
|
||||||
|
mcp = FastMCP("Tryton MCP Server", lifespan=server_lifespan)
|
||||||
|
|
||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
@@ -28,9 +35,7 @@ def create_schedule(
|
|||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|
||||||
client = Client(**settings.to_dict())
|
client = settings.get_client()
|
||||||
|
|
||||||
client.connect()
|
|
||||||
name = "model.naliia.schedule.create"
|
name = "model.naliia.schedule.create"
|
||||||
args = [example_schedule, {}]
|
args = [example_schedule, {}]
|
||||||
response = client.call(name, args)
|
response = client.call(name, args)
|
||||||
@@ -40,9 +45,7 @@ def create_schedule(
|
|||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
def find_service_centers():
|
def find_service_centers():
|
||||||
client = Client(**settings.to_dict())
|
client = settings.get_client()
|
||||||
|
|
||||||
client.connect()
|
|
||||||
|
|
||||||
name = "model.naliia.service_center.search_read"
|
name = "model.naliia.service_center.search_read"
|
||||||
args = [[[]], 0, None, None, ["name", "address.street"], {}]
|
args = [[[]], 0, None, None, ["name", "address.street"], {}]
|
||||||
@@ -53,9 +56,7 @@ def find_service_centers():
|
|||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
def find_products_and_services():
|
def find_products_and_services():
|
||||||
client = Client(**settings.to_dict())
|
client = settings.get_client()
|
||||||
|
|
||||||
client.connect()
|
|
||||||
|
|
||||||
name = "model.product.product.search_read"
|
name = "model.product.product.search_read"
|
||||||
args = [
|
args = [
|
||||||
|
|||||||
Reference in New Issue
Block a user