Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 25 additions & 1 deletion .github/workflows/sf_cli_integration.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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 \
Expand All @@ -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 \
Expand Down Expand Up @@ -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 \
Expand All @@ -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 \
Expand Down
4 changes: 2 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
31 changes: 28 additions & 3 deletions scripts/mock_sf_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand All @@ -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",
}
Expand Down Expand Up @@ -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
Expand All @@ -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()
56 changes: 56 additions & 0 deletions src/datacustomcode/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from __future__ import annotations

from enum import Enum
import os
from typing import (
TYPE_CHECKING,
Any,
Expand Down Expand Up @@ -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.

Expand All @@ -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.

Expand Down Expand Up @@ -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:
Comment thread
sbyrne-sf marked this conversation as resolved.
"""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)
Comment thread
sbyrne-sf marked this conversation as resolved.

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
12 changes: 12 additions & 0 deletions src/datacustomcode/io/writer/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Comment thread
sbyrne-sf marked this conversation as resolved.
"""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:
Expand Down
8 changes: 8 additions & 0 deletions src/datacustomcode/io/writer/csv.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
sbyrne-sf marked this conversation as resolved.

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)
7 changes: 7 additions & 0 deletions src/datacustomcode/io/writer/print.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,10 @@ def write_to_dlo(

dataframe.show()

def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None:
Comment thread
sbyrne-sf marked this conversation as resolved.
self.validate_dataframe_columns_against_dlo(dataframe, name)
dataframe.show()

def write_to_dmo(
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
) -> None:
Expand All @@ -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()
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Comment thread
sbyrne-sf marked this conversation as resolved.
# 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)
Comment thread
sbyrne-sf marked this conversation as resolved.


def transform(dataframe: DataFrame) -> DataFrame:
return dataframe.withColumn("description__c", upper(col("description__c")))


if __name__ == "__main__":
Expand Down
Loading
Loading