Merge pull request #137 from forcedotcom/noWriteModeV2

Added auto_write* methods to writers and clients.

Author: sbyrne-sf

.github/workflows/sf_cli_integration.yml

@@ -17,6 +17,18 @@ jobs:
     - name: Checkout code
       uses: actions/checkout@v4
 
+    - name: Set mock server TLS cert paths
+      run: |
+        echo "MOCK_SF_CERT_FILE=$RUNNER_TEMP/mock_sf_cert.pem" >> "$GITHUB_ENV"
+        echo "MOCK_SF_KEY_FILE=$RUNNER_TEMP/mock_sf_key.pem" >> "$GITHUB_ENV"
+
+    - name: Generate mock server TLS cert
+      run: |
+        openssl req -x509 -newkey rsa:2048 -nodes \
+          -keyout "$MOCK_SF_KEY_FILE" -out "$MOCK_SF_CERT_FILE" \
+          -days 1 -subj "/CN=localhost" \
+          -addext "subjectAltName=DNS:localhost,IP:127.0.0.1"
+
     - name: Set up Python 3.11
       uses: actions/setup-python@v5
       with:
@@ -73,7 +85,7 @@ jobs:
         sfdx_dir.mkdir(exist_ok=True)
         auth = {
             "accessToken": "00D000000000001AAA!fakeTokenForCITesting",
-            "instanceUrl": "http://localhost:8888",
+            "instanceUrl": "https://localhost:8888",
             "loginUrl": "https://login.salesforce.com",
             "orgId": "00D000000000001AAA",
             "userId": "005000000000001AAA",
@@ -164,6 +176,9 @@ jobs:
       # ── Script: run ───────────────────────────────────────────────────────────
 
     - name: '[script] run — sf data-code-extension script run --entrypoint testScript/payload/entrypoint.py -o dev1'
+      env:
+        NODE_EXTRA_CA_CERTS: ${{ env.MOCK_SF_CERT_FILE }}
+        REQUESTS_CA_BUNDLE: ${{ env.MOCK_SF_CERT_FILE }}
       run: |
         sf data-code-extension script run \
           --entrypoint testScript/payload/entrypoint.py \
@@ -175,6 +190,9 @@ jobs:
       # ── Script: deploy ───────────────────────────────────────────────────────
 
     - name: '[script] deploy — sf data-code-extension script deploy'
+      env:
+        NODE_EXTRA_CA_CERTS: ${{ env.MOCK_SF_CERT_FILE }}
+        REQUESTS_CA_BUNDLE: ${{ env.MOCK_SF_CERT_FILE }}
       run: |
         sf data-code-extension script deploy \
           --name test-script-deploy \
@@ -262,6 +280,9 @@ jobs:
       # ── Function: run ─────────────────────────────────────────────────────────
 
     - name: '[function] run — sf data-code-extension function run --entrypoint testFunction/payload/entrypoint.py  --test-with testFunction/payload/tests/test.json -o dev1'
+      env:
+        NODE_EXTRA_CA_CERTS: ${{ env.MOCK_SF_CERT_FILE }}
+        REQUESTS_CA_BUNDLE: ${{ env.MOCK_SF_CERT_FILE }}
       run: |
         sf data-code-extension function run \
           --entrypoint testFunction/payload/entrypoint.py \
@@ -273,6 +294,9 @@ jobs:
       # ── Function: deploy ─────────────────────────────────────────────────────
 
     - name: '[function] deploy — sf data-code-extension function deploy'
+      env:
+        NODE_EXTRA_CA_CERTS: ${{ env.MOCK_SF_CERT_FILE }}
+        REQUESTS_CA_BUNDLE: ${{ env.MOCK_SF_CERT_FILE }}
       run: |
         sf data-code-extension function deploy \
           --name test-function-deploy \

README.md

@@ -509,7 +509,7 @@ exploration.  Instead of running an entire script, one can run one code cell at
 
 You can read more about Jupyter Notebooks here: https://jupyter.org/
 
-1. Within the root project of your package folder, run `./jupyterlab.sh start`
+1. Within the root project of your package folder, run `./jupyterlab.sh start`. This prints an access token and opens an already-authenticated JupyterLab session in your browser. If the browser doesn't open automatically, copy the printed `http://localhost:8888/?token=...` URL into your browser.
 1. Double-click on "account.ipynb" file, which provides a starting point for a notebook
 1. Use shift+enter to execute each cell within the notebook.  Add/edit/delete cells of code as needed for your data exploration.
 1. Don't forget to run `./jupyterlab.sh stop` to stop the docker container
@@ -600,4 +600,4 @@ If you're using OAuth Tokens authentication, the initial configure will retrieve
 ## Other docs
 
 - [Troubleshooting](./docs/troubleshooting.md)
-- [For Contributors](./FOR_CONTRIBUTORS.md)
+- [Contributing](./CONTRIBUTING.md)

scripts/mock_sf_server.py

@@ -40,13 +40,25 @@
     python scripts/mock_sf_server.py          # listens on port 8888
     MOCK_SF_PORT=9000 python scripts/mock_sf_server.py
     python scripts/mock_sf_server.py 9000
+
+Serves TLS (the deploy path requires an HTTPS upload URL) using a pre-generated
+cert/key pair — this script does not generate one. Set ``MOCK_SF_CERT_FILE`` /
+``MOCK_SF_KEY_FILE`` to the pair's paths; generate a throwaway one with:
+
+    openssl req -x509 -newkey rsa:2048 -nodes -keyout key.pem -out cert.pem \\
+      -days 1 -subj "/CN=localhost" \\
+      -addext "subjectAltName=DNS:localhost,IP:127.0.0.1"
+
+Point clients at the cert so they trust it: ``NODE_EXTRA_CA_CERTS`` (CLI) and
+``REQUESTS_CA_BUNDLE`` (SDK).
 """
 
 from __future__ import annotations
 
 from http.server import BaseHTTPRequestHandler, HTTPServer
 import json
 import os
+import ssl
 import sys
 
 PORT = (
@@ -68,7 +80,7 @@
 
 _TOKEN_RESPONSE = {
     "access_token": "00D000000000001AAA!fakeAccessTokenForCITesting",
-    "instance_url": f"http://localhost:{PORT}",
+    "instance_url": f"https://localhost:{PORT}",
     "token_type": "Bearer",
     "scope": "api",
 }
@@ -135,7 +147,9 @@ def do_POST(self) -> None:
         elif path == _DATA_CUSTOM_CODE_PATH:
             # create_deployment() — return a presigned upload URL
             self._send_json(
-                {"fileUploadUrl": f"http://localhost:{PORT}/upload/fake-deployment.zip"}
+                {
+                    "fileUploadUrl": f"https://localhost:{PORT}/upload/fake-deployment.zip"
+                }
             )
         elif path == _DATA_TRANSFORMS_PATH:
             # create_data_transform() — script packages only
@@ -150,7 +164,18 @@ def do_PUT(self) -> None:
 
 
 if __name__ == "__main__":
+    cert_path = os.environ.get("MOCK_SF_CERT_FILE")
+    key_path = os.environ.get("MOCK_SF_KEY_FILE")
+    if not cert_path or not key_path:
+        sys.exit(
+            "MOCK_SF_CERT_FILE and MOCK_SF_KEY_FILE must both be set to an "
+            "existing TLS cert/key pair — see the module docstring."
+        )
+
     server = HTTPServer(("localhost", PORT), MockSFHandler)
     server.allow_reuse_address = True
-    print(f"[MOCK SF] Listening on http://localhost:{PORT}", flush=True)
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+    ctx.load_cert_chain(cert_path, key_path)
+    server.socket = ctx.wrap_socket(server.socket, server_side=True)
+    print(f"[MOCK SF] Listening on https://localhost:{PORT}", flush=True)
     server.serve_forever()

src/datacustomcode/client.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from enum import Enum
+import os
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -558,6 +559,15 @@ class StreamingClient(_BaseClient):
 
     _instance: ClassVar[Optional[StreamingClient]] = None
 
+    def read_dlo(self) -> PySparkDataFrame:
+        """Read the streamingSource
+
+        Returns:
+            A standard PySpark DataFrame from the streaming source DLO
+        """
+        self._record_dlo_access(_streaming_source_name())
+        return self._reader.read_dlo(_streaming_source_name())
+
     def read_dlo_deltas(self) -> PySparkDataFrame:
         """Read the streaming change feed (deltas) for a DLO from Data Cloud.
 
@@ -571,6 +581,14 @@ def read_dlo_deltas(self) -> PySparkDataFrame:
         self._record_dlo_access(_streaming_source_name())
         return self._reader.read_dlo_deltas()  # type: ignore[no-any-return]
 
+    def read_dmo(self) -> PySparkDataFrame:
+        """Read the streamingSource
+
+        Returns a standard PySpark DataFrame from the streaming source DMO
+        """
+        self._record_dmo_access(_streaming_source_name())
+        return self._reader.read_dmo(_streaming_source_name())
+
     def read_dmo_deltas(self) -> PySparkDataFrame:
         """Read the streaming change feed (deltas) for a DMO from Data Cloud.
 
@@ -598,3 +616,41 @@ def write_dlo_deltas(
         """
         self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
         return self._writer.write_dlo_deltas(name, dataframe, **kwargs)  # type: ignore[no-any-return]
+
+    def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None:
+        """Write a PySpark DataFrame to a DLO in Data Cloud automatically picking
+        the WriteMode.
+        For use with streaming transforms when running in rebuild or initial sync mode.
+        Args:
+            name: The name of the DLO to write to.
+            dataframe: The PySpark DataFrame to write.
+        """
+        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
+        return self._writer.auto_write_to_dlo(name, dataframe)
+
+    def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None:
+        """Write a PySpark DataFrame to a DMO in Data Cloud automatically picking
+        the WriteMode.
+        For use with streaming transforms when running in rebuild or initial sync mode.
+        Args:
+            name: The name of the DMO to write to.
+            dataframe: The PySpark DataFrame to write.
+        """
+        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
+        return self._writer.auto_write_to_dmo(name, dataframe)
+
+
+class RunMode(Enum):
+    BATCH = "BATCH"
+    INITIAL_SYNC = "INITIAL_SYNC"
+    REBUILD = "REBUILD"
+    DELTA_SYNC = "DELTA_SYNC"
+
+
+def get_run_mode() -> RunMode:
+    """Read and validate the BYOC_RUN_MODE env var; default to BATCH when unset."""
+    run_mode = os.getenv("BYOC_RUN_MODE", "BATCH").upper()
+    try:
+        return RunMode(run_mode)
+    except ValueError as exc:
+        raise ValueError("Set BYOC_RUN_MODE to a valid value") from exc

src/datacustomcode/io/writer/base.py

@@ -59,6 +59,18 @@ def write_to_dmo(
         self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
     ) -> None: ...
 
+    def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None:
+        """Write to a DLO automatically picking the write mode.
+        For use with streaming transforms when running in rebuild or initial sync mode.
+        """
+        raise NotImplementedError
+
+    def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None:
+        """Write to a DMO automatically picking the write mode.
+        For use with streaming transforms when running in rebuild or initial sync mode.
+        """
+        raise NotImplementedError
+
     def write_dlo_deltas(
         self, name: str, dataframe: PySparkDataFrame
     ) -> StreamingQuery:

src/datacustomcode/io/writer/csv.py

@@ -36,10 +36,18 @@ def write_to_dlo(
             name = f"{name}{SUFFIX}"
         dataframe.write.csv(name, mode=write_mode)
 
+    def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None:
+        # use overwrite since this is a local only writer
+        self.write_to_dlo(name, dataframe, WriteMode.OVERWRITE)
+
     def write_to_dmo(
         self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
     ) -> None:
         # Only add the suffix if it's not already there
         if not name.lower().endswith(SUFFIX):
             name = f"{name}{SUFFIX}"
         dataframe.write.csv(name, mode=write_mode)
+
+    def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None:
+        # use overwrite since this is a local only writer
+        self.write_to_dmo(name, dataframe, WriteMode.OVERWRITE)

src/datacustomcode/io/writer/print.py

@@ -122,6 +122,10 @@ def write_to_dlo(
 
         dataframe.show()
 
+    def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None:
+        self.validate_dataframe_columns_against_dlo(dataframe, name)
+        dataframe.show()
+
     def write_to_dmo(
         self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
     ) -> None:
@@ -130,3 +134,6 @@ def write_to_dmo(
         # so just show the dataframe.
 
         dataframe.show()
+
+    def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None:
+        dataframe.show()

src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py

@@ -3,46 +3,59 @@
 This example is the streaming counterpart to a normal batch entrypoint. Instead
 of a batch ``Client`` with ``read_dlo`` / ``write_to_dlo`` (which read and write
 a bounded snapshot), it uses a :class:`StreamingClient` and its streaming delta
-methods:
+methods.
 
-* ``client.read_dlo_deltas()`` returns a *streaming* DataFrame over the
-  Change Data Feed of the source DLO. Each row carries the source columns plus
-  change-feed metadata columns (``_record_type``, ``_commit_*``).
-* ``client.write_dlo_deltas(name, df)`` starts a streaming query that writes
-  each micro-batch to the target DLO and returns the ``StreamingQuery`` handle.
-  The runtime owns the trigger, and checkpoint location — the caller only
-  chooses the table.
+The first run of a streaming job will use the run mode INITIAL_SYNC which behaves
+like a batch run on the streaming source. A streaming transform can also use run
+mode REBUILD to do the same thing on demand. Note that these will process all
+source rows and overwrite the target.
 
 The transform in between is ordinary PySpark. Because the source is a change
 feed, keep the metadata columns on the DataFrame you hand to
 ``write_dlo_deltas`` — the sink relies on them to merge changes correctly.
 
-This entrypoint only runs inside the Data Cloud streaming (``DELTA_SYNC``)
-runtime; the local ``datacustomcode run`` readers/writers raise
+This entrypoint only runs inside the Data Cloud runtime;
+ the local ``datacustomcode run`` readers/writers raise
 ``NotImplementedError`` for the delta methods.
 """
 
+from pyspark.sql import DataFrame
 from pyspark.sql.functions import col, upper
 
-from datacustomcode.client import StreamingClient
+from datacustomcode.client import (
+    RunMode,
+    StreamingClient,
+    get_run_mode,
+)
 
 
 def main():
+    target_dlo = "Account_std_copy__dll"
     client = StreamingClient()
 
-    # Streaming DataFrame over the source DLO's change feed.
-    deltas = client.read_dlo_deltas()
-
-    # Ordinary PySpark transform.
-    transformed = deltas.withColumn("description__c", upper(col("description__c")))
-
-    # Start the streaming write. write_dlo_deltas returns the StreamingQuery;
-    # the trigger and checkpoint location are provided by the runtime.
-    query = client.write_dlo_deltas("Account_std_copy__dll", transformed)
-
-    # Drive the query's lifecycle. In the streaming runtime this blocks until
-    # the job is stopped by the platform.
-    query.awaitTermination()
+    if get_run_mode() == RunMode.DELTA_SYNC:
+        # Streaming DataFrame over the source DLO's change feed.
+        dataframe = client.read_dlo_deltas()
+        # Ordinary PySpark transform.
+        transformed = transform(dataframe)
+
+        # Start the streaming write. write_dlo_deltas returns the StreamingQuery;
+        # the trigger and checkpoint location are provided by the runtime.
+        query = client.write_dlo_deltas(target_dlo, transformed)
+
+        # Drive the query's lifecycle. In the streaming runtime this blocks until
+        # the job is stopped by the platform.
+        query.awaitTermination()
+    else:
+        # initial sync and rebuild read the entire streaming source DLO and
+        # write using a server-decided mode based on the run mode
+        dataframe = client.read_dlo()
+        transformed = transform(dataframe)
+        client.auto_write_to_dlo(target_dlo, transformed)
+
+
+def transform(dataframe: DataFrame) -> DataFrame:
+    return dataframe.withColumn("description__c", upper(col("description__c")))
 
 
 if __name__ == "__main__":