1616"""
1717
1818import re
19- from datetime import datetime
19+ from datetime import datetime , timezone
2020from typing import Any , Dict , List , Optional
2121
2222import jsonschema
@@ -105,29 +105,29 @@ class Base(DeclarativeBase):
105105 """Base class for all models."""
106106
107107
108- # Association table for tools and gateways (federation)
109- tool_gateway_table = Table (
110- "tool_gateway_association" ,
111- Base .metadata ,
112- Column ("tool_id" , Integer , ForeignKey ("tools.id" ), primary_key = True ),
113- Column ("gateway_id" , Integer , ForeignKey ("gateways.id" ), primary_key = True ),
114- )
115-
116- # Association table for resources and gateways (federation)
117- resource_gateway_table = Table (
118- "resource_gateway_association" ,
119- Base .metadata ,
120- Column ("resource_id" , Integer , ForeignKey ("resources.id" ), primary_key = True ),
121- Column ("gateway_id" , Integer , ForeignKey ("gateways.id" ), primary_key = True ),
122- )
123-
124- # Association table for prompts and gateways (federation)
125- prompt_gateway_table = Table (
126- "prompt_gateway_association" ,
127- Base .metadata ,
128- Column ("prompt_id" , Integer , ForeignKey ("prompts.id" ), primary_key = True ),
129- Column ("gateway_id" , Integer , ForeignKey ("gateways.id" ), primary_key = True ),
130- )
108+ # # Association table for tools and gateways (federation)
109+ # tool_gateway_table = Table(
110+ # "tool_gateway_association",
111+ # Base.metadata,
112+ # Column("tool_id", Integer, ForeignKey("tools.id"), primary_key=True),
113+ # Column("gateway_id", Integer, ForeignKey("gateways.id"), primary_key=True),
114+ # )
115+
116+ # # Association table for resources and gateways (federation)
117+ # resource_gateway_table = Table(
118+ # "resource_gateway_association",
119+ # Base.metadata,
120+ # Column("resource_id", Integer, ForeignKey("resources.id"), primary_key=True),
121+ # Column("gateway_id", Integer, ForeignKey("gateways.id"), primary_key=True),
122+ # )
123+
124+ # # Association table for prompts and gateways (federation)
125+ # prompt_gateway_table = Table(
126+ # "prompt_gateway_association",
127+ # Base.metadata,
128+ # Column("prompt_id", Integer, ForeignKey("prompts.id"), primary_key=True),
129+ # Column("gateway_id", Integer, ForeignKey("gateways.id"), primary_key=True),
130+ # )
131131
132132# Association table for servers and tools
133133server_tool_association = Table (
@@ -173,7 +173,7 @@ class ToolMetric(Base):
173173
174174 id : Mapped [int ] = mapped_column (primary_key = True )
175175 tool_id : Mapped [int ] = mapped_column (Integer , ForeignKey ("tools.id" ), nullable = False )
176- timestamp : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
176+ timestamp : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
177177 response_time : Mapped [float ] = mapped_column (Float , nullable = False )
178178 is_success : Mapped [bool ] = mapped_column (Boolean , nullable = False )
179179 error_message : Mapped [Optional [str ]] = mapped_column (Text , nullable = True )
@@ -199,7 +199,7 @@ class ResourceMetric(Base):
199199
200200 id : Mapped [int ] = mapped_column (primary_key = True )
201201 resource_id : Mapped [int ] = mapped_column (Integer , ForeignKey ("resources.id" ), nullable = False )
202- timestamp : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
202+ timestamp : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
203203 response_time : Mapped [float ] = mapped_column (Float , nullable = False )
204204 is_success : Mapped [bool ] = mapped_column (Boolean , nullable = False )
205205 error_message : Mapped [Optional [str ]] = mapped_column (Text , nullable = True )
@@ -225,7 +225,7 @@ class ServerMetric(Base):
225225
226226 id : Mapped [int ] = mapped_column (primary_key = True )
227227 server_id : Mapped [int ] = mapped_column (Integer , ForeignKey ("servers.id" ), nullable = False )
228- timestamp : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
228+ timestamp : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
229229 response_time : Mapped [float ] = mapped_column (Float , nullable = False )
230230 is_success : Mapped [bool ] = mapped_column (Boolean , nullable = False )
231231 error_message : Mapped [Optional [str ]] = mapped_column (Text , nullable = True )
@@ -251,7 +251,7 @@ class PromptMetric(Base):
251251
252252 id : Mapped [int ] = mapped_column (primary_key = True )
253253 prompt_id : Mapped [int ] = mapped_column (Integer , ForeignKey ("prompts.id" ), nullable = False )
254- timestamp : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
254+ timestamp : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
255255 response_time : Mapped [float ] = mapped_column (Float , nullable = False )
256256 is_success : Mapped [bool ] = mapped_column (Boolean , nullable = False )
257257 error_message : Mapped [Optional [str ]] = mapped_column (Text , nullable = True )
@@ -267,7 +267,7 @@ class Tool(Base):
267267 Supports both local tools and federated tools from other gateways.
268268 The integration_type field indicates the tool format:
269269 - "MCP" for MCP-compliant tools (default)
270- - "ICA " for non-MCP ICA Integrations
270+ - "REST " for REST tools
271271
272272 Additionally, this model provides computed properties for aggregated metrics based
273273 on the associated ToolMetric records. These include:
@@ -296,28 +296,28 @@ class Tool(Base):
296296
297297 id : Mapped [int ] = mapped_column (primary_key = True )
298298 name : Mapped [str ] = mapped_column (unique = True )
299- url : Mapped [str ]
299+ url : Mapped [str ] = mapped_column ( String , nullable = True )
300300 description : Mapped [Optional [str ]]
301301 integration_type : Mapped [str ] = mapped_column (default = "MCP" )
302+ request_type : Mapped [str ] = mapped_column (default = "SSE" )
302303 headers : Mapped [Optional [Dict [str , str ]]] = mapped_column (JSON )
303304 input_schema : Mapped [Dict [str , Any ]] = mapped_column (JSON )
304- created_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
305- updated_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow , onupdate = datetime .utcnow )
305+ created_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
306+ updated_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) , onupdate = lambda : datetime .now ( timezone . utc ) )
306307 is_active : Mapped [bool ] = mapped_column (default = True )
307308 jsonpath_filter : Mapped [str ] = mapped_column (default = "" )
308309
309310 # Request type and authentication fields
310- request_type : Mapped [str ] = mapped_column (default = "SSE" )
311311 auth_type : Mapped [Optional [str ]] = mapped_column (default = None ) # "basic", "bearer", or None
312312 auth_value : Mapped [Optional [str ]] = mapped_column (default = None )
313313
314314 # Federation relationship with a local gateway
315315 gateway_id : Mapped [Optional [int ]] = mapped_column (ForeignKey ("gateways.id" ))
316- gateway = relationship ("Gateway" , back_populates = "tools" )
317- federated_with = relationship ("Gateway" , secondary = tool_gateway_table , back_populates = "federated_tools" )
316+ gateway : Mapped [ "Gateway" ] = relationship ("Gateway" , back_populates = "tools" )
317+ # federated_with = relationship("Gateway", secondary=tool_gateway_table, back_populates="federated_tools")
318318
319319 # Many-to-many relationship with Servers
320- servers = relationship ("Server" , secondary = server_tool_association , back_populates = "tools" )
320+ servers : Mapped [ List [ "Server" ]] = relationship ("Server" , secondary = server_tool_association , back_populates = "tools" )
321321
322322 # Relationship with ToolMetric records
323323 metrics : Mapped [List ["ToolMetric" ]] = relationship ("ToolMetric" , back_populates = "tool" , cascade = "all, delete-orphan" )
@@ -479,8 +479,8 @@ class Resource(Base):
479479 mime_type : Mapped [Optional [str ]]
480480 size : Mapped [Optional [int ]]
481481 template : Mapped [Optional [str ]] # URI template for parameterized resources
482- created_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
483- updated_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow , onupdate = datetime .utcnow )
482+ created_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
483+ updated_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) , onupdate = lambda : datetime .now ( timezone . utc ) )
484484 is_active : Mapped [bool ] = mapped_column (default = True )
485485 metrics : Mapped [List ["ResourceMetric" ]] = relationship ("ResourceMetric" , back_populates = "resource" , cascade = "all, delete-orphan" )
486486
@@ -492,11 +492,11 @@ class Resource(Base):
492492 subscriptions : Mapped [List ["ResourceSubscription" ]] = relationship ("ResourceSubscription" , back_populates = "resource" , cascade = "all, delete-orphan" )
493493
494494 gateway_id : Mapped [Optional [int ]] = mapped_column (ForeignKey ("gateways.id" ))
495- gateway = relationship ("Gateway" , back_populates = "resources" )
496- federated_with = relationship ("Gateway" , secondary = resource_gateway_table , back_populates = "federated_resources" )
495+ gateway : Mapped [ "Gateway" ] = relationship ("Gateway" , back_populates = "resources" )
496+ # federated_with = relationship("Gateway", secondary=resource_gateway_table, back_populates="federated_resources")
497497
498498 # Many-to-many relationship with Servers
499- servers = relationship ("Server" , secondary = server_resource_association , back_populates = "resources" )
499+ servers : Mapped [ List [ "Server" ]] = relationship ("Server" , secondary = server_resource_association , back_populates = "resources" )
500500
501501 @property
502502 def content (self ) -> ResourceContent :
@@ -636,7 +636,7 @@ class ResourceSubscription(Base):
636636 id : Mapped [int ] = mapped_column (primary_key = True )
637637 resource_id : Mapped [int ] = mapped_column (ForeignKey ("resources.id" ))
638638 subscriber_id : Mapped [str ] # Client identifier
639- created_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
639+ created_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
640640 last_notification : Mapped [Optional [datetime ]] = mapped_column (DateTime )
641641
642642 resource : Mapped ["Resource" ] = relationship (back_populates = "subscriptions" )
@@ -667,17 +667,17 @@ class Prompt(Base):
667667 description : Mapped [Optional [str ]]
668668 template : Mapped [str ] = mapped_column (Text )
669669 argument_schema : Mapped [Dict [str , Any ]] = mapped_column (JSON )
670- created_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
671- updated_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow , onupdate = datetime .utcnow )
670+ created_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
671+ updated_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) , onupdate = lambda : datetime .now ( timezone . utc ) )
672672 is_active : Mapped [bool ] = mapped_column (default = True )
673673 metrics : Mapped [List ["PromptMetric" ]] = relationship ("PromptMetric" , back_populates = "prompt" , cascade = "all, delete-orphan" )
674674
675675 gateway_id : Mapped [Optional [int ]] = mapped_column (ForeignKey ("gateways.id" ))
676- gateway = relationship ("Gateway" , back_populates = "prompts" )
677- federated_with = relationship ("Gateway" , secondary = prompt_gateway_table , back_populates = "federated_prompts" )
676+ gateway : Mapped [ "Gateway" ] = relationship ("Gateway" , back_populates = "prompts" )
677+ # federated_with = relationship("Gateway", secondary=prompt_gateway_table, back_populates="federated_prompts")
678678
679679 # Many-to-many relationship with Servers
680- servers = relationship ("Server" , secondary = server_prompt_association , back_populates = "prompts" )
680+ servers : Mapped [ List [ "Server" ]] = relationship ("Server" , secondary = server_prompt_association , back_populates = "prompts" )
681681
682682 def validate_arguments (self , args : Dict [str , str ]) -> None :
683683 """
@@ -816,15 +816,15 @@ class Server(Base):
816816 name : Mapped [str ] = mapped_column (unique = True )
817817 description : Mapped [Optional [str ]]
818818 icon : Mapped [Optional [str ]]
819- created_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
820- updated_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow , onupdate = datetime .utcnow )
819+ created_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
820+ updated_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) , onupdate = lambda : datetime .now ( timezone . utc ) )
821821 is_active : Mapped [bool ] = mapped_column (default = True )
822822 metrics : Mapped [List ["ServerMetric" ]] = relationship ("ServerMetric" , back_populates = "server" , cascade = "all, delete-orphan" )
823823
824824 # Many-to-many relationships for associated items
825- tools = relationship ("Tool" , secondary = server_tool_association , back_populates = "servers" )
826- resources = relationship ("Resource" , secondary = server_resource_association , back_populates = "servers" )
827- prompts = relationship ("Prompt" , secondary = server_prompt_association , back_populates = "servers" )
825+ tools : Mapped [ List [ "Tool" ]] = relationship ("Tool" , secondary = server_tool_association , back_populates = "servers" )
826+ resources : Mapped [ List [ "Resource" ]] = relationship ("Resource" , secondary = server_resource_association , back_populates = "servers" )
827+ prompts : Mapped [ List [ "Prompt" ]] = relationship ("Prompt" , secondary = server_prompt_association , back_populates = "servers" )
828828
829829 @property
830830 def execution_count (self ) -> int :
@@ -934,8 +934,8 @@ class Gateway(Base):
934934 url : Mapped [str ]
935935 description : Mapped [Optional [str ]]
936936 capabilities : Mapped [Dict [str , Any ]] = mapped_column (JSON )
937- created_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow )
938- updated_at : Mapped [datetime ] = mapped_column (DateTime , default = datetime .utcnow , onupdate = datetime .utcnow )
937+ created_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) )
938+ updated_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ) , default = lambda : datetime .now ( timezone . utc ) , onupdate = lambda : datetime .now ( timezone . utc ) )
939939 is_active : Mapped [bool ] = mapped_column (default = True )
940940 last_seen : Mapped [Optional [datetime ]]
941941
@@ -948,14 +948,14 @@ class Gateway(Base):
948948 # Relationship with local resources this gateway provides
949949 resources : Mapped [List ["Resource" ]] = relationship (back_populates = "gateway" , cascade = "all, delete-orphan" )
950950
951- # Tools federated from this gateway
952- federated_tools : Mapped [List ["Tool" ]] = relationship (secondary = tool_gateway_table , back_populates = "federated_with" )
951+ # # Tools federated from this gateway
952+ # federated_tools: Mapped[List["Tool"]] = relationship(secondary=tool_gateway_table, back_populates="federated_with")
953953
954- # Prompts federated from this resource
955- federated_resources : Mapped [List ["Resource" ]] = relationship (secondary = resource_gateway_table , back_populates = "federated_with" )
954+ # # Prompts federated from this resource
955+ # federated_resources: Mapped[List["Resource"]] = relationship(secondary=resource_gateway_table, back_populates="federated_with")
956956
957- # Prompts federated from this gateway
958- federated_prompts : Mapped [List ["Prompt" ]] = relationship (secondary = prompt_gateway_table , back_populates = "federated_with" )
957+ # # Prompts federated from this gateway
958+ # federated_prompts: Mapped[List["Prompt"]] = relationship(secondary=prompt_gateway_table, back_populates="federated_with")
959959
960960 # Authorizations
961961 auth_type : Mapped [Optional [str ]] = mapped_column (default = None ) # "basic", "bearer", "headers" or None
@@ -968,21 +968,25 @@ class SessionRecord(Base):
968968 __tablename__ = "mcp_sessions"
969969
970970 session_id : Mapped [str ] = mapped_column (primary_key = True )
971- created_at : Mapped [datetime ] = mapped_column (default = func .now ()) # pylint: disable=not-callable
972- last_accessed : Mapped [datetime ] = mapped_column (default = func .now (), onupdate = func .now ()) # pylint: disable=not-callable
971+ created_at : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ), default = lambda : datetime .now (timezone . utc )) # pylint: disable=not-callable
972+ last_accessed : Mapped [datetime ] = mapped_column (DateTime ( timezone = True ), default = lambda : datetime .now (timezone . utc ), onupdate = lambda : datetime .now (timezone . utc )) # pylint: disable=not-callable
973973 data : Mapped [str ] = mapped_column (String , nullable = True )
974974
975+ messages : Mapped [List ["SessionMessageRecord" ]] = relationship ("SessionMessageRecord" , back_populates = "session" , cascade = "all, delete-orphan" )
976+
975977
976978class SessionMessageRecord (Base ):
977979 """ORM model for messages from SSE client."""
978980
979981 __tablename__ = "mcp_messages"
980982
981983 id : Mapped [int ] = mapped_column (primary_key = True , autoincrement = True )
982- session_id : Mapped [str ] = mapped_column (String )
984+ session_id : Mapped [str ] = mapped_column (ForeignKey ( "mcp_sessions.session_id" ) )
983985 message : Mapped [str ] = mapped_column (String , nullable = True )
984- created_at : Mapped [datetime ] = mapped_column (default = func .now ()) # pylint: disable=not-callable
985- last_accessed : Mapped [datetime ] = mapped_column (default = func .now (), onupdate = func .now ()) # pylint: disable=not-callable
986+ created_at : Mapped [datetime ] = mapped_column (DateTime (timezone = True ), default = lambda : datetime .now (timezone .utc )) # pylint: disable=not-callable
987+ last_accessed : Mapped [datetime ] = mapped_column (DateTime (timezone = True ), default = lambda : datetime .now (timezone .utc ), onupdate = lambda : datetime .now (timezone .utc )) # pylint: disable=not-callable
988+
989+ session : Mapped ["SessionRecord" ] = relationship ("SessionRecord" , back_populates = "messages" )
986990
987991
988992# Event listeners for validation
@@ -1083,6 +1087,7 @@ def init_db():
10831087 Exception: If database initialization fails.
10841088 """
10851089 try :
1090+ # Base.metadata.drop_all(bind=engine)
10861091 Base .metadata .create_all (bind = engine )
10871092 except SQLAlchemyError as e :
10881093 raise Exception (f"Failed to initialize database: { str (e )} " )
0 commit comments