channels.py 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
  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.dict(),
  53. "id": str(uuid.uuid4()),
  54. "user_id": user_id,
  55. "created_at": int(time.time()),
  56. "updated_at": int(time.time()),
  57. }
  58. )
  59. new_channel = Channel(**channel.model_dump())
  60. db.add(new_channel)
  61. db.commit()
  62. return channel
  63. def get_channels(self) -> list[ChannelModel]:
  64. with get_db() as db:
  65. channels = db.query(Channel).all()
  66. return [ChannelModel.model_validate(channel) for channel in channels]
  67. def get_channels_by_user_id(
  68. self, user_id: str, permission: str = "read"
  69. ) -> list[ChannelModel]:
  70. channels = self.get_channels()
  71. return [
  72. channel
  73. for channel in channels
  74. if channel.user_id == user_id
  75. or has_access(user_id, permission, channel.access_control)
  76. ]
  77. def get_channel_by_id(self, id: str) -> Optional[ChannelModel]:
  78. with get_db() as db:
  79. channel = db.query(Channel).filter(Channel.id == id).first()
  80. return ChannelModel.model_validate(channel) if channel else None
  81. def update_channel_by_id(
  82. self, id: str, form_data: ChannelForm
  83. ) -> Optional[ChannelModel]:
  84. with get_db() as db:
  85. channel = db.query(Channel).filter(Channel.id == id).first()
  86. if not channel:
  87. return None
  88. channel.name = form_data.name
  89. channel.data = form_data.data
  90. channel.meta = form_data.meta
  91. channel.access_control = form_data.access_control
  92. channel.updated_at = int(time.time())
  93. db.commit()
  94. return ChannelModel.model_validate(channel) if channel else None
  95. def delete_channel_by_id(self, id: str):
  96. with get_db() as db:
  97. db.query(Channel).filter(Channel.id == id).delete()
  98. db.commit()
  99. return True
  100. Channels = ChannelTable()