channels.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. import json
  2. import time
  3. import uuid
  4. from typing import Optional
  5. from open_webui.internal.db import Base, get_db
  6. from open_webui.utils.access_control import has_access
  7. from pydantic import BaseModel, ConfigDict
  8. from sqlalchemy import BigInteger, Boolean, Column, String, Text, JSON
  9. from sqlalchemy import or_, func, select, and_, text
  10. from sqlalchemy.sql import exists
  11. ####################
  12. # Channel DB Schema
  13. ####################
  14. class Channel(Base):
  15. __tablename__ = "channel"
  16. id = Column(Text, primary_key=True)
  17. user_id = Column(Text)
  18. name = Column(Text)
  19. description = Column(Text, nullable=True)
  20. data = Column(JSON, nullable=True)
  21. meta = Column(JSON, nullable=True)
  22. access_control = Column(JSON, nullable=True)
  23. created_at = Column(BigInteger)
  24. updated_at = Column(BigInteger)
  25. class ChannelModel(BaseModel):
  26. model_config = ConfigDict(from_attributes=True)
  27. id: str
  28. user_id: str
  29. description: Optional[str] = None
  30. name: str
  31. data: Optional[dict] = None
  32. meta: Optional[dict] = None
  33. access_control: Optional[dict] = None
  34. created_at: int # timestamp in epoch
  35. updated_at: int # timestamp in epoch
  36. ####################
  37. # Forms
  38. ####################
  39. class ChannelForm(BaseModel):
  40. name: str
  41. description: Optional[str] = None
  42. data: Optional[dict] = None
  43. meta: Optional[dict] = None
  44. access_control: Optional[dict] = None
  45. class ChannelTable:
  46. def insert_new_channel(
  47. self, form_data: ChannelForm, user_id: str
  48. ) -> Optional[ChannelModel]:
  49. with get_db() as db:
  50. channel = ChannelModel(
  51. **{
  52. **form_data.model_dump(),
  53. "name": form_data.name.lower(),
  54. "id": str(uuid.uuid4()),
  55. "user_id": user_id,
  56. "created_at": int(time.time_ns()),
  57. "updated_at": int(time.time_ns()),
  58. }
  59. )
  60. new_channel = Channel(**channel.model_dump())
  61. db.add(new_channel)
  62. db.commit()
  63. return channel
  64. def get_channels(self) -> list[ChannelModel]:
  65. with get_db() as db:
  66. channels = db.query(Channel).all()
  67. return [ChannelModel.model_validate(channel) for channel in channels]
  68. def get_channels_by_user_id(
  69. self, user_id: str, permission: str = "read"
  70. ) -> list[ChannelModel]:
  71. channels = self.get_channels()
  72. return [
  73. channel
  74. for channel in channels
  75. if channel.user_id == user_id
  76. or has_access(user_id, permission, channel.access_control)
  77. ]
  78. def get_channel_by_id(self, id: str) -> Optional[ChannelModel]:
  79. with get_db() as db:
  80. channel = db.query(Channel).filter(Channel.id == id).first()
  81. return ChannelModel.model_validate(channel) if channel else None
  82. def update_channel_by_id(
  83. self, id: str, form_data: ChannelForm
  84. ) -> Optional[ChannelModel]:
  85. with get_db() as db:
  86. channel = db.query(Channel).filter(Channel.id == id).first()
  87. if not channel:
  88. return None
  89. channel.name = form_data.name
  90. channel.data = form_data.data
  91. channel.meta = form_data.meta
  92. channel.access_control = form_data.access_control
  93. channel.updated_at = int(time.time_ns())
  94. db.commit()
  95. return ChannelModel.model_validate(channel) if channel else None
  96. def delete_channel_by_id(self, id: str):
  97. with get_db() as db:
  98. db.query(Channel).filter(Channel.id == id).delete()
  99. db.commit()
  100. return True
  101. Channels = ChannelTable()