welcome to onyx

backend/onyx/redis/redis_connector.py

@@ -0,0 +1,119 @@
+import time
+
+import redis
+
+from onyx.db.models import SearchSettings
+from onyx.redis.redis_connector_delete import RedisConnectorDelete
+from onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync
+from onyx.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync
+from onyx.redis.redis_connector_index import RedisConnectorIndex
+from onyx.redis.redis_connector_prune import RedisConnectorPrune
+from onyx.redis.redis_connector_stop import RedisConnectorStop
+from onyx.redis.redis_pool import get_redis_client
+
+
+class RedisConnector:
+    """Composes several classes to simplify interacting with a connector and its
+    associated background tasks / associated redis interactions."""
+
+    def __init__(self, tenant_id: str | None, id: int) -> None:
+        self.tenant_id: str | None = tenant_id
+        self.id: int = id
+        self.redis: redis.Redis = get_redis_client(tenant_id=tenant_id)
+
+        self.stop = RedisConnectorStop(tenant_id, id, self.redis)
+        self.prune = RedisConnectorPrune(tenant_id, id, self.redis)
+        self.delete = RedisConnectorDelete(tenant_id, id, self.redis)
+        self.permissions = RedisConnectorPermissionSync(tenant_id, id, self.redis)
+        self.external_group_sync = RedisConnectorExternalGroupSync(
+            tenant_id, id, self.redis
+        )
+
+    def new_index(self, search_settings_id: int) -> RedisConnectorIndex:
+        return RedisConnectorIndex(
+            self.tenant_id, self.id, search_settings_id, self.redis
+        )
+
+    def wait_for_indexing_termination(
+        self,
+        search_settings_list: list[SearchSettings],
+        timeout: float = 15.0,
+    ) -> bool:
+        """
+        Returns True if all indexing for the given redis connector is finished within the given timeout.
+        Returns False if the timeout is exceeded
+
+        This check does not guarantee that current indexings being terminated
+        won't get restarted midflight
+        """
+
+        finished = False
+
+        start = time.monotonic()
+
+        while True:
+            still_indexing = False
+            for search_settings in search_settings_list:
+                redis_connector_index = self.new_index(search_settings.id)
+                if redis_connector_index.fenced:
+                    still_indexing = True
+                    break
+
+            if not still_indexing:
+                finished = True
+                break
+
+            now = time.monotonic()
+            if now - start > timeout:
+                break
+
+            time.sleep(1)
+            continue
+
+        return finished
+
+    @staticmethod
+    def get_id_from_fence_key(key: str) -> str | None:
+        """
+        Extracts the object ID from a fence key in the format `PREFIX_fence_X`.
+
+        Args:
+            key (str): The fence key string.
+
+        Returns:
+            Optional[int]: The extracted ID if the key is in the correct format, otherwise None.
+        """
+        parts = key.split("_")
+        if len(parts) != 3:
+            return None
+
+        object_id = parts[2]
+        return object_id
+
+    @staticmethod
+    def get_id_from_task_id(task_id: str) -> str | None:
+        """
+        Extracts the object ID from a task ID string.
+
+        This method assumes the task ID is formatted as `prefix_objectid_suffix`, where:
+        - `prefix` is an arbitrary string (e.g., the name of the task or entity),
+        - `objectid` is the ID you want to extract,
+        - `suffix` is another arbitrary string (e.g., a UUID).
+
+        Example:
+            If the input `task_id` is `documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc`,
+            this method will return the string `"1"`.
+
+        Args:
+            task_id (str): The task ID string from which to extract the object ID.
+
+        Returns:
+            str | None: The extracted object ID if the task ID is in the correct format, otherwise None.
+        """
+        # example: task_id=documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc
+        parts = task_id.split("_")
+        if len(parts) != 3:
+            return None
+
+        object_id = parts[1]
+        return object_id

backend/onyx/redis/redis_connector_credential_pair.py

@@ -0,0 +1,119 @@
+import time
+from typing import cast
+from uuid import uuid4
+
+from celery import Celery
+from redis import Redis
+from redis.lock import Lock as RedisLock
+from sqlalchemy.orm import Session
+
+from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
+from onyx.configs.constants import OnyxCeleryPriority
+from onyx.configs.constants import OnyxCeleryQueues
+from onyx.configs.constants import OnyxCeleryTask
+from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
+from onyx.db.document import (
+    construct_document_select_for_connector_credential_pair_by_needs_sync,
+)
+from onyx.db.models import Document
+from onyx.redis.redis_object_helper import RedisObjectHelper
+
+
+class RedisConnectorCredentialPair(RedisObjectHelper):
+    """This class is used to scan documents by cc_pair in the db and collect them into
+    a unified set for syncing.
+
+    It differs from the other redis helpers in that the taskset used spans
+    all connectors and is not per connector."""
+
+    PREFIX = "connectorsync"
+    FENCE_PREFIX = PREFIX + "_fence"
+    TASKSET_PREFIX = PREFIX + "_taskset"
+
+    def __init__(self, tenant_id: str | None, id: int) -> None:
+        super().__init__(tenant_id, str(id))
+
+        # documents that should be skipped
+        self.skip_docs: set[str] = set()
+
+    @classmethod
+    def get_fence_key(cls) -> str:
+        return RedisConnectorCredentialPair.FENCE_PREFIX
+
+    @classmethod
+    def get_taskset_key(cls) -> str:
+        return RedisConnectorCredentialPair.TASKSET_PREFIX
+
+    @property
+    def taskset_key(self) -> str:
+        """Notice that this is intentionally reusing the same taskset for all
+        connector syncs"""
+        # example: connector_taskset
+        return f"{self.TASKSET_PREFIX}"
+
+    def set_skip_docs(self, skip_docs: set[str]) -> None:
+        # documents that should be skipped. Note that this classes updates
+        # the list on the fly
+        self.skip_docs = skip_docs
+
+    def generate_tasks(
+        self,
+        celery_app: Celery,
+        db_session: Session,
+        redis_client: Redis,
+        lock: RedisLock,
+        tenant_id: str | None,
+    ) -> tuple[int, int] | None:
+        last_lock_time = time.monotonic()
+
+        async_results = []
+        cc_pair = get_connector_credential_pair_from_id(int(self._id), db_session)
+        if not cc_pair:
+            return None
+
+        stmt = construct_document_select_for_connector_credential_pair_by_needs_sync(
+            cc_pair.connector_id, cc_pair.credential_id
+        )
+
+        num_docs = 0
+
+        for doc in db_session.scalars(stmt).yield_per(1):
+            doc = cast(Document, doc)
+            current_time = time.monotonic()
+            if current_time - last_lock_time >= (
+                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
+            ):
+                lock.reacquire()
+                last_lock_time = current_time
+
+            num_docs += 1
+
+            # check if we should skip the document (typically because it's already syncing)
+            if doc.id in self.skip_docs:
+                continue
+
+            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            # the key for the result is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            # we prefix the task id so it's easier to keep track of who created the task
+            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            custom_task_id = f"{self.task_id_prefix}_{uuid4()}"
+
+            # add to the tracking taskset in redis BEFORE creating the celery task.
+            # note that for the moment we are using a single taskset key, not differentiated by cc_pair id
+            redis_client.sadd(
+                RedisConnectorCredentialPair.get_taskset_key(), custom_task_id
+            )
+
+            # Priority on sync's triggered by new indexing should be medium
+            result = celery_app.send_task(
+                OnyxCeleryTask.VESPA_METADATA_SYNC_TASK,
+                kwargs=dict(document_id=doc.id, tenant_id=tenant_id),
+                queue=OnyxCeleryQueues.VESPA_METADATA_SYNC,
+                task_id=custom_task_id,
+                priority=OnyxCeleryPriority.MEDIUM,
+            )
+
+            async_results.append(result)
+            self.skip_docs.add(doc.id)
+
+        return len(async_results), num_docs

backend/onyx/redis/redis_connector_delete.py

@@ -0,0 +1,151 @@
+import time
+from datetime import datetime
+from typing import cast
+from uuid import uuid4
+
+import redis
+from celery import Celery
+from pydantic import BaseModel
+from redis.lock import Lock as RedisLock
+from sqlalchemy.orm import Session
+
+from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
+from onyx.configs.constants import OnyxCeleryPriority
+from onyx.configs.constants import OnyxCeleryQueues
+from onyx.configs.constants import OnyxCeleryTask
+from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
+from onyx.db.document import construct_document_select_for_connector_credential_pair
+from onyx.db.models import Document as DbDocument
+
+
+class RedisConnectorDeletePayload(BaseModel):
+    num_tasks: int | None
+    submitted: datetime
+
+
+class RedisConnectorDelete:
+    """Manages interactions with redis for deletion tasks. Should only be accessed
+    through RedisConnector."""
+
+    PREFIX = "connectordeletion"
+    FENCE_PREFIX = f"{PREFIX}_fence"  # "connectordeletion_fence"
+    TASKSET_PREFIX = f"{PREFIX}_taskset"  # "connectordeletion_taskset"
+
+    def __init__(self, tenant_id: str | None, id: int, redis: redis.Redis) -> None:
+        self.tenant_id: str | None = tenant_id
+        self.id = id
+        self.redis = redis
+
+        self.fence_key: str = f"{self.FENCE_PREFIX}_{id}"
+        self.taskset_key = f"{self.TASKSET_PREFIX}_{id}"
+
+    def taskset_clear(self) -> None:
+        self.redis.delete(self.taskset_key)
+
+    def get_remaining(self) -> int:
+        # todo: move into fence
+        remaining = cast(int, self.redis.scard(self.taskset_key))
+        return remaining
+
+    @property
+    def fenced(self) -> bool:
+        if self.redis.exists(self.fence_key):
+            return True
+
+        return False
+
+    @property
+    def payload(self) -> RedisConnectorDeletePayload | None:
+        # read related data and evaluate/print task progress
+        fence_bytes = cast(bytes, self.redis.get(self.fence_key))
+        if fence_bytes is None:
+            return None
+
+        fence_str = fence_bytes.decode("utf-8")
+        payload = RedisConnectorDeletePayload.model_validate_json(cast(str, fence_str))
+
+        return payload
+
+    def set_fence(self, payload: RedisConnectorDeletePayload | None) -> None:
+        if not payload:
+            self.redis.delete(self.fence_key)
+            return
+
+        self.redis.set(self.fence_key, payload.model_dump_json())
+
+    def _generate_task_id(self) -> str:
+        # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
+        # we prefix the task id so it's easier to keep track of who created the task
+        # aka "connectordeletion_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
+
+        return f"{self.PREFIX}_{self.id}_{uuid4()}"
+
+    def generate_tasks(
+        self,
+        celery_app: Celery,
+        db_session: Session,
+        lock: RedisLock,
+    ) -> int | None:
+        """Returns None if the cc_pair doesn't exist.
+        Otherwise, returns an int with the number of generated tasks."""
+        last_lock_time = time.monotonic()
+
+        async_results = []
+        cc_pair = get_connector_credential_pair_from_id(int(self.id), db_session)
+        if not cc_pair:
+            return None
+
+        stmt = construct_document_select_for_connector_credential_pair(
+            cc_pair.connector_id, cc_pair.credential_id
+        )
+        for doc_temp in db_session.scalars(stmt).yield_per(1):
+            doc: DbDocument = doc_temp
+            current_time = time.monotonic()
+            if current_time - last_lock_time >= (
+                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
+            ):
+                lock.reacquire()
+                last_lock_time = current_time
+
+            custom_task_id = self._generate_task_id()
+
+            # add to the tracking taskset in redis BEFORE creating the celery task.
+            # note that for the moment we are using a single taskset key, not differentiated by cc_pair id
+            self.redis.sadd(self.taskset_key, custom_task_id)
+
+            # Priority on sync's triggered by new indexing should be medium
+            result = celery_app.send_task(
+                OnyxCeleryTask.DOCUMENT_BY_CC_PAIR_CLEANUP_TASK,
+                kwargs=dict(
+                    document_id=doc.id,
+                    connector_id=cc_pair.connector_id,
+                    credential_id=cc_pair.credential_id,
+                    tenant_id=self.tenant_id,
+                ),
+                queue=OnyxCeleryQueues.CONNECTOR_DELETION,
+                task_id=custom_task_id,
+                priority=OnyxCeleryPriority.MEDIUM,
+            )
+
+            async_results.append(result)
+
+        return len(async_results)
+
+    def reset(self) -> None:
+        self.redis.delete(self.taskset_key)
+        self.redis.delete(self.fence_key)
+
+    @staticmethod
+    def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:
+        taskset_key = f"{RedisConnectorDelete.TASKSET_PREFIX}_{id}"
+        r.srem(taskset_key, task_id)
+        return
+
+    @staticmethod
+    def reset_all(r: redis.Redis) -> None:
+        """Deletes all redis values for all connectors"""
+        for key in r.scan_iter(RedisConnectorDelete.TASKSET_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorDelete.FENCE_PREFIX + "*"):
+            r.delete(key)

backend/onyx/redis/redis_connector_doc_perm_sync.py

@@ -0,0 +1,200 @@
+import time
+from datetime import datetime
+from typing import cast
+from uuid import uuid4
+
+import redis
+from celery import Celery
+from pydantic import BaseModel
+from redis.lock import Lock as RedisLock
+
+from onyx.access.models import DocExternalAccess
+from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
+from onyx.configs.constants import OnyxCeleryPriority
+from onyx.configs.constants import OnyxCeleryQueues
+from onyx.configs.constants import OnyxCeleryTask
+
+
+class RedisConnectorPermissionSyncPayload(BaseModel):
+    started: datetime | None
+    celery_task_id: str | None
+
+
+class RedisConnectorPermissionSync:
+    """Manages interactions with redis for doc permission sync tasks. Should only be accessed
+    through RedisConnector."""
+
+    PREFIX = "connectordocpermissionsync"
+
+    FENCE_PREFIX = f"{PREFIX}_fence"
+
+    # phase 1 - geneartor task and progress signals
+    GENERATORTASK_PREFIX = f"{PREFIX}+generator"  # connectorpermissions+generator
+    GENERATOR_PROGRESS_PREFIX = (
+        PREFIX + "_generator_progress"
+    )  # connectorpermissions_generator_progress
+    GENERATOR_COMPLETE_PREFIX = (
+        PREFIX + "_generator_complete"
+    )  # connectorpermissions_generator_complete
+
+    TASKSET_PREFIX = f"{PREFIX}_taskset"  # connectorpermissions_taskset
+    SUBTASK_PREFIX = f"{PREFIX}+sub"  # connectorpermissions+sub
+
+    def __init__(self, tenant_id: str | None, id: int, redis: redis.Redis) -> None:
+        self.tenant_id: str | None = tenant_id
+        self.id = id
+        self.redis = redis
+
+        self.fence_key: str = f"{self.FENCE_PREFIX}_{id}"
+        self.generator_task_key = f"{self.GENERATORTASK_PREFIX}_{id}"
+        self.generator_progress_key = f"{self.GENERATOR_PROGRESS_PREFIX}_{id}"
+        self.generator_complete_key = f"{self.GENERATOR_COMPLETE_PREFIX}_{id}"
+
+        self.taskset_key = f"{self.TASKSET_PREFIX}_{id}"
+
+        self.subtask_prefix: str = f"{self.SUBTASK_PREFIX}_{id}"
+
+    def taskset_clear(self) -> None:
+        self.redis.delete(self.taskset_key)
+
+    def generator_clear(self) -> None:
+        self.redis.delete(self.generator_progress_key)
+        self.redis.delete(self.generator_complete_key)
+
+    def get_remaining(self) -> int:
+        remaining = cast(int, self.redis.scard(self.taskset_key))
+        return remaining
+
+    def get_active_task_count(self) -> int:
+        """Count of active permission sync tasks"""
+        count = 0
+        for _ in self.redis.scan_iter(RedisConnectorPermissionSync.FENCE_PREFIX + "*"):
+            count += 1
+        return count
+
+    @property
+    def fenced(self) -> bool:
+        if self.redis.exists(self.fence_key):
+            return True
+
+        return False
+
+    @property
+    def payload(self) -> RedisConnectorPermissionSyncPayload | None:
+        # read related data and evaluate/print task progress
+        fence_bytes = cast(bytes, self.redis.get(self.fence_key))
+        if fence_bytes is None:
+            return None
+
+        fence_str = fence_bytes.decode("utf-8")
+        payload = RedisConnectorPermissionSyncPayload.model_validate_json(
+            cast(str, fence_str)
+        )
+
+        return payload
+
+    def set_fence(
+        self,
+        payload: RedisConnectorPermissionSyncPayload | None,
+    ) -> None:
+        if not payload:
+            self.redis.delete(self.fence_key)
+            return
+
+        self.redis.set(self.fence_key, payload.model_dump_json())
+
+    @property
+    def generator_complete(self) -> int | None:
+        """the fence payload is an int representing the starting number of
+        permission sync tasks to be processed ... just after the generator completes."""
+        fence_bytes = self.redis.get(self.generator_complete_key)
+        if fence_bytes is None:
+            return None
+
+        if fence_bytes == b"None":
+            return None
+
+        fence_int = int(cast(bytes, fence_bytes).decode())
+        return fence_int
+
+    @generator_complete.setter
+    def generator_complete(self, payload: int | None) -> None:
+        """Set the payload to an int to set the fence, otherwise if None it will
+        be deleted"""
+        if payload is None:
+            self.redis.delete(self.generator_complete_key)
+            return
+
+        self.redis.set(self.generator_complete_key, payload)
+
+    def generate_tasks(
+        self,
+        celery_app: Celery,
+        lock: RedisLock | None,
+        new_permissions: list[DocExternalAccess],
+        source_string: str,
+        connector_id: int,
+        credential_id: int,
+    ) -> int | None:
+        last_lock_time = time.monotonic()
+        async_results = []
+
+        # Create a task for each document permission sync
+        for doc_perm in new_permissions:
+            current_time = time.monotonic()
+            if lock and current_time - last_lock_time >= (
+                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
+            ):
+                lock.reacquire()
+                last_lock_time = current_time
+            # Add task for document permissions sync
+            custom_task_id = f"{self.subtask_prefix}_{uuid4()}"
+            self.redis.sadd(self.taskset_key, custom_task_id)
+
+            result = celery_app.send_task(
+                OnyxCeleryTask.UPDATE_EXTERNAL_DOCUMENT_PERMISSIONS_TASK,
+                kwargs=dict(
+                    tenant_id=self.tenant_id,
+                    serialized_doc_external_access=doc_perm.to_dict(),
+                    source_string=source_string,
+                    connector_id=connector_id,
+                    credential_id=credential_id,
+                ),
+                queue=OnyxCeleryQueues.DOC_PERMISSIONS_UPSERT,
+                task_id=custom_task_id,
+                priority=OnyxCeleryPriority.MEDIUM,
+            )
+            async_results.append(result)
+
+        return len(async_results)
+
+    def reset(self) -> None:
+        self.redis.delete(self.generator_progress_key)
+        self.redis.delete(self.generator_complete_key)
+        self.redis.delete(self.taskset_key)
+        self.redis.delete(self.fence_key)
+
+    @staticmethod
+    def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:
+        taskset_key = f"{RedisConnectorPermissionSync.TASKSET_PREFIX}_{id}"
+        r.srem(taskset_key, task_id)
+        return
+
+    @staticmethod
+    def reset_all(r: redis.Redis) -> None:
+        """Deletes all redis values for all connectors"""
+        for key in r.scan_iter(RedisConnectorPermissionSync.TASKSET_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(
+            RedisConnectorPermissionSync.GENERATOR_COMPLETE_PREFIX + "*"
+        ):
+            r.delete(key)
+
+        for key in r.scan_iter(
+            RedisConnectorPermissionSync.GENERATOR_PROGRESS_PREFIX + "*"
+        ):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorPermissionSync.FENCE_PREFIX + "*"):
+            r.delete(key)

backend/onyx/redis/redis_connector_ext_group_sync.py

@@ -0,0 +1,158 @@
+from datetime import datetime
+from typing import cast
+
+import redis
+from celery import Celery
+from pydantic import BaseModel
+from redis.lock import Lock as RedisLock
+from sqlalchemy.orm import Session
+
+
+class RedisConnectorExternalGroupSyncPayload(BaseModel):
+    started: datetime | None
+    celery_task_id: str | None
+
+
+class RedisConnectorExternalGroupSync:
+    """Manages interactions with redis for external group syncing tasks. Should only be accessed
+    through RedisConnector."""
+
+    PREFIX = "connectorexternalgroupsync"
+
+    FENCE_PREFIX = f"{PREFIX}_fence"
+
+    # phase 1 - geneartor task and progress signals
+    GENERATORTASK_PREFIX = f"{PREFIX}+generator"  # connectorexternalgroupsync+generator
+    GENERATOR_PROGRESS_PREFIX = (
+        PREFIX + "_generator_progress"
+    )  # connectorexternalgroupsync_generator_progress
+    GENERATOR_COMPLETE_PREFIX = (
+        PREFIX + "_generator_complete"
+    )  # connectorexternalgroupsync_generator_complete
+
+    TASKSET_PREFIX = f"{PREFIX}_taskset"  # connectorexternalgroupsync_taskset
+    SUBTASK_PREFIX = f"{PREFIX}+sub"  # connectorexternalgroupsync+sub
+
+    def __init__(self, tenant_id: str | None, id: int, redis: redis.Redis) -> None:
+        self.tenant_id: str | None = tenant_id
+        self.id = id
+        self.redis = redis
+
+        self.fence_key: str = f"{self.FENCE_PREFIX}_{id}"
+        self.generator_task_key = f"{self.GENERATORTASK_PREFIX}_{id}"
+        self.generator_progress_key = f"{self.GENERATOR_PROGRESS_PREFIX}_{id}"
+        self.generator_complete_key = f"{self.GENERATOR_COMPLETE_PREFIX}_{id}"
+
+        self.taskset_key = f"{self.TASKSET_PREFIX}_{id}"
+
+        self.subtask_prefix: str = f"{self.SUBTASK_PREFIX}_{id}"
+
+    def taskset_clear(self) -> None:
+        self.redis.delete(self.taskset_key)
+
+    def generator_clear(self) -> None:
+        self.redis.delete(self.generator_progress_key)
+        self.redis.delete(self.generator_complete_key)
+
+    def get_remaining(self) -> int:
+        # todo: move into fence
+        remaining = cast(int, self.redis.scard(self.taskset_key))
+        return remaining
+
+    def get_active_task_count(self) -> int:
+        """Count of active external group syncing tasks"""
+        count = 0
+        for _ in self.redis.scan_iter(
+            RedisConnectorExternalGroupSync.FENCE_PREFIX + "*"
+        ):
+            count += 1
+        return count
+
+    @property
+    def fenced(self) -> bool:
+        if self.redis.exists(self.fence_key):
+            return True
+
+        return False
+
+    @property
+    def payload(self) -> RedisConnectorExternalGroupSyncPayload | None:
+        # read related data and evaluate/print task progress
+        fence_bytes = cast(bytes, self.redis.get(self.fence_key))
+        if fence_bytes is None:
+            return None
+
+        fence_str = fence_bytes.decode("utf-8")
+        payload = RedisConnectorExternalGroupSyncPayload.model_validate_json(
+            cast(str, fence_str)
+        )
+
+        return payload
+
+    def set_fence(
+        self,
+        payload: RedisConnectorExternalGroupSyncPayload | None,
+    ) -> None:
+        if not payload:
+            self.redis.delete(self.fence_key)
+            return
+
+        self.redis.set(self.fence_key, payload.model_dump_json())
+
+    @property
+    def generator_complete(self) -> int | None:
+        """the fence payload is an int representing the starting number of
+        external group syncing tasks to be processed ... just after the generator completes.
+        """
+        fence_bytes = self.redis.get(self.generator_complete_key)
+        if fence_bytes is None:
+            return None
+
+        if fence_bytes == b"None":
+            return None
+
+        fence_int = int(cast(bytes, fence_bytes).decode())
+        return fence_int
+
+    @generator_complete.setter
+    def generator_complete(self, payload: int | None) -> None:
+        """Set the payload to an int to set the fence, otherwise if None it will
+        be deleted"""
+        if payload is None:
+            self.redis.delete(self.generator_complete_key)
+            return
+
+        self.redis.set(self.generator_complete_key, payload)
+
+    def generate_tasks(
+        self,
+        celery_app: Celery,
+        db_session: Session,
+        lock: RedisLock | None,
+    ) -> int | None:
+        pass
+
+    @staticmethod
+    def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:
+        taskset_key = f"{RedisConnectorExternalGroupSync.TASKSET_PREFIX}_{id}"
+        r.srem(taskset_key, task_id)
+        return
+
+    @staticmethod
+    def reset_all(r: redis.Redis) -> None:
+        """Deletes all redis values for all connectors"""
+        for key in r.scan_iter(RedisConnectorExternalGroupSync.TASKSET_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(
+            RedisConnectorExternalGroupSync.GENERATOR_COMPLETE_PREFIX + "*"
+        ):
+            r.delete(key)
+
+        for key in r.scan_iter(
+            RedisConnectorExternalGroupSync.GENERATOR_PROGRESS_PREFIX + "*"
+        ):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorExternalGroupSync.FENCE_PREFIX + "*"):
+            r.delete(key)

backend/onyx/redis/redis_connector_index.py

@@ -0,0 +1,159 @@
+from datetime import datetime
+from typing import cast
+from uuid import uuid4
+
+import redis
+from pydantic import BaseModel
+
+
+class RedisConnectorIndexPayload(BaseModel):
+    index_attempt_id: int | None
+    started: datetime | None
+    submitted: datetime
+    celery_task_id: str | None
+
+
+class RedisConnectorIndex:
+    """Manages interactions with redis for indexing tasks. Should only be accessed
+    through RedisConnector."""
+
+    PREFIX = "connectorindexing"
+    FENCE_PREFIX = f"{PREFIX}_fence"  # "connectorindexing_fence"
+    GENERATOR_TASK_PREFIX = PREFIX + "+generator"  # "connectorindexing+generator_fence"
+    GENERATOR_PROGRESS_PREFIX = (
+        PREFIX + "_generator_progress"
+    )  # connectorindexing_generator_progress
+    GENERATOR_COMPLETE_PREFIX = (
+        PREFIX + "_generator_complete"
+    )  # connectorindexing_generator_complete
+
+    GENERATOR_LOCK_PREFIX = "da_lock:indexing"
+
+    TERMINATE_PREFIX = PREFIX + "_terminate"  # connectorindexing_terminate
+
+    def __init__(
+        self,
+        tenant_id: str | None,
+        id: int,
+        search_settings_id: int,
+        redis: redis.Redis,
+    ) -> None:
+        self.tenant_id: str | None = tenant_id
+        self.id = id
+        self.search_settings_id = search_settings_id
+        self.redis = redis
+
+        self.fence_key: str = f"{self.FENCE_PREFIX}_{id}/{search_settings_id}"
+        self.generator_progress_key = (
+            f"{self.GENERATOR_PROGRESS_PREFIX}_{id}/{search_settings_id}"
+        )
+        self.generator_complete_key = (
+            f"{self.GENERATOR_COMPLETE_PREFIX}_{id}/{search_settings_id}"
+        )
+        self.generator_lock_key = (
+            f"{self.GENERATOR_LOCK_PREFIX}_{id}/{search_settings_id}"
+        )
+        self.terminate_key = f"{self.TERMINATE_PREFIX}_{id}/{search_settings_id}"
+
+    @classmethod
+    def fence_key_with_ids(cls, cc_pair_id: int, search_settings_id: int) -> str:
+        return f"{cls.FENCE_PREFIX}_{cc_pair_id}/{search_settings_id}"
+
+    def generate_generator_task_id(self) -> str:
+        # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
+        # we prefix the task id so it's easier to keep track of who created the task
+        # aka "connectorindexing+generator_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
+
+        return f"{self.GENERATOR_TASK_PREFIX}_{self.id}/{self.search_settings_id}_{uuid4()}"
+
+    @property
+    def fenced(self) -> bool:
+        if self.redis.exists(self.fence_key):
+            return True
+
+        return False
+
+    @property
+    def payload(self) -> RedisConnectorIndexPayload | None:
+        # read related data and evaluate/print task progress
+        fence_bytes = cast(bytes, self.redis.get(self.fence_key))
+        if fence_bytes is None:
+            return None
+
+        fence_str = fence_bytes.decode("utf-8")
+        payload = RedisConnectorIndexPayload.model_validate_json(cast(str, fence_str))
+
+        return payload
+
+    def set_fence(
+        self,
+        payload: RedisConnectorIndexPayload | None,
+    ) -> None:
+        if not payload:
+            self.redis.delete(self.fence_key)
+            return
+
+        self.redis.set(self.fence_key, payload.model_dump_json())
+
+    def terminating(self, celery_task_id: str) -> bool:
+        if self.redis.exists(f"{self.terminate_key}_{celery_task_id}"):
+            return True
+
+        return False
+
+    def set_terminate(self, celery_task_id: str) -> None:
+        """This sets a signal. It does not block!"""
+        # We shouldn't need very long to terminate the spawned task.
+        # 10 minute TTL is good.
+        self.redis.set(f"{self.terminate_key}_{celery_task_id}", 0, ex=600)
+
+    def set_generator_complete(self, payload: int | None) -> None:
+        if not payload:
+            self.redis.delete(self.generator_complete_key)
+            return
+
+        self.redis.set(self.generator_complete_key, payload)
+
+    def generator_clear(self) -> None:
+        self.redis.delete(self.generator_progress_key)
+        self.redis.delete(self.generator_complete_key)
+
+    def get_progress(self) -> int | None:
+        """Returns None if the key doesn't exist. The"""
+        # TODO: move into fence?
+        bytes = self.redis.get(self.generator_progress_key)
+        if bytes is None:
+            return None
+
+        progress = int(cast(int, bytes))
+        return progress
+
+    def get_completion(self) -> int | None:
+        # TODO: move into fence?
+        bytes = self.redis.get(self.generator_complete_key)
+        if bytes is None:
+            return None
+
+        status = int(cast(int, bytes))
+        return status
+
+    def reset(self) -> None:
+        self.redis.delete(self.generator_lock_key)
+        self.redis.delete(self.generator_progress_key)
+        self.redis.delete(self.generator_complete_key)
+        self.redis.delete(self.fence_key)
+
+    @staticmethod
+    def reset_all(r: redis.Redis) -> None:
+        """Deletes all redis values for all connectors"""
+        for key in r.scan_iter(RedisConnectorIndex.GENERATOR_LOCK_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorIndex.GENERATOR_COMPLETE_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorIndex.GENERATOR_PROGRESS_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorIndex.FENCE_PREFIX + "*"):
+            r.delete(key)

backend/onyx/redis/redis_connector_prune.py

@@ -0,0 +1,179 @@
+import time
+from typing import cast
+from uuid import uuid4
+
+import redis
+from celery import Celery
+from redis.lock import Lock as RedisLock
+from sqlalchemy.orm import Session
+
+from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
+from onyx.configs.constants import OnyxCeleryPriority
+from onyx.configs.constants import OnyxCeleryQueues
+from onyx.configs.constants import OnyxCeleryTask
+from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
+
+
+class RedisConnectorPrune:
+    """Manages interactions with redis for pruning tasks. Should only be accessed
+    through RedisConnector."""
+
+    PREFIX = "connectorpruning"
+
+    FENCE_PREFIX = f"{PREFIX}_fence"
+
+    # phase 1 - geneartor task and progress signals
+    GENERATORTASK_PREFIX = f"{PREFIX}+generator"  # connectorpruning+generator
+    GENERATOR_PROGRESS_PREFIX = (
+        PREFIX + "_generator_progress"
+    )  # connectorpruning_generator_progress
+    GENERATOR_COMPLETE_PREFIX = (
+        PREFIX + "_generator_complete"
+    )  # connectorpruning_generator_complete
+
+    TASKSET_PREFIX = f"{PREFIX}_taskset"  # connectorpruning_taskset
+    SUBTASK_PREFIX = f"{PREFIX}+sub"  # connectorpruning+sub
+
+    def __init__(self, tenant_id: str | None, id: int, redis: redis.Redis) -> None:
+        self.tenant_id: str | None = tenant_id
+        self.id = id
+        self.redis = redis
+
+        self.fence_key: str = f"{self.FENCE_PREFIX}_{id}"
+        self.generator_task_key = f"{self.GENERATORTASK_PREFIX}_{id}"
+        self.generator_progress_key = f"{self.GENERATOR_PROGRESS_PREFIX}_{id}"
+        self.generator_complete_key = f"{self.GENERATOR_COMPLETE_PREFIX}_{id}"
+
+        self.taskset_key = f"{self.TASKSET_PREFIX}_{id}"
+
+        self.subtask_prefix: str = f"{self.SUBTASK_PREFIX}_{id}"
+
+    def taskset_clear(self) -> None:
+        self.redis.delete(self.taskset_key)
+
+    def generator_clear(self) -> None:
+        self.redis.delete(self.generator_progress_key)
+        self.redis.delete(self.generator_complete_key)
+
+    def get_remaining(self) -> int:
+        # todo: move into fence
+        remaining = cast(int, self.redis.scard(self.taskset_key))
+        return remaining
+
+    def get_active_task_count(self) -> int:
+        """Count of active pruning tasks"""
+        count = 0
+        for key in self.redis.scan_iter(RedisConnectorPrune.FENCE_PREFIX + "*"):
+            count += 1
+        return count
+
+    @property
+    def fenced(self) -> bool:
+        if self.redis.exists(self.fence_key):
+            return True
+
+        return False
+
+    def set_fence(self, value: bool) -> None:
+        if not value:
+            self.redis.delete(self.fence_key)
+            return
+
+        self.redis.set(self.fence_key, 0)
+
+    @property
+    def generator_complete(self) -> int | None:
+        """the fence payload is an int representing the starting number of
+        pruning tasks to be processed ... just after the generator completes."""
+        fence_bytes = self.redis.get(self.generator_complete_key)
+        if fence_bytes is None:
+            return None
+
+        fence_int = cast(int, fence_bytes)
+        return fence_int
+
+    @generator_complete.setter
+    def generator_complete(self, payload: int | None) -> None:
+        """Set the payload to an int to set the fence, otherwise if None it will
+        be deleted"""
+        if payload is None:
+            self.redis.delete(self.generator_complete_key)
+            return
+
+        self.redis.set(self.generator_complete_key, payload)
+
+    def generate_tasks(
+        self,
+        documents_to_prune: set[str],
+        celery_app: Celery,
+        db_session: Session,
+        lock: RedisLock | None,
+    ) -> int | None:
+        last_lock_time = time.monotonic()
+
+        async_results = []
+        cc_pair = get_connector_credential_pair_from_id(int(self.id), db_session)
+        if not cc_pair:
+            return None
+
+        for doc_id in documents_to_prune:
+            current_time = time.monotonic()
+            if lock and current_time - last_lock_time >= (
+                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
+            ):
+                lock.reacquire()
+                last_lock_time = current_time
+
+            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            # the actual redis key is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            # we prefix the task id so it's easier to keep track of who created the task
+            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            custom_task_id = f"{self.subtask_prefix}_{uuid4()}"
+
+            # add to the tracking taskset in redis BEFORE creating the celery task.
+            self.redis.sadd(self.taskset_key, custom_task_id)
+
+            # Priority on sync's triggered by new indexing should be medium
+            result = celery_app.send_task(
+                OnyxCeleryTask.DOCUMENT_BY_CC_PAIR_CLEANUP_TASK,
+                kwargs=dict(
+                    document_id=doc_id,
+                    connector_id=cc_pair.connector_id,
+                    credential_id=cc_pair.credential_id,
+                    tenant_id=self.tenant_id,
+                ),
+                queue=OnyxCeleryQueues.CONNECTOR_DELETION,
+                task_id=custom_task_id,
+                priority=OnyxCeleryPriority.MEDIUM,
+            )
+
+            async_results.append(result)
+
+        return len(async_results)
+
+    def reset(self) -> None:
+        self.redis.delete(self.generator_progress_key)
+        self.redis.delete(self.generator_complete_key)
+        self.redis.delete(self.taskset_key)
+        self.redis.delete(self.fence_key)
+
+    @staticmethod
+    def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:
+        taskset_key = f"{RedisConnectorPrune.TASKSET_PREFIX}_{id}"
+        r.srem(taskset_key, task_id)
+        return
+
+    @staticmethod
+    def reset_all(r: redis.Redis) -> None:
+        """Deletes all redis values for all connectors"""
+        for key in r.scan_iter(RedisConnectorPrune.TASKSET_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorPrune.GENERATOR_COMPLETE_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorPrune.GENERATOR_PROGRESS_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisConnectorPrune.FENCE_PREFIX + "*"):
+            r.delete(key)

backend/onyx/redis/redis_connector_stop.py

@@ -0,0 +1,34 @@
+import redis
+
+
+class RedisConnectorStop:
+    """Manages interactions with redis for stop signaling. Should only be accessed
+    through RedisConnector."""
+
+    FENCE_PREFIX = "connectorstop_fence"
+
+    def __init__(self, tenant_id: str | None, id: int, redis: redis.Redis) -> None:
+        self.tenant_id: str | None = tenant_id
+        self.id: int = id
+        self.redis = redis
+
+        self.fence_key: str = f"{self.FENCE_PREFIX}_{id}"
+
+    @property
+    def fenced(self) -> bool:
+        if self.redis.exists(self.fence_key):
+            return True
+
+        return False
+
+    def set_fence(self, value: bool) -> None:
+        if not value:
+            self.redis.delete(self.fence_key)
+            return
+
+        self.redis.set(self.fence_key, 0)
+
+    @staticmethod
+    def reset_all(r: redis.Redis) -> None:
+        for key in r.scan_iter(RedisConnectorStop.FENCE_PREFIX + "*"):
+            r.delete(key)

backend/onyx/redis/redis_document_set.py

@@ -0,0 +1,101 @@
+import time
+from typing import cast
+from uuid import uuid4
+
+import redis
+from celery import Celery
+from redis import Redis
+from redis.lock import Lock as RedisLock
+from sqlalchemy.orm import Session
+
+from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
+from onyx.configs.constants import OnyxCeleryPriority
+from onyx.configs.constants import OnyxCeleryQueues
+from onyx.configs.constants import OnyxCeleryTask
+from onyx.db.document_set import construct_document_select_by_docset
+from onyx.redis.redis_object_helper import RedisObjectHelper
+
+
+class RedisDocumentSet(RedisObjectHelper):
+    PREFIX = "documentset"
+    FENCE_PREFIX = PREFIX + "_fence"
+    TASKSET_PREFIX = PREFIX + "_taskset"
+
+    def __init__(self, tenant_id: str | None, id: int) -> None:
+        super().__init__(tenant_id, str(id))
+
+    @property
+    def fenced(self) -> bool:
+        if self.redis.exists(self.fence_key):
+            return True
+
+        return False
+
+    def set_fence(self, payload: int | None) -> None:
+        if payload is None:
+            self.redis.delete(self.fence_key)
+            return
+
+        self.redis.set(self.fence_key, payload)
+
+    @property
+    def payload(self) -> int | None:
+        bytes = self.redis.get(self.fence_key)
+        if bytes is None:
+            return None
+
+        progress = int(cast(int, bytes))
+        return progress
+
+    def generate_tasks(
+        self,
+        celery_app: Celery,
+        db_session: Session,
+        redis_client: Redis,
+        lock: RedisLock,
+        tenant_id: str | None,
+    ) -> tuple[int, int] | None:
+        last_lock_time = time.monotonic()
+
+        async_results = []
+        stmt = construct_document_select_by_docset(int(self._id), current_only=False)
+        for doc in db_session.scalars(stmt).yield_per(1):
+            current_time = time.monotonic()
+            if current_time - last_lock_time >= (
+                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
+            ):
+                lock.reacquire()
+                last_lock_time = current_time
+
+            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            # the key for the result is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            # we prefix the task id so it's easier to keep track of who created the task
+            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            custom_task_id = f"{self.task_id_prefix}_{uuid4()}"
+
+            # add to the set BEFORE creating the task.
+            redis_client.sadd(self.taskset_key, custom_task_id)
+
+            result = celery_app.send_task(
+                OnyxCeleryTask.VESPA_METADATA_SYNC_TASK,
+                kwargs=dict(document_id=doc.id, tenant_id=tenant_id),
+                queue=OnyxCeleryQueues.VESPA_METADATA_SYNC,
+                task_id=custom_task_id,
+                priority=OnyxCeleryPriority.LOW,
+            )
+
+            async_results.append(result)
+
+        return len(async_results), len(async_results)
+
+    def reset(self) -> None:
+        self.redis.delete(self.taskset_key)
+        self.redis.delete(self.fence_key)
+
+    @staticmethod
+    def reset_all(r: redis.Redis) -> None:
+        for key in r.scan_iter(RedisDocumentSet.TASKSET_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + "*"):
+            r.delete(key)

backend/onyx/redis/redis_object_helper.py

@@ -0,0 +1,97 @@
+from abc import ABC
+from abc import abstractmethod
+
+from celery import Celery
+from redis import Redis
+from redis.lock import Lock as RedisLock
+from sqlalchemy.orm import Session
+
+from onyx.redis.redis_pool import get_redis_client
+
+
+class RedisObjectHelper(ABC):
+    PREFIX = "base"
+    FENCE_PREFIX = PREFIX + "_fence"
+    TASKSET_PREFIX = PREFIX + "_taskset"
+
+    def __init__(self, tenant_id: str | None, id: str):
+        self._tenant_id: str | None = tenant_id
+        self._id: str = id
+        self.redis = get_redis_client(tenant_id=tenant_id)
+
+    @property
+    def task_id_prefix(self) -> str:
+        return f"{self.PREFIX}_{self._id}"
+
+    @property
+    def fence_key(self) -> str:
+        # example: documentset_fence_1
+        return f"{self.FENCE_PREFIX}_{self._id}"
+
+    @property
+    def taskset_key(self) -> str:
+        # example: documentset_taskset_1
+        return f"{self.TASKSET_PREFIX}_{self._id}"
+
+    @staticmethod
+    def get_id_from_fence_key(key: str) -> str | None:
+        """
+        Extracts the object ID from a fence key in the format `PREFIX_fence_X`.
+
+        Args:
+            key (str): The fence key string.
+
+        Returns:
+            Optional[int]: The extracted ID if the key is in the correct format, otherwise None.
+        """
+        parts = key.split("_")
+        if len(parts) != 3:
+            return None
+
+        object_id = parts[2]
+        return object_id
+
+    @staticmethod
+    def get_id_from_task_id(task_id: str) -> str | None:
+        """
+        Extracts the object ID from a task ID string.
+
+        This method assumes the task ID is formatted as `prefix_objectid_suffix`, where:
+        - `prefix` is an arbitrary string (e.g., the name of the task or entity),
+        - `objectid` is the ID you want to extract,
+        - `suffix` is another arbitrary string (e.g., a UUID).
+
+        Example:
+            If the input `task_id` is `documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc`,
+            this method will return the string `"1"`.
+
+        Args:
+            task_id (str): The task ID string from which to extract the object ID.
+
+        Returns:
+            str | None: The extracted object ID if the task ID is in the correct format, otherwise None.
+        """
+        # example: task_id=documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc
+        parts = task_id.split("_")
+        if len(parts) != 3:
+            return None
+
+        object_id = parts[1]
+        return object_id
+
+    @abstractmethod
+    def generate_tasks(
+        self,
+        celery_app: Celery,
+        db_session: Session,
+        redis_client: Redis,
+        lock: RedisLock,
+        tenant_id: str | None,
+    ) -> tuple[int, int] | None:
+        """First element should be the number of actual tasks generated, second should
+        be the number of docs that were candidates to be synced for the cc pair.
+
+        The need for this is when we are syncing stale docs referenced by multiple
+        connectors. In a single pass across multiple cc pairs, we only want a task
+        for be created for a particular document id the first time we see it.
+        The rest can be skipped."""

backend/onyx/redis/redis_pool.py

@@ -0,0 +1,198 @@
+import functools
+import threading
+from collections.abc import Callable
+from typing import Any
+from typing import Optional
+
+import redis
+from redis.client import Redis
+
+from onyx.configs.app_configs import REDIS_DB_NUMBER
+from onyx.configs.app_configs import REDIS_HEALTH_CHECK_INTERVAL
+from onyx.configs.app_configs import REDIS_HOST
+from onyx.configs.app_configs import REDIS_PASSWORD
+from onyx.configs.app_configs import REDIS_POOL_MAX_CONNECTIONS
+from onyx.configs.app_configs import REDIS_PORT
+from onyx.configs.app_configs import REDIS_SSL
+from onyx.configs.app_configs import REDIS_SSL_CA_CERTS
+from onyx.configs.app_configs import REDIS_SSL_CERT_REQS
+from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+class TenantRedis(redis.Redis):
+    def __init__(self, tenant_id: str, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.tenant_id: str = tenant_id
+
+    def _prefixed(self, key: str | bytes | memoryview) -> str | bytes | memoryview:
+        prefix: str = f"{self.tenant_id}:"
+        if isinstance(key, str):
+            if key.startswith(prefix):
+                return key
+            else:
+                return prefix + key
+        elif isinstance(key, bytes):
+            prefix_bytes = prefix.encode()
+            if key.startswith(prefix_bytes):
+                return key
+            else:
+                return prefix_bytes + key
+        elif isinstance(key, memoryview):
+            key_bytes = key.tobytes()
+            prefix_bytes = prefix.encode()
+            if key_bytes.startswith(prefix_bytes):
+                return key
+            else:
+                return memoryview(prefix_bytes + key_bytes)
+        else:
+            raise TypeError(f"Unsupported key type: {type(key)}")
+
+    def _prefix_method(self, method: Callable) -> Callable:
+        @functools.wraps(method)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if "name" in kwargs:
+                kwargs["name"] = self._prefixed(kwargs["name"])
+            elif len(args) > 0:
+                args = (self._prefixed(args[0]),) + args[1:]
+            return method(*args, **kwargs)
+
+        return wrapper
+
+    def _prefix_scan_iter(self, method: Callable) -> Callable:
+        @functools.wraps(method)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Prefix the match pattern if provided
+            if "match" in kwargs:
+                kwargs["match"] = self._prefixed(kwargs["match"])
+            elif len(args) > 0:
+                args = (self._prefixed(args[0]),) + args[1:]
+
+            # Get the iterator
+            iterator = method(*args, **kwargs)
+
+            # Remove prefix from returned keys
+            prefix = f"{self.tenant_id}:".encode()
+            prefix_len = len(prefix)
+
+            for key in iterator:
+                if isinstance(key, bytes) and key.startswith(prefix):
+                    yield key[prefix_len:]
+                else:
+                    yield key
+
+        return wrapper
+
+    def __getattribute__(self, item: str) -> Any:
+        original_attr = super().__getattribute__(item)
+        methods_to_wrap = [
+            "lock",
+            "unlock",
+            "get",
+            "set",
+            "delete",
+            "exists",
+            "incrby",
+            "hset",
+            "hget",
+            "getset",
+            "owned",
+            "reacquire",
+            "create_lock",
+            "startswith",
+            "sadd",
+            "srem",
+            "scard",
+        ]  # Regular methods that need simple prefixing
+
+        if item == "scan_iter":
+            return self._prefix_scan_iter(original_attr)
+        elif item in methods_to_wrap and callable(original_attr):
+            return self._prefix_method(original_attr)
+        return original_attr
+
+
+class RedisPool:
+    _instance: Optional["RedisPool"] = None
+    _lock: threading.Lock = threading.Lock()
+    _pool: redis.BlockingConnectionPool
+
+    def __new__(cls) -> "RedisPool":
+        if not cls._instance:
+            with cls._lock:
+                if not cls._instance:
+                    cls._instance = super(RedisPool, cls).__new__(cls)
+                    cls._instance._init_pool()
+        return cls._instance
+
+    def _init_pool(self) -> None:
+        self._pool = RedisPool.create_pool(ssl=REDIS_SSL)
+
+    def get_client(self, tenant_id: str | None) -> Redis:
+        if tenant_id is None:
+            tenant_id = "public"
+        return TenantRedis(tenant_id, connection_pool=self._pool)
+
+    @staticmethod
+    def create_pool(
+        host: str = REDIS_HOST,
+        port: int = REDIS_PORT,
+        db: int = REDIS_DB_NUMBER,
+        password: str = REDIS_PASSWORD,
+        max_connections: int = REDIS_POOL_MAX_CONNECTIONS,
+        ssl_ca_certs: str | None = REDIS_SSL_CA_CERTS,
+        ssl_cert_reqs: str = REDIS_SSL_CERT_REQS,
+        ssl: bool = False,
+    ) -> redis.BlockingConnectionPool:
+        """We use BlockingConnectionPool because it will block and wait for a connection
+        rather than error if max_connections is reached. This is far more deterministic
+        behavior and aligned with how we want to use Redis."""
+
+        # Using ConnectionPool is not well documented.
+        # Useful examples: https://github.com/redis/redis-py/issues/780
+        if ssl:
+            return redis.BlockingConnectionPool(
+                host=host,
+                port=port,
+                db=db,
+                password=password,
+                max_connections=max_connections,
+                timeout=None,
+                health_check_interval=REDIS_HEALTH_CHECK_INTERVAL,
+                socket_keepalive=True,
+                socket_keepalive_options=REDIS_SOCKET_KEEPALIVE_OPTIONS,
+                connection_class=redis.SSLConnection,
+                ssl_ca_certs=ssl_ca_certs,
+                ssl_cert_reqs=ssl_cert_reqs,
+            )
+
+        return redis.BlockingConnectionPool(
+            host=host,
+            port=port,
+            db=db,
+            password=password,
+            max_connections=max_connections,
+            timeout=None,
+            health_check_interval=REDIS_HEALTH_CHECK_INTERVAL,
+            socket_keepalive=True,
+            socket_keepalive_options=REDIS_SOCKET_KEEPALIVE_OPTIONS,
+        )
+
+
+redis_pool = RedisPool()
+
+
+def get_redis_client(*, tenant_id: str | None) -> Redis:
+    return redis_pool.get_client(tenant_id)
+
+
+# # Usage example
+# redis_pool = RedisPool()
+# redis_client = redis_pool.get_client()
+
+# # Example of setting and getting a value
+# redis_client.set('key', 'value')
+# value = redis_client.get('key')
+# print(value.decode())  # Output: 'value'

backend/onyx/redis/redis_usergroup.py

@@ -0,0 +1,114 @@
+import time
+from typing import cast
+from uuid import uuid4
+
+import redis
+from celery import Celery
+from redis import Redis
+from redis.lock import Lock as RedisLock
+from sqlalchemy.orm import Session
+
+from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
+from onyx.configs.constants import OnyxCeleryPriority
+from onyx.configs.constants import OnyxCeleryQueues
+from onyx.configs.constants import OnyxCeleryTask
+from onyx.redis.redis_object_helper import RedisObjectHelper
+from onyx.utils.variable_functionality import fetch_versioned_implementation
+from onyx.utils.variable_functionality import global_version
+
+
+class RedisUserGroup(RedisObjectHelper):
+    PREFIX = "usergroup"
+    FENCE_PREFIX = PREFIX + "_fence"
+    TASKSET_PREFIX = PREFIX + "_taskset"
+
+    def __init__(self, tenant_id: str | None, id: int) -> None:
+        super().__init__(tenant_id, str(id))
+
+    @property
+    def fenced(self) -> bool:
+        if self.redis.exists(self.fence_key):
+            return True
+
+        return False
+
+    def set_fence(self, payload: int | None) -> None:
+        if payload is None:
+            self.redis.delete(self.fence_key)
+            return
+
+        self.redis.set(self.fence_key, payload)
+
+    @property
+    def payload(self) -> int | None:
+        bytes = self.redis.get(self.fence_key)
+        if bytes is None:
+            return None
+
+        progress = int(cast(int, bytes))
+        return progress
+
+    def generate_tasks(
+        self,
+        celery_app: Celery,
+        db_session: Session,
+        redis_client: Redis,
+        lock: RedisLock,
+        tenant_id: str | None,
+    ) -> tuple[int, int] | None:
+        last_lock_time = time.monotonic()
+
+        async_results = []
+
+        if not global_version.is_ee_version():
+            return 0, 0
+
+        try:
+            construct_document_select_by_usergroup = fetch_versioned_implementation(
+                "onyx.db.user_group",
+                "construct_document_select_by_usergroup",
+            )
+        except ModuleNotFoundError:
+            return 0, 0
+
+        stmt = construct_document_select_by_usergroup(int(self._id))
+        for doc in db_session.scalars(stmt).yield_per(1):
+            current_time = time.monotonic()
+            if current_time - last_lock_time >= (
+                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
+            ):
+                lock.reacquire()
+                last_lock_time = current_time
+
+            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            # the key for the result is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            # we prefix the task id so it's easier to keep track of who created the task
+            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
+            custom_task_id = f"{self.task_id_prefix}_{uuid4()}"
+
+            # add to the set BEFORE creating the task.
+            redis_client.sadd(self.taskset_key, custom_task_id)
+
+            result = celery_app.send_task(
+                OnyxCeleryTask.VESPA_METADATA_SYNC_TASK,
+                kwargs=dict(document_id=doc.id, tenant_id=tenant_id),
+                queue=OnyxCeleryQueues.VESPA_METADATA_SYNC,
+                task_id=custom_task_id,
+                priority=OnyxCeleryPriority.LOW,
+            )
+
+            async_results.append(result)
+
+        return len(async_results), len(async_results)
+
+    def reset(self) -> None:
+        self.redis.delete(self.taskset_key)
+        self.redis.delete(self.fence_key)
+
+    @staticmethod
+    def reset_all(r: redis.Redis) -> None:
+        for key in r.scan_iter(RedisUserGroup.TASKSET_PREFIX + "*"):
+            r.delete(key)
+
+        for key in r.scan_iter(RedisUserGroup.FENCE_PREFIX + "*"):
+            r.delete(key)