Add custom SOQL condition handling (#16)

piotrekkr · web-flow · commit 24a7ed251571 · 2025-06-02T17:06:29.000+02:00
diff --git a/config.example.yaml b/config.example.yaml
@@ -73,6 +73,14 @@ objects:
     #        Check project readme file for more details.
     dir_name_field: LinkedEntity.Username
 
+    # Purpose: Use value of this field as an extra condition in SOQL query when fetching objects
+    # Required: NO
+    # Default: None
+    # Notes: In contrary to `modified_date_*`, it filters objects and not files attached to them.
+    #        It can filter objects by any field of the object that is allowed by Salesforce.
+    extra_soql_condition: "MyCustomField__c = 'MyValue'"
+
+
   # Example 1: Process ContentDocument objects attached to `Account` object, older than datetime
   Account:
     modified_date_lt: 2023-08-01T00:00:00Z
diff --git a/src/salesforce_archivist/archivist.py b/src/salesforce_archivist/archivist.py
@@ -25,6 +25,7 @@ class ArchivistObject(BaseModel):
     modified_date_lt: Optional[datetime.datetime] = None
     modified_date_gt: Optional[datetime.datetime] = None
     dir_name_field: Optional[str] = None
+    extra_soql_condition: Optional[str] = None
 
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, type(self)):
diff --git a/src/salesforce_archivist/salesforce/salesforce.py b/src/salesforce_archivist/salesforce/salesforce.py
@@ -52,35 +52,36 @@ def _get_content_document_list_query(self) -> str:
         select_list = ["LinkedEntityId", "ContentDocumentId", "LinkedEntity.Type"]
         if self._archivist_obj.dir_name_field is not None and self._archivist_obj.dir_name_field not in select_list:
             select_list.append(self._archivist_obj.dir_name_field)
-        where_conditions = []
+        where = []
+        linked_id_where = "LinkedEntityId IN (SELECT Id FROM {obj_type})".format(obj_type=self._archivist_obj.obj_type)
+        if self._archivist_obj.extra_soql_condition is not None:
+            linked_id_where = "LinkedEntityId IN (SELECT Id FROM {obj_type} WHERE {where})".format(
+                obj_type=self._archivist_obj.obj_type, where=self._archivist_obj.extra_soql_condition
+            )
+        where.append(linked_id_where)
+        doc_id_where = []
         if self._archivist_obj.modified_date_lt is not None:
-            where_conditions.append(
+            doc_id_where.append(
                 "ContentDocument.ContentModifiedDate < {date}".format(
                     date=self._archivist_obj.modified_date_lt.strftime("%Y-%m-%dT%H:%M:%SZ")
                 )
             )
         if self._archivist_obj.modified_date_gt is not None:
-            where_conditions.append(
+            doc_id_where.append(
                 "ContentDocument.ContentModifiedDate > {date}".format(
                     date=self._archivist_obj.modified_date_gt.strftime("%Y-%m-%dT%H:%M:%SZ")
                 )
             )
-        where = ""
-        if len(where_conditions):
-            where = "WHERE {}".format(" AND ".join(where_conditions))
-
-        # Using WHERE IN and not using filter on `LinkedEntity.Type` is done because of SF restrictions like:
-        #
-        #   Implementation restriction: ContentDocumentLink requires a filter by a single ID on ContentDocumentId
-        #   or LinkedEntityId using the equals operator or multiple ID's using the IN operator.
-        #
-        #   Implementation restriction: filtering on non-id fields is only permitted when filtering
-        #   by ContentDocumentLink.LinkedEntityId using the equal operator.
+        if len(doc_id_where):
+            where.append(
+                "ContentDocumentId IN (SELECT Id FROM ContentDocument WHERE {where})".format(
+                    where=" AND ".join(doc_id_where)
+                )
+            )
 
-        return (
-            "SELECT {fields} FROM ContentDocumentLink "
-            "WHERE ContentDocumentId IN (SELECT Id FROM ContentDocument {where})"
-        ).format(fields=", ".join(select_list), where=where)
+        return ("SELECT {fields} " "FROM ContentDocumentLink " "WHERE {where}").format(
+            fields=", ".join(select_list), where=" AND ".join(where)
+        )
 
     def download_content_document_link_list(
         self,
@@ -187,6 +188,8 @@ def _get_attachment_list_query(self) -> str:
                     date=self._archivist_obj.modified_date_gt.strftime("%Y-%m-%dT%H:%M:%SZ")
                 )
             )
+        if self._archivist_obj.extra_soql_condition is not None:
+            where_conditions.append(self._archivist_obj.extra_soql_condition)
         where = ""
         if len(where_conditions):
             where = "WHERE {}".format(" AND ".join(where_conditions))
diff --git a/test/salesforce/test_salesforce.py b/test/salesforce/test_salesforce.py
@@ -18,13 +18,15 @@
 
 
 @pytest.mark.parametrize(
-    "modified_date_lt, modified_date_gt, dir_name_field, expected_query",
+    "modified_date_lt, modified_date_gt, dir_name_field, extra_soql_condition, expected_query",
     [
         (
             None,
             None,
             None,
-            "SELECT LinkedEntityId, ContentDocumentId, LinkedEntity.Type FROM ContentDocumentLink WHERE ContentDocumentId IN (SELECT Id FROM ContentDocument )",
+            None,
+            "SELECT LinkedEntityId, ContentDocumentId, LinkedEntity.Type FROM ContentDocumentLink "
+            "WHERE LinkedEntityId IN (SELECT Id FROM User)",
         ),
         (
             datetime(
@@ -39,10 +41,14 @@
             ),
             None,
             None,
+            None,
             (
                 "SELECT LinkedEntityId, ContentDocumentId, LinkedEntity.Type "
                 "FROM ContentDocumentLink "
-                "WHERE ContentDocumentId IN (SELECT Id FROM ContentDocument WHERE ContentDocument.ContentModifiedDate < 2024-01-01T00:00:00Z)"
+                "WHERE LinkedEntityId IN (SELECT Id FROM User) "
+                "AND ContentDocumentId IN ("
+                "SELECT Id FROM ContentDocument WHERE ContentDocument.ContentModifiedDate < 2024-01-01T00:00:00Z"
+                ")"
             ),
         ),
         (
@@ -67,11 +73,51 @@
                 tzinfo=timezone.utc,
             ),
             None,
+            None,
             (
                 "SELECT LinkedEntityId, ContentDocumentId, LinkedEntity.Type "
                 "FROM ContentDocumentLink "
-                "WHERE ContentDocumentId IN ("
-                "SELECT Id FROM ContentDocument WHERE ContentDocument.ContentModifiedDate < 2024-01-01T00:00:00Z AND ContentDocument.ContentModifiedDate > 2023-01-01T00:00:00Z"
+                "WHERE LinkedEntityId IN (SELECT Id FROM User) "
+                "AND ContentDocumentId IN ("
+                "SELECT Id "
+                "FROM ContentDocument "
+                "WHERE ContentDocument.ContentModifiedDate < 2024-01-01T00:00:00Z "
+                "AND ContentDocument.ContentModifiedDate > 2023-01-01T00:00:00Z"
+                ")"
+            ),
+        ),
+        (
+            datetime(
+                year=2024,
+                month=1,
+                day=1,
+                hour=0,
+                minute=0,
+                second=0,
+                microsecond=0,
+                tzinfo=timezone.utc,
+            ),
+            datetime(
+                year=2023,
+                month=1,
+                day=1,
+                hour=0,
+                minute=0,
+                second=0,
+                microsecond=0,
+                tzinfo=timezone.utc,
+            ),
+            "DirField",
+            None,
+            (
+                "SELECT LinkedEntityId, ContentDocumentId, LinkedEntity.Type, DirField "
+                "FROM ContentDocumentLink "
+                "WHERE LinkedEntityId IN (SELECT Id FROM User) "
+                "AND ContentDocumentId IN ("
+                "SELECT Id "
+                "FROM ContentDocument "
+                "WHERE ContentDocument.ContentModifiedDate < 2024-01-01T00:00:00Z "
+                "AND ContentDocument.ContentModifiedDate > 2023-01-01T00:00:00Z"
                 ")"
             ),
         ),
@@ -97,11 +143,16 @@
                 tzinfo=timezone.utc,
             ),
             "DirField",
+            "MyCustomField__c = 'MyValue'",
             (
                 "SELECT LinkedEntityId, ContentDocumentId, LinkedEntity.Type, DirField "
                 "FROM ContentDocumentLink "
-                "WHERE ContentDocumentId IN ("
-                "SELECT Id FROM ContentDocument WHERE ContentDocument.ContentModifiedDate < 2024-01-01T00:00:00Z AND ContentDocument.ContentModifiedDate > 2023-01-01T00:00:00Z"
+                "WHERE LinkedEntityId IN (SELECT Id FROM User WHERE MyCustomField__c = 'MyValue') "
+                "AND ContentDocumentId IN ("
+                "SELECT Id "
+                "FROM ContentDocument "
+                "WHERE ContentDocument.ContentModifiedDate < 2024-01-01T00:00:00Z "
+                "AND ContentDocument.ContentModifiedDate > 2023-01-01T00:00:00Z"
                 ")"
             ),
         ),
@@ -111,6 +162,7 @@ def test_download_content_document_link_list_queries(
     modified_date_lt: datetime | None,
     modified_date_gt: datetime | None,
     dir_name_field: str | None,
+    extra_soql_condition: str | None,
     expected_query: str,
 ):
     client = Mock()
@@ -123,6 +175,7 @@ def test_download_content_document_link_list_queries(
             modified_date_lt=modified_date_lt,
             modified_date_gt=modified_date_gt,
             dir_name_field=dir_name_field,
+            extra_soql_condition=extra_soql_condition,
         )
         salesforce = Salesforce(archivist_obj=archivist_obj, client=client, max_api_usage_percent=50)
         salesforce.download_content_document_link_list(
@@ -136,9 +189,10 @@ def test_download_content_document_link_list_queries(
 
 
 @pytest.mark.parametrize(
-    "modified_date_lt, modified_date_gt, expected_query",
+    "modified_date_lt, modified_date_gt, extra_soql_condition, expected_query",
     [
         (
+            None,
             None,
             None,
             "SELECT Id, ParentId, BodyLength, Name FROM Attachment",
@@ -155,6 +209,7 @@ def test_download_content_document_link_list_queries(
                 tzinfo=timezone.utc,
             ),
             None,
+            None,
             (
                 "SELECT Id, ParentId, BodyLength, Name "
                 "FROM Attachment "
@@ -182,17 +237,47 @@ def test_download_content_document_link_list_queries(
                 microsecond=0,
                 tzinfo=timezone.utc,
             ),
+            None,
             (
                 "SELECT Id, ParentId, BodyLength, Name "
                 "FROM Attachment "
                 "WHERE LastModifiedDate < 2024-01-01T00:00:00Z AND LastModifiedDate > 2023-01-01T00:00:00Z"
             ),
         ),
+        (
+            datetime(
+                year=2024,
+                month=1,
+                day=1,
+                hour=0,
+                minute=0,
+                second=0,
+                microsecond=0,
+                tzinfo=timezone.utc,
+            ),
+            datetime(
+                year=2023,
+                month=1,
+                day=1,
+                hour=0,
+                minute=0,
+                second=0,
+                microsecond=0,
+                tzinfo=timezone.utc,
+            ),
+            "MyCustomField__c = 'MyValue'",
+            (
+                "SELECT Id, ParentId, BodyLength, Name "
+                "FROM Attachment "
+                "WHERE LastModifiedDate < 2024-01-01T00:00:00Z AND LastModifiedDate > 2023-01-01T00:00:00Z AND MyCustomField__c = 'MyValue'"
+            ),
+        ),
     ],
 )
 def test_download_attachment_list_queries(
     modified_date_lt: datetime | None,
     modified_date_gt: datetime | None,
+    extra_soql_condition: str | None,
     expected_query: str,
 ):
     client = Mock()
@@ -204,6 +289,7 @@ def test_download_attachment_list_queries(
             obj_type="Attachment",
             modified_date_lt=modified_date_lt,
             modified_date_gt=modified_date_gt,
+            extra_soql_condition=extra_soql_condition,
         )
         salesforce = Salesforce(archivist_obj=archivist_obj, client=client, max_api_usage_percent=50)
         salesforce.download_attachment_list(
diff --git a/test/test_archivist.py b/test/test_archivist.py
