feat: support for profile run parameters

README.md

@@ -98,14 +98,15 @@ Mandatatory parameters for RudderstackProfilesOperator:
 * profile_id: This is the [profiles id](https://www.rudderstack.com/docs/api/profiles-api/#run-project) for the profiles project to run.
 * connection_id: The Airflow connection to use for connecting to the Rudderstack API.	Default value is `rudderstack_default`.
 
-RudderstackRETLOperator exposes other configurable parameters as well. Mostly default values for them would be recommended.
+RudderstackProfilesOperator exposes other configurable parameters as well. Mostly default values for them would be recommended.
 
 * request_max_retries: The maximum number of times requests to the RudderStack API should be retried before failng.
 * request_retry_delay: Time (in seconds) to wait between each request retry.
 * request_timeout: Time (in seconds) after which the requests to RudderStack are declared timed out.
 * poll_interval: Time (in seconds) for polling status of triggered job.
 * poll_timeout: Time (in seconds) after which the polling for a triggered job is declared timed out.
 * wait_for_completion: Boolean if execution run should poll and wait till completion of sync. Default value is True.
+* parameters: Additional parameters to pass to the profiles run command, as supported by the API endpoint. Default value is `None`.
 
 
 ## Contribute

rudder_airflow_provider/hooks/rudderstack.py

@@ -2,7 +2,7 @@
 import time
 import datetime
 from importlib.metadata import PackageNotFoundError, version
-from typing import Any, Dict, Mapping, Optional
+from typing import Any, Dict, Mapping, Optional, List
 from urllib.parse import urljoin
 from airflow.exceptions import AirflowException
 from airflow.providers.http.hooks.http import HttpHook
@@ -247,7 +247,7 @@ def __init__(
         self.poll_timeout = poll_timeout
         self.poll_interval = poll_interval
 
-    def start_profile_run(self, profile_id: str):
+    def start_profile_run(self, profile_id: str, parameters: Optional[List[str]] = None) -> str:
         """Triggers a profile run and returns runId if successful, else raises Failure.
 
         Args:
@@ -258,6 +258,7 @@ def start_profile_run(self, profile_id: str):
         self.log.info(f"Triggering profile run for profile id: {profile_id}")
         return self.make_request(
             endpoint=f"/v2/sources/{profile_id}/start",
+            data={"parameters": parameters} if parameters else None,
         )["runId"]
 
     def poll_profile_run(self, profile_id: str, run_id: str) -> Dict[str, Any]:

rudder_airflow_provider/operators/rudderstack.py

@@ -1,6 +1,6 @@
 import logging
 from airflow.models import baseoperator
-from typing import Optional
+from typing import Optional, List
 from rudder_airflow_provider.hooks.rudderstack import (
     RudderStackRETLHook,
     RudderStackProfilesHook,
@@ -69,7 +69,7 @@ def execute(self, context):
 
 
 class RudderstackProfilesOperator(baseoperator.BaseOperator):
-    template_fields = "profile_id"
+    template_fields = ("profile_id", "parameters")
 
     """
         Rudderstack operator for Profiles
@@ -87,6 +87,7 @@ def __init__(
         request_max_retries: int = DEFAULT_REQUEST_MAX_RETRIES,
         poll_timeout: float = None,
         poll_interval: float = DEFAULT_POLL_INTERVAL_SECONDS,
+        parameters: Optional[List[str]] = None,
         **kwargs,
     ):
         """
@@ -101,6 +102,7 @@ def __init__(
         self.request_max_retries = request_max_retries
         self.poll_timeout = poll_timeout
         self.poll_interval = poll_interval
+        self.parameters = parameters
 
     def execute(self, context):
         rs_profiles_hook = RudderStackProfilesHook(
@@ -111,7 +113,7 @@ def execute(self, context):
             poll_timeout=self.poll_timeout,
             poll_interval=self.poll_interval,
         )
-        profile_run_id = rs_profiles_hook.start_profile_run(self.profile_id)
+        profile_run_id = rs_profiles_hook.start_profile_run(self.profile_id, self.parameters)
         if self.wait_for_completion:
             self.log.info(
                 f"Poll and wait for profiles run to finish for profilesId: {self.profile_id}, runId: {profile_run_id}"

rudder_airflow_provider/test/hooks/test_rudderstack_hook.py

@@ -204,6 +204,24 @@ def test_start_profile_run(mock_request, mock_connection, airflow_connection):
         timeout=30,
     )
 
+# RudderStackProfilesHook tests
+@patch("airflow.providers.http.hooks.http.HttpHook.get_connection")
+@patch("requests.request")
+def test_start_profile_run_parameters(mock_request, mock_connection, airflow_connection):
+    mock_request.return_value = MagicMock(
+        status_code=200, json=lambda: {"runId": TEST_PROFILE_RUN_ID}
+    )
+    mock_connection.return_value = airflow_connection
+    profiles_hook = RudderStackProfilesHook(TEST_AIRFLOW_CONN_ID)
+    run_id = profiles_hook.start_profile_run(TEST_PROFILE_ID, ["--rebase_incremental"])
+    assert run_id == TEST_PROFILE_RUN_ID
+    mock_request.assert_called_once_with(
+        method="POST",
+        url=f"{TEST_BASE_URL}/v2/sources/{TEST_PROFILE_ID}/start",
+        json={'parameters': ['--rebase_incremental']},
+        headers=ANY,
+        timeout=30,
+    )
 
 @patch("airflow.providers.http.hooks.http.HttpHook.get_connection")
 def test_start_sprofile_run_invalid_parameters(mock_connection, airflow_connection):

rudder_airflow_provider/test/operators/test_rudderstack_operator.py

@@ -58,7 +58,7 @@ def test_profiles_operator_execute_without_wait(mock_profile_run, mock_poll_prof
                                        wait_for_completion=False,
                                        task_id='some-task-id')
     profiles_operator.execute(context=None)
-    mock_profile_run.assert_called_once_with(TEST_PROFILE_ID)
+    mock_profile_run.assert_called_once_with(TEST_PROFILE_ID, None)
     mock_poll_profile_run.assert_not_called()
 
 
@@ -77,7 +77,7 @@ def test_profiles_operator_execute_with_wait(mock_profile_run, mock_poll_profile
                                        task_id='some-task-id')
     profiles_operator.execute(context=None)
 
-    mock_profile_run.assert_called_once_with(TEST_PROFILE_ID)
+    mock_profile_run.assert_called_once_with(TEST_PROFILE_ID, None)
     mock_poll_profile_run.assert_called_once_with(TEST_PROFILE_ID, TEST_PROFILES_RUN_ID)
 
 if __name__ == "__main__":