channels.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  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. data = Column(JSON, nullable=True)
  20. meta = Column(JSON, nullable=True)
  21. access_control = Column(JSON, nullable=True)
  22. created_at = Column(BigInteger)
  23. updated_at = Column(BigInteger)
  24. class ChannelModel(BaseModel):
  25. model_config = ConfigDict(from_attributes=True)
  26. id: str
  27. user_id: str
  28. name: str
  29. data: Optional[dict] = None
  30. meta: Optional[dict] = None
  31. access_control: Optional[dict] = None
  32. created_at: int # timestamp in epoch
  33. updated_at: int # timestamp in epoch
  34. ####################
  35. # Forms
  36. ####################
  37. class ChannelForm(BaseModel):
  38. name: str
  39. data: Optional[dict] = None
  40. meta: Optional[dict] = None
  41. access_control: Optional[dict] = None
  42. class ChannelTable:
  43. def insert_new_channel(
  44. self, form_data: ChannelForm, user_id: str
  45. ) -> Optional[ChannelModel]:
  46. with get_db() as db:
  47. channel = ChannelModel(
  48. **{
  49. **form_data.dict(),
  50. "id": str(uuid.uuid4()),
  51. "user_id": user_id,
  52. "created_at": int(time.time()),
  53. "updated_at": int(time.time()),
  54. }
  55. )
  56. new_channel = Channel(**channel.model_dump())
  57. db.add(new_channel)
  58. db.commit()
  59. return channel
  60. def get_channels(self) -> list[ChannelModel]:
  61. with get_db() as db:
  62. channels = db.query(Channel).all()
  63. return [ChannelModel.model_validate(channel) for channel in channels]
  64. def get_channels_by_user_id(
  65. self, user_id: str, permission: str = "read"
  66. ) -> list[ChannelModel]:
  67. channels = self.get_channels()
  68. return [
  69. channel
  70. for channel in channels
  71. if channel.user_id == user_id
  72. or has_access(user_id, permission, channel.access_control)
  73. ]
  74. def get_channel_by_id(self, id: str) -> Optional[ChannelModel]:
  75. with get_db() as db:
  76. channel = db.query(Channel).filter(Channel.id == id).first()
  77. return ChannelModel.model_validate(channel) if channel else None
  78. def update_channel_by_id(
  79. self, id: str, form_data: ChannelForm
  80. ) -> Optional[ChannelModel]:
  81. with get_db() as db:
  82. channel = db.query(Channel).filter(Channel.id == id).first()
  83. if not channel:
  84. return None
  85. channel.name = form_data.name
  86. channel.data = form_data.data
  87. channel.meta = form_data.meta
  88. channel.access_control = form_data.access_control
  89. channel.updated_at = int(time.time())
  90. db.commit()
  91. return ChannelModel.model_validate(channel) if channel else None
  92. def delete_channel_by_id(self, id: str):
  93. with get_db() as db:
  94. db.query(Channel).filter(Channel.id == id).delete()
  95. db.commit()
  96. return True
  97. Channels = ChannelTable()