Kevin Hu commited on
Commit
e6cd231
·
1 Parent(s): 9aadaa2

add owner check for team work (#2892)

Browse files

### What problem does this PR solve?

#2834

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

agent/component/exesql.py CHANGED
@@ -45,6 +45,9 @@ class ExeSQLParam(ComponentParamBase):
45
  self.check_positive_integer(self.port, "IP Port")
46
  self.check_empty(self.password, "Database password")
47
  self.check_positive_integer(self.top_n, "Number of records")
 
 
 
48
 
49
 
50
  class ExeSQL(ComponentBase, ABC):
 
45
  self.check_positive_integer(self.port, "IP Port")
46
  self.check_empty(self.password, "Database password")
47
  self.check_positive_integer(self.top_n, "Number of records")
48
+ if self.database == "rag_flow":
49
+ if self.host == "ragflow-mysql": raise ValueError("The host is not accessible.")
50
+ if self.password == "infini_rag_flow": raise ValueError("The host is not accessible.")
51
 
52
 
53
  class ExeSQL(ComponentBase, ABC):
api/apps/document_app.py CHANGED
@@ -209,9 +209,17 @@ def list_docs():
209
 
210
 
211
  @manager.route('/infos', methods=['POST'])
 
212
  def docinfos():
213
  req = request.json
214
  doc_ids = req["doc_ids"]
 
 
 
 
 
 
 
215
  docs = DocumentService.get_by_ids(doc_ids)
216
  return get_json_result(data=list(docs.dicts()))
217
 
@@ -242,11 +250,17 @@ def thumbnails():
242
  def change_status():
243
  req = request.json
244
  if str(req["status"]) not in ["0", "1"]:
245
- get_json_result(
246
  data=False,
247
  retmsg='"Status" must be either 0 or 1!',
248
  retcode=RetCode.ARGUMENT_ERROR)
249
 
 
 
 
 
 
 
250
  try:
251
  e, doc = DocumentService.get_by_id(req["doc_id"])
252
  if not e:
@@ -285,6 +299,15 @@ def rm():
285
  req = request.json
286
  doc_ids = req["doc_id"]
287
  if isinstance(doc_ids, str): doc_ids = [doc_ids]
 
 
 
 
 
 
 
 
 
288
  root_folder = FileService.get_root_folder(current_user.id)
289
  pf_id = root_folder["id"]
290
  FileService.init_knowledgebase_docs(pf_id, current_user.id)
@@ -323,6 +346,13 @@ def rm():
323
  @validate_request("doc_ids", "run")
324
  def run():
325
  req = request.json
 
 
 
 
 
 
 
326
  try:
327
  for id in req["doc_ids"]:
328
  info = {"run": str(req["run"]), "progress": 0}
@@ -356,6 +386,12 @@ def run():
356
  @validate_request("doc_id", "name")
357
  def rename():
358
  req = request.json
 
 
 
 
 
 
359
  try:
360
  e, doc = DocumentService.get_by_id(req["doc_id"])
361
  if not e:
@@ -416,6 +452,13 @@ def get(doc_id):
416
  @validate_request("doc_id", "parser_id")
417
  def change_parser():
418
  req = request.json
 
 
 
 
 
 
 
419
  try:
420
  e, doc = DocumentService.get_by_id(req["doc_id"])
421
  if not e:
 
209
 
210
 
211
  @manager.route('/infos', methods=['POST'])
212
+ @login_required
213
  def docinfos():
214
  req = request.json
215
  doc_ids = req["doc_ids"]
216
+ for doc_id in doc_ids:
217
+ if not DocumentService.accessible(doc_id, current_user.id):
218
+ return get_json_result(
219
+ data=False,
220
+ retmsg='No authorization.',
221
+ retcode=RetCode.AUTHENTICATION_ERROR
222
+ )
223
  docs = DocumentService.get_by_ids(doc_ids)
224
  return get_json_result(data=list(docs.dicts()))
225
 
 
250
  def change_status():
251
  req = request.json
252
  if str(req["status"]) not in ["0", "1"]:
253
+ return get_json_result(
254
  data=False,
255
  retmsg='"Status" must be either 0 or 1!',
256
  retcode=RetCode.ARGUMENT_ERROR)
257
 
258
+ if not DocumentService.accessible(req["doc_id"], current_user.id):
259
+ return get_json_result(
260
+ data=False,
261
+ retmsg='No authorization.',
262
+ retcode=RetCode.AUTHENTICATION_ERROR)
263
+
264
  try:
265
  e, doc = DocumentService.get_by_id(req["doc_id"])
266
  if not e:
 
299
  req = request.json
300
  doc_ids = req["doc_id"]
301
  if isinstance(doc_ids, str): doc_ids = [doc_ids]
302
+
303
+ for doc_id in doc_ids:
304
+ if not DocumentService.accessible4deletion(doc_id, current_user.id):
305
+ return get_json_result(
306
+ data=False,
307
+ retmsg='No authorization.',
308
+ retcode=RetCode.AUTHENTICATION_ERROR
309
+ )
310
+
311
  root_folder = FileService.get_root_folder(current_user.id)
312
  pf_id = root_folder["id"]
313
  FileService.init_knowledgebase_docs(pf_id, current_user.id)
 
346
  @validate_request("doc_ids", "run")
347
  def run():
348
  req = request.json
349
+ for doc_id in req["doc_ids"]:
350
+ if not DocumentService.accessible(doc_id, current_user.id):
351
+ return get_json_result(
352
+ data=False,
353
+ retmsg='No authorization.',
354
+ retcode=RetCode.AUTHENTICATION_ERROR
355
+ )
356
  try:
357
  for id in req["doc_ids"]:
358
  info = {"run": str(req["run"]), "progress": 0}
 
386
  @validate_request("doc_id", "name")
387
  def rename():
388
  req = request.json
389
+ if not DocumentService.accessible(req["doc_id"], current_user.id):
390
+ return get_json_result(
391
+ data=False,
392
+ retmsg='No authorization.',
393
+ retcode=RetCode.AUTHENTICATION_ERROR
394
+ )
395
  try:
396
  e, doc = DocumentService.get_by_id(req["doc_id"])
397
  if not e:
 
452
  @validate_request("doc_id", "parser_id")
453
  def change_parser():
454
  req = request.json
455
+
456
+ if not DocumentService.accessible(req["doc_id"], current_user.id):
457
+ return get_json_result(
458
+ data=False,
459
+ retmsg='No authorization.',
460
+ retcode=RetCode.AUTHENTICATION_ERROR
461
+ )
462
  try:
463
  e, doc = DocumentService.get_by_id(req["doc_id"])
464
  if not e:
api/apps/kb_app.py CHANGED
@@ -13,7 +13,6 @@
13
  # See the License for the specific language governing permissions and
14
  # limitations under the License.
15
  #
16
- from elasticsearch_dsl import Q
17
  from flask import request
18
  from flask_login import login_required, current_user
19
 
@@ -23,14 +22,12 @@ from api.db.services.file2document_service import File2DocumentService
23
  from api.db.services.file_service import FileService
24
  from api.db.services.user_service import TenantService, UserTenantService
25
  from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
26
- from api.utils import get_uuid, get_format_time
27
- from api.db import StatusEnum, UserTenantRole, FileSource
28
  from api.db.services.knowledgebase_service import KnowledgebaseService
29
- from api.db.db_models import Knowledgebase, File
30
- from api.settings import stat_logger, RetCode
31
  from api.utils.api_utils import get_json_result
32
- from rag.nlp import search
33
- from rag.utils.es_conn import ELASTICSEARCH
34
 
35
 
36
  @manager.route('/create', methods=['post'])
@@ -65,6 +62,12 @@ def create():
65
  def update():
66
  req = request.json
67
  req["name"] = req["name"].strip()
 
 
 
 
 
 
68
  try:
69
  if not KnowledgebaseService.query(
70
  created_by=current_user.id, id=req["kb_id"]):
@@ -139,6 +142,12 @@ def list_kbs():
139
  @validate_request("kb_id")
140
  def rm():
141
  req = request.json
 
 
 
 
 
 
142
  try:
143
  kbs = KnowledgebaseService.query(
144
  created_by=current_user.id, id=req["kb_id"])
 
13
  # See the License for the specific language governing permissions and
14
  # limitations under the License.
15
  #
 
16
  from flask import request
17
  from flask_login import login_required, current_user
18
 
 
22
  from api.db.services.file_service import FileService
23
  from api.db.services.user_service import TenantService, UserTenantService
24
  from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
25
+ from api.utils import get_uuid
26
+ from api.db import StatusEnum, FileSource
27
  from api.db.services.knowledgebase_service import KnowledgebaseService
28
+ from api.db.db_models import File
29
+ from api.settings import RetCode
30
  from api.utils.api_utils import get_json_result
 
 
31
 
32
 
33
  @manager.route('/create', methods=['post'])
 
62
  def update():
63
  req = request.json
64
  req["name"] = req["name"].strip()
65
+ if not KnowledgebaseService.accessible4deletion(req["kb_id"], current_user.id):
66
+ return get_json_result(
67
+ data=False,
68
+ retmsg='No authorization.',
69
+ retcode=RetCode.AUTHENTICATION_ERROR
70
+ )
71
  try:
72
  if not KnowledgebaseService.query(
73
  created_by=current_user.id, id=req["kb_id"]):
 
142
  @validate_request("kb_id")
143
  def rm():
144
  req = request.json
145
+ if not KnowledgebaseService.accessible4deletion(req["kb_id"], current_user.id):
146
+ return get_json_result(
147
+ data=False,
148
+ retmsg='No authorization.',
149
+ retcode=RetCode.AUTHENTICATION_ERROR
150
+ )
151
  try:
152
  kbs = KnowledgebaseService.query(
153
  created_by=current_user.id, id=req["kb_id"])
api/db/services/document_service.py CHANGED
@@ -38,7 +38,7 @@ from rag.utils.storage_factory import STORAGE_IMPL
38
  from rag.nlp import search, rag_tokenizer
39
 
40
  from api.db import FileType, TaskStatus, ParserType, LLMType
41
- from api.db.db_models import DB, Knowledgebase, Tenant, Task
42
  from api.db.db_models import Document
43
  from api.db.services.common_service import CommonService
44
  from api.db.services.knowledgebase_service import KnowledgebaseService
@@ -263,6 +263,33 @@ class DocumentService(CommonService):
263
  return
264
  return docs[0]["tenant_id"]
265
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
266
  @classmethod
267
  @DB.connection_context()
268
  def get_embd_id(cls, doc_id):
 
38
  from rag.nlp import search, rag_tokenizer
39
 
40
  from api.db import FileType, TaskStatus, ParserType, LLMType
41
+ from api.db.db_models import DB, Knowledgebase, Tenant, Task, UserTenant
42
  from api.db.db_models import Document
43
  from api.db.services.common_service import CommonService
44
  from api.db.services.knowledgebase_service import KnowledgebaseService
 
263
  return
264
  return docs[0]["tenant_id"]
265
 
266
+ @classmethod
267
+ @DB.connection_context()
268
+ def accessible(cls, doc_id, user_id):
269
+ docs = cls.model.select(
270
+ cls.model.id).join(
271
+ Knowledgebase, on=(
272
+ Knowledgebase.id == cls.model.kb_id)
273
+ ).join(UserTenant, on=(UserTenant.tenant_id == Knowledgebase.tenant_id)
274
+ ).where(cls.model.id == doc_id, UserTenant.user_id == user_id).paginate(0, 1)
275
+ docs = docs.dicts()
276
+ if not docs:
277
+ return False
278
+ return True
279
+
280
+ @classmethod
281
+ @DB.connection_context()
282
+ def accessible4deletion(cls, doc_id, user_id):
283
+ docs = cls.model.select(
284
+ cls.model.id).join(
285
+ Knowledgebase, on=(
286
+ Knowledgebase.id == cls.model.kb_id)
287
+ ).where(cls.model.id == doc_id, Knowledgebase.created_by == user_id).paginate(0, 1)
288
+ docs = docs.dicts()
289
+ if not docs:
290
+ return False
291
+ return True
292
+
293
  @classmethod
294
  @DB.connection_context()
295
  def get_embd_id(cls, doc_id):
api/db/services/knowledgebase_service.py CHANGED
@@ -14,7 +14,7 @@
14
  # limitations under the License.
15
  #
16
  from api.db import StatusEnum, TenantPermission
17
- from api.db.db_models import Knowledgebase, DB, Tenant, User
18
  from api.db.services.common_service import CommonService
19
 
20
 
@@ -182,3 +182,25 @@ class KnowledgebaseService(CommonService):
182
  kbs = kbs.paginate(page_number, items_per_page)
183
 
184
  return list(kbs.dicts())
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  # limitations under the License.
15
  #
16
  from api.db import StatusEnum, TenantPermission
17
+ from api.db.db_models import Knowledgebase, DB, Tenant, User, UserTenant
18
  from api.db.services.common_service import CommonService
19
 
20
 
 
182
  kbs = kbs.paginate(page_number, items_per_page)
183
 
184
  return list(kbs.dicts())
185
+
186
+ @classmethod
187
+ @DB.connection_context()
188
+ def accessible(cls, kb_id, user_id):
189
+ docs = cls.model.select(
190
+ cls.model.id).join(UserTenant, on=(UserTenant.tenant_id == Knowledgebase.tenant_id)
191
+ ).where(cls.model.id == kb_id, UserTenant.user_id == user_id).paginate(0, 1)
192
+ docs = docs.dicts()
193
+ if not docs:
194
+ return False
195
+ return True
196
+
197
+ @classmethod
198
+ @DB.connection_context()
199
+ def accessible4deletion(cls, kb_id, user_id):
200
+ docs = cls.model.select(
201
+ cls.model.id).where(cls.model.id == kb_id, cls.model.created_by == user_id).paginate(0, 1)
202
+ docs = docs.dicts()
203
+ if not docs:
204
+ return False
205
+ return True
206
+
api/settings.py CHANGED
@@ -14,6 +14,7 @@
14
  # limitations under the License.
15
  #
16
  import os
 
17
  from enum import IntEnum, Enum
18
  from api.utils.file_utils import get_project_base_directory
19
  from api.utils.log_utils import LoggerFactory, getLogger
@@ -143,9 +144,8 @@ HTTP_PORT = get_base_config(RAG_FLOW_SERVICE_NAME, {}).get("http_port")
143
 
144
  SECRET_KEY = get_base_config(
145
  RAG_FLOW_SERVICE_NAME,
146
- {}).get(
147
- "secret_key",
148
- "infiniflow")
149
  TOKEN_EXPIRE_IN = get_base_config(
150
  RAG_FLOW_SERVICE_NAME, {}).get(
151
  "token_expires_in", 3600)
 
14
  # limitations under the License.
15
  #
16
  import os
17
+ from datetime import date
18
  from enum import IntEnum, Enum
19
  from api.utils.file_utils import get_project_base_directory
20
  from api.utils.log_utils import LoggerFactory, getLogger
 
144
 
145
  SECRET_KEY = get_base_config(
146
  RAG_FLOW_SERVICE_NAME,
147
+ {}).get("secret_key", str(date.today()))
148
+
 
149
  TOKEN_EXPIRE_IN = get_base_config(
150
  RAG_FLOW_SERVICE_NAME, {}).get(
151
  "token_expires_in", 3600)