mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
sanitize inputs for id/name
This commit is contained in:
+26
-49
@@ -1,15 +1,18 @@
|
||||
"""Database models for basic-memory."""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import List, Optional
|
||||
from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, UniqueConstraint
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase
|
||||
from sqlalchemy.ext.asyncio import AsyncAttrs
|
||||
|
||||
from basic_memory.utils import normalize_entity_id
|
||||
from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, UniqueConstraint
|
||||
from sqlalchemy.ext.asyncio import AsyncAttrs
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase
|
||||
|
||||
from basic_memory.utils import sanitize_name
|
||||
|
||||
|
||||
class Base(AsyncAttrs, DeclarativeBase):
|
||||
"""Base class for all models"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@@ -24,54 +27,45 @@ class Entity(Base):
|
||||
- A description (optional)
|
||||
- A list of observations
|
||||
"""
|
||||
|
||||
__tablename__ = "entity"
|
||||
__table_args__ = (
|
||||
UniqueConstraint('entity_type', 'name', name='uix_entity_type_name'),
|
||||
)
|
||||
__table_args__ = (UniqueConstraint("entity_type", "name", name="uix_entity_type_name"),)
|
||||
|
||||
id: Mapped[str] = mapped_column(String, primary_key=True)
|
||||
name: Mapped[str] = mapped_column(String)
|
||||
entity_type: Mapped[str] = mapped_column(String)
|
||||
description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime,
|
||||
server_default=text('CURRENT_TIMESTAMP')
|
||||
)
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime,
|
||||
server_default=text('CURRENT_TIMESTAMP'),
|
||||
onupdate=text('CURRENT_TIMESTAMP')
|
||||
DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP")
|
||||
)
|
||||
|
||||
# Relationships
|
||||
observations: Mapped[List["Observation"]] = relationship(
|
||||
"Observation",
|
||||
back_populates="entity",
|
||||
cascade="all, delete-orphan"
|
||||
"Observation", back_populates="entity", cascade="all, delete-orphan"
|
||||
)
|
||||
outgoing_relations: Mapped[List["Relation"]] = relationship(
|
||||
"Relation",
|
||||
foreign_keys="[Relation.from_id]",
|
||||
back_populates="from_entity",
|
||||
cascade="all, delete-orphan"
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
incoming_relations: Mapped[List["Relation"]] = relationship(
|
||||
"Relation",
|
||||
foreign_keys="[Relation.to_id]",
|
||||
back_populates="to_entity",
|
||||
cascade="all, delete-orphan"
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
@property
|
||||
def relations(self):
|
||||
return self.outgoing_relations + self.incoming_relations
|
||||
|
||||
|
||||
@classmethod
|
||||
def generate_id(cls, entity_type: str, name: str) -> str:
|
||||
"""Generate a filesystem path-based ID for this entity."""
|
||||
# Use common normalization for filesystem safety
|
||||
safe_name = normalize_entity_id(name)
|
||||
safe_name = sanitize_name(name)
|
||||
return f"{entity_type}/{safe_name}"
|
||||
|
||||
def get_file_path(self) -> str:
|
||||
@@ -90,26 +84,19 @@ class Observation(Base):
|
||||
- Can be added or removed independently
|
||||
- Should be atomic (one fact per observation)
|
||||
"""
|
||||
|
||||
__tablename__ = "observation"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
||||
entity_id: Mapped[str] = mapped_column( # Reference Entity.id which is text
|
||||
String,
|
||||
ForeignKey("entity.id", ondelete="CASCADE"),
|
||||
index=True
|
||||
String, ForeignKey("entity.id", ondelete="CASCADE"), index=True
|
||||
)
|
||||
content: Mapped[str] = mapped_column(String)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime,
|
||||
server_default=text('CURRENT_TIMESTAMP')
|
||||
)
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
|
||||
context: Mapped[str | None] = mapped_column(String, nullable=True)
|
||||
|
||||
# Relationships
|
||||
entity: Mapped[Entity] = relationship(
|
||||
"Entity",
|
||||
back_populates="observations"
|
||||
)
|
||||
entity: Mapped[Entity] = relationship("Entity", back_populates="observations")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
content = self.content[:50] + "..." if len(self.content) > 50 else self.content
|
||||
@@ -121,37 +108,27 @@ class Relation(Base):
|
||||
Relations define directed connections between entities.
|
||||
They are always stored in active voice and describe how entities interact or relate to each other.
|
||||
"""
|
||||
|
||||
__tablename__ = "relation"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
||||
from_id: Mapped[str] = mapped_column( # Reference Entity.id which is text
|
||||
String,
|
||||
ForeignKey("entity.id", ondelete="CASCADE"),
|
||||
index=True
|
||||
String, ForeignKey("entity.id", ondelete="CASCADE"), index=True
|
||||
)
|
||||
to_id: Mapped[str] = mapped_column( # Reference Entity.id which is text
|
||||
String,
|
||||
ForeignKey("entity.id", ondelete="CASCADE"),
|
||||
index=True
|
||||
String, ForeignKey("entity.id", ondelete="CASCADE"), index=True
|
||||
)
|
||||
relation_type: Mapped[str] = mapped_column(String)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime,
|
||||
server_default=text('CURRENT_TIMESTAMP')
|
||||
)
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
|
||||
context: Mapped[str | None] = mapped_column(String, nullable=True)
|
||||
|
||||
# Relationships
|
||||
from_entity: Mapped[Entity] = relationship(
|
||||
"Entity",
|
||||
foreign_keys=[from_id],
|
||||
back_populates="outgoing_relations"
|
||||
"Entity", foreign_keys=[from_id], back_populates="outgoing_relations"
|
||||
)
|
||||
to_entity: Mapped[Entity] = relationship(
|
||||
"Entity",
|
||||
foreign_keys=[to_id],
|
||||
back_populates="incoming_relations"
|
||||
"Entity", foreign_keys=[to_id], back_populates="incoming_relations"
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Relation(id={self.id}, from='{self.from_id}', type='{self.relation_type}', to='{self.to_id}')"
|
||||
return f"Relation(id={self.id}, from='{self.from_id}', type='{self.relation_type}', to='{self.to_id}')"
|
||||
|
||||
Reference in New Issue
Block a user