Cognition integration provider (#167)

* build: third party integration first commit

* chore: update enums

* perf: add integration acess

* perf: rename to integration

* perf: add last_extraction column to integration

* perf: update integration delta

* perf: update integration access to list types

* perf: add integration_types to integration access

* perf: add integration_types to integration access

* perf: rename last_extraction to extract_history

* fix: store enum.value instead of enum

* fix: integration.project_id nullable

* fix: nulable column instead of foreignkey

* fix: enum values

* perf: task cancellation

* fix: keyword arguments

* perf: integration record

* perf: add tokenizer

* perf: add update integration access

* perf: update integration endpoints

* perf: add get endpoint

* Oidc field in the users table

* perf: add org_id to integration provider

* perf: add org_id support to integration

* perf: add record delta criteria

* fix: task execution finish on failed integration

* perf: add integration finished_at

* perf: add started_at

* fix: started_at - finished_at syntax error

* perf: add integration records

* perf: add integration tables

* perf: update integrations delta

* perf: add sharepoint integration

* perf: update integration objects

* perf: expand IntegrationSharepoint

* fix: integration.started_at
perf: add unique constraints to integrations

* perf: integration data types

* perf: unique constraint names

* Reset finished at for new integrations

* perf: update integration objects

* perf: add integration delta deletion

* basic models

* perf: last_synced_at integration column

* perf: add is_synced column

* chore: add typing

* perf: add sync columns

* fix model

* perf: add get_all integrations

* chore: add todo comment

* permission test

* projects with access management

* perf: add sharepoint db bo

* deactivate mock up

* get by user id

* enable list payload

* perf: integration update

* perf: tech discussion feedback

* perf: get integrations updates

* perf: integration updates

* perf: introduce managers

* chore: typing

* perf: access + check for updates

* perf: update integration

* perf: add delta url to sharepoint integration

* fix: move delta_url to cognitionintegration

* perf: integration updates

* perf: add updated_by + delta_criteria

* perf: add delta_criteria field

* perf: check for status improvement

* add meta_data group

* format

* model

* sync action

* rename action

* perf: dynamic 'by' record grouping

* style: arguments newlines in function definitions

* new getters

* delete groups

* fix: rm get project by name and org id

* perf: pr review comments

* perf: pr review comments

* fix: paginated query

* perf: update integrations model

* perf: update to_snake_case regex compilation

* access management

* perf: db model update

* perf: pr comments

* perf: remove unnecessary checks

* fix: nullable error message

* perf: pr comments

* delete

* perf: move IntegrationMetadata enum to integration_objects.helper

* Oidc field in the users table (#170)

Co-authored-by: andhreljaKern <andrea.hrelja@kern.ai>

* fix: metadata helper function

* perf: update integration task type name

* style: formatting

* perf: pr review comments

* perf: add REFINERY_ATTRIBUTE_ACCESS constants

* fix: query builder for record_ids

* perf: monitor.set_integration_task_to_failed with state

* perf: add early exit for deleted integrations

* perf: add early exit for task execution

* sharepoint active queue

* unique by name and integration

* typing

* improve sql

* merge

* model

* fix group

* remove unique

* perf: add file_properties integration column

* perf: update default state for set_integration_task_to_failed

* Adds option filter for pid

* PR comments

---------

Co-authored-by: Lina <lina.lumburovska@kern.ai>
Co-authored-by: LennartSchmidtKern <lennart.schmidt@kern.ai>
Co-authored-by: lumburovskalina <104008550+lumburovskalina@users.noreply.github.com>
Co-authored-by: JWittmeyer <jens.wittmeyer@kern.ai>

Author: 4 people

business_objects/embedding.py

@@ -321,8 +321,11 @@ def __build_payload_selector(
             if (
                 data_type != enums.DataTypes.TEXT.value
                 and data_type != enums.DataTypes.LLM_RESPONSE.value
+                and data_type != enums.DataTypes.PERMISSION.value
             ):
                 payload_selector += f"'{attr}', (r.\"data\"->>'{attr}')::{data_type}"
+            if data_type == enums.DataTypes.PERMISSION.value:
+                payload_selector += f"'{attr}', r.\"data\"->'{attr}'"
             else:
                 payload_selector += f"'{attr}', r.\"data\"->>'{attr}'"
         payload_selector = f"json_build_object({payload_selector}) payload"
@@ -391,7 +394,8 @@ def get_tensors_and_attributes_for_qdrant(
     WHERE et.project_id = '{project_id}' AND et.embedding_id = '{embedding_id}'
     """
     if record_ids:
-        query += f" AND r.id IN ('{','.join(record_ids)}')"
+        _record_ids = "','".join(record_ids)
+        query += f" AND r.id IN ('{_record_ids}')"
 
     return general.execute_all(query)
 

business_objects/monitor.py

@@ -1,4 +1,5 @@
 from typing import Any, List, Optional
+import datetime
 from . import general
 from .. import enums
 from ..models import TaskQueue, Organization
@@ -9,6 +10,7 @@
     markdown_file as markdown_file_db_bo,
     file_extraction as file_extraction_db_bo,
     file_transformation as file_transformation_db_bo,
+    integration as integration_db_bo,
 )
 
 FILE_CACHING_IN_PROGRESS_STATES = [
@@ -197,6 +199,27 @@ def set_parse_cognition_file_task_to_failed(
     general.commit()
 
 
+def set_integration_task_to_failed(
+    integration_id: str,
+    is_synced: bool = False,
+    error_message: Optional[str] = None,
+    state: Optional[
+        enums.CognitionMarkdownFileState
+    ] = enums.CognitionMarkdownFileState.FAILED,
+    with_commit: bool = True,
+) -> None:
+    # argument `state` is a workaround for cognition-gateway/api/routes/integrations.delete_many
+    integration_db_bo.update(
+        id=integration_id,
+        state=state,
+        finished_at=datetime.datetime.now(datetime.timezone.utc),
+        is_synced=is_synced,
+        error_message=error_message,
+        last_synced_at=datetime.datetime.now(datetime.timezone.utc),
+        with_commit=with_commit,
+    )
+
+
 def __select_running_information_source_payloads(
     project_id: Optional[str] = None,
     only_running: bool = False,

business_objects/project.py

@@ -1,16 +1,14 @@
 from typing import List, Optional, Any, Dict, Union, Set
-from sqlalchemy.sql import func
-from sqlalchemy import cast, Integer
+from sqlalchemy.sql import func, cast
 from sqlalchemy.sql.functions import coalesce
-
-
+from sqlalchemy import Integer
 from . import general, attribute
-
 from .. import enums
 from ..session import session
-from ..models import (
-    Project,
-    Record,
+from ..models import Project, Record, Attribute
+from ..integration_objects.helper import (
+    REFINERY_ATTRIBUTE_ACCESS_GROUPS,
+    REFINERY_ATTRIBUTE_ACCESS_USERS,
 )
 from ..util import prevent_sql_injection
 
@@ -118,12 +116,18 @@ def __build_sql_data_slices_by_project(project_id: str) -> str:
         project.id = '{project_id}'::UUID; """
 
 
-def get_dropdown_list_project_list(org_id: str) -> List[Dict[str, str]]:
+def get_dropdown_list_project_list(
+    org_id: str, project_id: Optional[str] = None
+) -> List[Dict[str, str]]:
     org_id = prevent_sql_injection(org_id, isinstance(org_id, str))
+    prj_filter = ""
+    if project_id:
+        project_id = prevent_sql_injection(project_id, isinstance(project_id, str))
+        prj_filter = f"AND p.id = '{project_id}'"
     query = f"""
     SELECT array_agg(jsonb_build_object('value', p.id,'name',p.NAME))
     FROM public.project p
-    WHERE p.organization_id = '{org_id}' AND p.status != '{enums.ProjectStatus.HIDDEN.value}'
+    WHERE p.organization_id = '{org_id}' AND p.status != '{enums.ProjectStatus.HIDDEN.value}' {prj_filter}
     """
     values = general.execute_first(query)
 
@@ -155,6 +159,56 @@ def get_all(organization_id: str) -> List[Project]:
     )
 
 
+def get_all_with_access_management(org_id: str) -> List[Dict[str, Any]]:
+    org_id_safe = prevent_sql_injection(org_id, isinstance(org_id, str))
+
+    hidden_status = enums.ProjectStatus.HIDDEN.value
+    permission_data_type = enums.DataTypes.PERMISSION.value
+    automatically_created_state = enums.AttributeState.AUTOMATICALLY_CREATED.value
+    access_groups_attr = REFINERY_ATTRIBUTE_ACCESS_GROUPS
+    access_users_attr = REFINERY_ATTRIBUTE_ACCESS_USERS
+
+    query = f"""
+    SELECT DISTINCT
+            p.*,
+            COALESCE((ci.config -> 'extract_kwargs' ->> 'sync_sharepoint_permissions')::BOOLEAN,FALSE) AS is_sharepoint_sync_active
+        FROM
+            public.project p
+        JOIN
+            public.attribute a ON p.id = a.project_id
+        LEFT JOIN
+            cognition.integration ci ON p.id = ci.project_id
+        WHERE
+            p.organization_id = '{org_id_safe}'
+            AND p.status != '{hidden_status}'
+            AND a.name IN ('{access_groups_attr}', '{access_users_attr}')
+            AND a.user_created = FALSE
+            AND a.data_type = '{permission_data_type}'
+            AND a.state = '{automatically_created_state}';
+    """
+
+    values = general.execute_all(query)
+    return values
+
+
+def check_access_management_active(project_id: str) -> bool:
+    return (
+        session.query(Project)
+        .join(Attribute, Project.id == Attribute.project_id)
+        .filter(
+            Project.id == project_id,
+            Attribute.name.in_(
+                [REFINERY_ATTRIBUTE_ACCESS_GROUPS, REFINERY_ATTRIBUTE_ACCESS_USERS]
+            ),
+            Attribute.user_created == False,
+            Attribute.data_type == enums.DataTypes.PERMISSION.value,
+            Attribute.state == enums.AttributeState.AUTOMATICALLY_CREATED.value,
+        )
+        .count()
+        > 0
+    )
+
+
 def get_all_by_user_organization_id(organization_id: str) -> List[Project]:
     projects = (
         session.query(Project).filter(Project.organization_id == organization_id).all()

business_objects/record.py

@@ -1,6 +1,6 @@
 from __future__ import with_statement
 from typing import List, Dict, Any, Optional, Tuple, Iterable
-from sqlalchemy import cast, Text
+from sqlalchemy import cast, Text, String
 from sqlalchemy.orm.attributes import flag_modified
 from sqlalchemy.sql.expression import bindparam
 from sqlalchemy import update
@@ -15,6 +15,10 @@
     Attribute,
     RecordTokenized,
 )
+from ..integration_objects.helper import (
+    REFINERY_ATTRIBUTE_ACCESS_GROUPS,
+    REFINERY_ATTRIBUTE_ACCESS_USERS,
+)
 from ..session import session
 from ..util import prevent_sql_injection
 
@@ -609,7 +613,7 @@ def count_missing_tokenized_records(project_id: str) -> int:
     query = f"""
     SELECT COUNT(*)
     FROM (
-        {get_records_without_tokenization(project_id, None, query_only = True)}
+        {get_records_without_tokenization(project_id, None, query_only=True)}
     ) record_query
     """
     return general.execute_first(query)[0]
@@ -807,6 +811,29 @@ def delete_user_created_attribute(
     general.flush_or_commit(with_commit)
 
 
+def delete_access_management_attributes(
+    project_id: str, with_commit: bool = True
+) -> None:
+    access_groups_attribute_item = attribute.get_by_name(
+        project_id, REFINERY_ATTRIBUTE_ACCESS_GROUPS
+    )
+    access_users_attribute_item = attribute.get_by_name(
+        project_id, REFINERY_ATTRIBUTE_ACCESS_USERS
+    )
+
+    if access_users_attribute_item and access_groups_attribute_item:
+        record_items = get_all(project_id=project_id)
+        for i, record_item in enumerate(record_items):
+            if record_item.data.get(access_groups_attribute_item.name):
+                del record_item.data[access_groups_attribute_item.name]
+            if record_item.data.get(access_users_attribute_item.name):
+                del record_item.data[access_users_attribute_item.name]
+            flag_modified(record_item, "data")
+            if (i + 1) % 1000 == 0:
+                general.flush_or_commit(with_commit)
+        general.flush_or_commit(with_commit)
+
+
 def delete_duplicated_rats(with_commit: bool = False) -> None:
     # no project so run for all to prevent expensive join with record table
     query = """
@@ -925,3 +952,19 @@ def get_first_no_text_column(project_id: str, record_id: str) -> str:
     WHERE r.project_id = '{project_id}' AND r.id = '{record_id}'
     """
     return general.execute_first(query)[0]
+
+
+def get_record_ids_by_running_ids(project_id: str, running_ids: List[int]) -> List[str]:
+    return [
+        row[0]
+        for row in (
+            session.query(cast(Record.id, String))
+            .filter(
+                Record.project_id == project_id,
+                Record.data[attribute.get_running_id_name(project_id)]
+                .as_integer()
+                .in_(running_ids),
+            )
+            .all()
+        )
+    ]

cognition_objects/group.py

@@ -0,0 +1,124 @@
+from datetime import datetime
+from typing import List, Optional
+from ..business_objects import general
+from ..session import session
+from ..models import CognitionGroup
+
+
+def get(group_id: str) -> CognitionGroup:
+    return session.query(CognitionGroup).filter(CognitionGroup.id == group_id).first()
+
+
+def get_with_organization_id(organization_id: str, group_id: str) -> CognitionGroup:
+    return (
+        session.query(CognitionGroup)
+        .filter(
+            CognitionGroup.organization_id == organization_id,
+            CognitionGroup.id == group_id,
+        )
+        .first()
+    )
+
+
+def get_all(organization_id: str) -> List[CognitionGroup]:
+    return (
+        session.query(CognitionGroup)
+        .filter(CognitionGroup.organization_id == organization_id)
+        .order_by(CognitionGroup.name.asc())
+        .all()
+    )
+
+
+def get_all_by_integration_id(
+    organization_id: str, integration_id: str
+) -> List[CognitionGroup]:
+    integration_id_json = CognitionGroup.meta_data.op("->>")("integration_id")
+
+    return (
+        session.query(CognitionGroup)
+        .filter(
+            CognitionGroup.organization_id == organization_id,
+            integration_id_json == integration_id,
+        )
+        .order_by(CognitionGroup.name.asc())
+        .all()
+    )
+
+
+def get_all_by_integration_id_permission_grouped(
+    organization_id: str, integration_id: str
+) -> List[CognitionGroup]:
+    integration_id_json = CognitionGroup.meta_data.op("->>")("integration_id")
+
+    integration_groups = (
+        session.query(CognitionGroup)
+        .filter(
+            CognitionGroup.organization_id == organization_id,
+            integration_id_json == integration_id,
+        )
+        .all()
+    )
+
+    return {group.meta_data.get("permission_id"): group for group in integration_groups}
+
+
+def get_by_name_and_integration(
+    organization_id: str, integration_id: str, name: str
+) -> CognitionGroup:
+    integration_id_json = CognitionGroup.meta_data.op("->>")("integration_id")
+    return (
+        session.query(CognitionGroup)
+        .filter(
+            CognitionGroup.organization_id == organization_id,
+            CognitionGroup.name == name,
+            integration_id_json == integration_id,
+        )
+        .first()
+    )
+
+
+def create_group(
+    organization_id: str,
+    name: str,
+    description: str,
+    created_by: str,
+    created_at: Optional[datetime] = None,
+    with_commit: bool = True,
+    meta_data: Optional[dict] = None,
+) -> CognitionGroup:
+    group = CognitionGroup(
+        organization_id=organization_id,
+        name=name,
+        description=description,
+        created_by=created_by,
+        created_at=created_at,
+        meta_data=meta_data,
+    )
+    general.add(group, with_commit)
+    return group
+
+
+def update_group(
+    group_id: str,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    with_commit: bool = True,
+    meta_data: Optional[dict] = None,
+) -> CognitionGroup:
+    group = get(group_id)
+
+    if name is not None:
+        group.name = name
+    if description is not None:
+        group.description = description
+    if meta_data is not None:
+        group.meta_data = meta_data
+    general.flush_or_commit(with_commit)
+
+    return group
+
+
+def delete(organization_id: str, group_id: str, with_commit: bool = True) -> None:
+    group = get_with_organization_id(organization_id, group_id)
+    if group:
+        general.delete(group, with_commit)