release-v-1-15-0 (#228)

* Users into organization count is not correct(deleted users are counted) * null for empty org id * Searches in kratos if cache doesnt hold values * Fixes progress 0 log * Adds alembic to migrate the db * Fixes faulty gates similarity search display * adds soo indicator to mapped request * remove print * cancel task doesn’t work * Adds return for config chagnes * Adds session handler for starlett middleware * Hotfix session connection loss * submodules ref --------- Co-authored-by: JWittmeyer <[email protected]> Co-authored-by: anmarhindi <[email protected]>
code-kern-ai · Jun 13, 2024 · 27da64b · 27da64b
1 parent be59c33
commit 27da64b
Show file tree

Hide file tree

Showing 16 changed files with 219 additions and 49 deletions.
diff --git a/alembic/versions/706d5611a73e_ensured_user_cascade_behaviour.py b/alembic/versions/706d5611a73e_ensured_user_cascade_behaviour.py
@@ -0,0 +1,84 @@
+"""ensured user cascade behaviour
+
+Revision ID: 706d5611a73e
+Revises: 194838aa0431
+Create Date: 2024-06-12 09:29:07.617462
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '706d5611a73e'
+down_revision = '194838aa0431'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint('admin_message_archived_by_fkey', 'admin_message', type_='foreignkey')
+    op.drop_constraint('admin_message_created_by_fkey', 'admin_message', type_='foreignkey')
+    op.create_foreign_key(None, 'admin_message', 'user', ['archived_by'], ['id'], ondelete='SET NULL')
+    op.create_foreign_key(None, 'admin_message', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('agreement_user_id_fkey', 'agreement', type_='foreignkey')
+    op.create_foreign_key(None, 'agreement', 'user', ['user_id'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('comment_data_created_by_fkey', 'comment_data', type_='foreignkey')
+    op.create_foreign_key(None, 'comment_data', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('data_slice_created_by_fkey', 'data_slice', type_='foreignkey')
+    op.create_foreign_key(None, 'data_slice', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('embedding_created_by_fkey', 'embedding', type_='foreignkey')
+    op.create_foreign_key(None, 'embedding', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('information_source_created_by_fkey', 'information_source', type_='foreignkey')
+    op.create_foreign_key(None, 'information_source', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('labeling_access_link_created_by_fkey', 'labeling_access_link', type_='foreignkey')
+    op.create_foreign_key(None, 'labeling_access_link', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('labeling_task_label_created_by_fkey', 'labeling_task_label', type_='foreignkey')
+    op.create_foreign_key(None, 'labeling_task_label', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('project_created_by_fkey', 'project', type_='foreignkey')
+    op.create_foreign_key(None, 'project', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('record_label_association_created_by_fkey', 'record_label_association', type_='foreignkey')
+    op.create_foreign_key(None, 'record_label_association', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.create_index(op.f('ix_task_queue_created_by'), 'task_queue', ['created_by'], unique=False)
+    op.drop_constraint('task_queue_created_by_fkey', 'task_queue', type_='foreignkey')
+    op.create_foreign_key(None, 'task_queue', 'user', ['created_by'], ['id'], ondelete='SET NULL')
+    op.drop_constraint('upload_task_user_id_fkey', 'upload_task', type_='foreignkey')
+    op.create_foreign_key(None, 'upload_task', 'user', ['user_id'], ['id'], ondelete='SET NULL')
+    op.create_index(op.f('ix_cognition_macro_execution_organization_id'), 'macro_execution', ['organization_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_execution_link_organization_id'), 'macro_execution_link', ['organization_id'], unique=False, schema='cognition')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_cognition_macro_execution_link_organization_id'), table_name='macro_execution_link', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_execution_organization_id'), table_name='macro_execution', schema='cognition')
+    op.drop_constraint(None, 'upload_task', type_='foreignkey')
+    op.create_foreign_key('upload_task_user_id_fkey', 'upload_task', 'user', ['user_id'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'task_queue', type_='foreignkey')
+    op.create_foreign_key('task_queue_created_by_fkey', 'task_queue', 'user', ['created_by'], ['id'])
+    op.drop_index(op.f('ix_task_queue_created_by'), table_name='task_queue')
+    op.drop_constraint(None, 'record_label_association', type_='foreignkey')
+    op.create_foreign_key('record_label_association_created_by_fkey', 'record_label_association', 'user', ['created_by'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'project', type_='foreignkey')
+    op.create_foreign_key('project_created_by_fkey', 'project', 'user', ['created_by'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'labeling_task_label', type_='foreignkey')
+    op.create_foreign_key('labeling_task_label_created_by_fkey', 'labeling_task_label', 'user', ['created_by'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'labeling_access_link', type_='foreignkey')
+    op.create_foreign_key('labeling_access_link_created_by_fkey', 'labeling_access_link', 'user', ['created_by'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'information_source', type_='foreignkey')
+    op.create_foreign_key('information_source_created_by_fkey', 'information_source', 'user', ['created_by'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'embedding', type_='foreignkey')
+    op.create_foreign_key('embedding_created_by_fkey', 'embedding', 'user', ['created_by'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'data_slice', type_='foreignkey')
+    op.create_foreign_key('data_slice_created_by_fkey', 'data_slice', 'user', ['created_by'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'comment_data', type_='foreignkey')
+    op.create_foreign_key('comment_data_created_by_fkey', 'comment_data', 'user', ['created_by'], ['id'])
+    op.drop_constraint(None, 'agreement', type_='foreignkey')
+    op.create_foreign_key('agreement_user_id_fkey', 'agreement', 'user', ['user_id'], ['id'], ondelete='CASCADE')
+    op.drop_constraint(None, 'admin_message', type_='foreignkey')
+    op.drop_constraint(None, 'admin_message', type_='foreignkey')
+    op.create_foreign_key('admin_message_created_by_fkey', 'admin_message', 'user', ['created_by'], ['id'])
+    op.create_foreign_key('admin_message_archived_by_fkey', 'admin_message', 'user', ['archived_by'], ['id'])
+    # ### end Alembic commands ###
diff --git a/api/transfer.py b/api/transfer.py
@@ -2,7 +2,6 @@
 import traceback
 import time
 from typing import Optional, Dict
-from time import sleep
 from starlette.endpoints import HTTPEndpoint
 from starlette.responses import PlainTextResponse, JSONResponse
 from controller.embedding.manager import recreate_embeddings
@@ -299,22 +298,25 @@ def put(self, request) -> PlainTextResponse:
                 cognition_prj.organization_id, cognition_prj.created_by, True
             ).id
         )
+        cached = {str(e.id): str(e.created_by) for e in execution_entries}
 
         def queue_tasks():
             token = general.get_ctx_token()
-            for entry in execution_entries:
-                task_queue_manager.add_task(
-                    refinery_prj_id,
-                    TaskType.RUN_COGNITION_MACRO,
-                    entry.created_by,
-                    {
-                        "macro_id": macro_id,
-                        "execution_id": str(entry.id),
-                        "execution_group_id": group_id,
-                    },
-                )
-            general.commit()
-            general.remove_and_refresh_session(token, False)
+            try:
+                for exec_id in cached:
+                    task_queue_manager.add_task(
+                        refinery_prj_id,
+                        TaskType.RUN_COGNITION_MACRO,
+                        cached[exec_id],
+                        {
+                            "macro_id": macro_id,
+                            "execution_id": exec_id,
+                            "execution_group_id": group_id,
+                        },
+                    )
+                general.commit()
+            finally:
+                general.remove_and_refresh_session(token, False)
 
         daemon.run(queue_tasks)
 

diff --git a/app.py b/app.py
@@ -2,6 +2,7 @@
 
 from fastapi import FastAPI
 from api.healthcheck import Healthcheck
+from starlette.middleware import Middleware
 from api.misc import IsDemoRest, IsManagedRest
 from api.project import ProjectDetails, ProjectCreationFromWorkflow
 from api.transfer import (
@@ -36,6 +37,7 @@
 from fast_api.routes.weak_supervision import router as weak_supervision_router
 from fast_api.routes.labeling_tasks import router as labeling_tasks_router
 from middleware.database_session import handle_db_session
+from middleware.starlette_tmp_middleware import DatabaseSessionHandler
 from starlette.applications import Starlette
 from starlette.routing import Route, Mount
 
@@ -151,11 +153,8 @@
 
 fastapi_app.middleware("http")(handle_db_session)
 
-
-app = Starlette(routes=routes)
-
-# middleware = [Middleware(DatabaseSessionHandler)]
-# app = Starlette(routes=routes, middleware=middleware)
+middleware = [Middleware(DatabaseSessionHandler)]
+app = Starlette(routes=routes, middleware=middleware)
 
 init_task_queues()
 check_in_deletion_projects()

diff --git a/controller/auth/kratos.py b/controller/auth/kratos.py
@@ -3,6 +3,7 @@
 import requests
 import logging
 from datetime import datetime, timedelta
+from urllib.parse import quote
 
 logging.basicConfig(level=logging.INFO)
 logger: logging.Logger = logging.getLogger(__name__)
@@ -122,6 +123,19 @@ def get_userid_from_mail(user_mail: str) -> str:
             continue
         if values[key]["simple"]["mail"] == user_mail:
             return key
+    # not in cached values, try search kratos
+    return __search_kratos_for_user_mail(user_mail)
+
+
+def __search_kratos_for_user_mail(user_mail: str) -> str:
+    request = requests.get(
+        f"{KRATOS_ADMIN_URL}/identities?preview_credentials_identifier_similar={quote(user_mail)}"
+    )
+    if request.ok:
+        identities = request.json()
+        for i in identities:
+            if i["traits"]["email"].lower() == user_mail.lower():
+                return i["id"]
     return None
 
 

diff --git a/controller/auth/manager.py b/controller/auth/manager.py
@@ -17,7 +17,7 @@
 from controller.misc import manager as misc_manager
 import sqlalchemy
 
-DEV_USER_ID = "59e8dfca-ce56-44df-a8c7-5f05c61da499"
+DEV_USER_ID = "741df1c2-a531-43b6-b259-df23bc78e9a2"
 
 
 def get_organization_id_by_info(info) -> Organization:

diff --git a/controller/organization/manager.py b/controller/organization/manager.py
@@ -72,6 +72,11 @@ def get_all_users(
         all_users, column_whitelist=USER_INFO_WHITELIST
     )
     all_users_expanded = kratos.expand_user_mail_name(all_users_dict)
+    all_users_expanded = [
+        user
+        for user in all_users_expanded
+        if user["firstName"] is not None and user["lastName"] is not None
+    ]
     return all_users_expanded
 
 

diff --git a/controller/payload/manager.py b/controller/payload/manager.py
@@ -1,7 +1,6 @@
 from typing import Any, Dict, List, Optional
 from controller.payload import payload_scheduler
 from fast_api.types import (
-    LabelingFunctionSampleRecordWrapper,
     LabelingFunctionSampleRecords,
 )
 from submodules.model import InformationSourcePayload, enums
@@ -74,7 +73,7 @@ def get_labeling_function_on_10_records(
             }
             for record_item in sample_records
         ],
-        "containerLogs": container_logs,
+        "containerLogs": [log for log in container_logs if "progress: " not in log],
         "codeHasErrors": code_has_errors,
     }
 

diff --git a/controller/user/manager.py b/controller/user/manager.py
@@ -127,6 +127,13 @@ def get_mapped_sorted_paginated_users(
             get_user["role"] = active_user_by_id["role"]
             get_user["organization"] = active_user_by_id["organizationName"]
 
+            public_meta = get_user["metadata_public"]
+            get_user["sso_provider"] = (
+                public_meta.get("registration_scope", {}).get("provider_id", None)
+                if public_meta
+                else None
+            )
+
             final_users.append(get_user)
             save_len_final_users += 1
 

diff --git a/fast_api/models.py b/fast_api/models.py
@@ -419,3 +419,9 @@ class MappedSortedPaginatedUsers(BaseModel):
 
 class DeleteUserBody(BaseModel):
     user_id: StrictStr
+
+
+class CancelTaskBody(BaseModel):
+    project_id: StrictStr
+    task_id: StrictStr
+    task_type: StrictStr
diff --git a/fast_api/routes/client_response.py b/fast_api/routes/client_response.py
@@ -1,6 +1,6 @@
 import json
 from typing import Any, Optional
-from fastapi.responses import JSONResponse, Response
+from fastapi.responses import JSONResponse, PlainTextResponse, Response
 from fastapi import status
 from sqlalchemy.engine.row import Row
 from submodules.model.models import Base
@@ -50,6 +50,12 @@ def wrap_content_for_frontend(content: Any):
 )
 
 
+GENERIC_FAILURE_RESPONSE = PlainTextResponse(
+    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+    content="An error occurred",
+)
+
+
 def get_silent_success() -> JSONResponse:
     return SILENT_SUCCESS_RESPONSE
 

diff --git a/fast_api/routes/misc.py b/fast_api/routes/misc.py
@@ -2,6 +2,7 @@
 from fastapi import APIRouter, Body, Request
 from exceptions.exceptions import ProjectAccessError
 from fast_api.models import (
+    CancelTaskBody,
     ModelProviderDeleteModelBody,
     ModelProviderDownloadModelBody,
 )
@@ -108,9 +109,15 @@ def get_all_tasks(request: Request, only_running: bool):
 
 
 @router.post("/cancel-task")
-def cancel_task(request: Request, project_id: str, task_id: str, task_type: str):
+def cancel_task(
+    request: Request,
+    body: CancelTaskBody = Body(...),
+):
 
     auth.check_admin_access(request.state.info)
+    task_type = body.task_type
+    project_id = body.project_id
+    task_id = body.task_id
 
     if task_type == enums.TaskType.ATTRIBUTE_CALCULATION.value:
         controller_manager.cancel_attribute_calculation(project_id, task_id)

diff --git a/fast_api/routes/organization.py b/fast_api/routes/organization.py
@@ -76,7 +76,9 @@ def get_user_info_extended(request: Request):
     data = {
         "userInfo": {
             "id": str(user.id),
-            "organizationId": str(user.organization_id),
+            "organizationId": (
+                str(user.organization_id) if user.organization_id else None
+            ),
             "firstName": name.get("first"),
             "lastName": name.get("last"),
             "mail": mail,
@@ -216,11 +218,13 @@ def update_config(request: Request, body: UpdateConfigBody = Body(...)):
         print(
             "config should only be changed for open source/local version to prevent limit issues"
         )
+        return
     misc.update_config(body.dict_str)
     misc.refresh_config()
     orgs = organization.get_all()
     if not orgs or len(orgs) != 1:
         print("local version should only have one organization")
+        return
 
     for org in orgs:
         # send to all so all are notified about the change
@@ -275,6 +279,7 @@ def get_all_organizations(request: Request):
                                 }
                             }
                             for user in org.users
+                            if resolve_user_mail_by_id(user.id) is not None
                         ]
                     },
                     "maxRows": org.max_rows,

diff --git a/fast_api/routes/project.py b/fast_api/routes/project.py
@@ -305,11 +305,18 @@ def labeling_tasks_by_project_id(project_id: str) -> str:
     "/{project_id}/labeling-tasks-by-project-id-with-embeddings",
     dependencies=[Depends(auth_manager.check_project_access_dep)],
 )
-def labeling_tasks_by_project_id_with_embeddings(project_id: str) -> str:
+def labeling_tasks_by_project_id_with_embeddings(
+    project_id: str, only_on_attribute: bool = False
+) -> str:
     embeddings = get_all_embeddings_by_project_id(project_id)
 
     embeddings_edges = []
     for embedding in embeddings:
+        if (
+            only_on_attribute
+            and embedding.type != enums.EmbeddingType.ON_ATTRIBUTE.value
+        ):
+            continue
         attribute = attr_manager.get_attribute(project_id, embedding.attribute_id)
         embeddings_edges.append(
             {