Skip to content

spark_session

Celonis Spark Session Builder.

Provides SparkSession creation connected to Storium Spark Proxy.

GRPC_URL_ENV module-attribute 

GRPC_URL_ENV = 'GRPC_URL'

INVALID_GRPC_URL_MESSAGE module-attribute 

INVALID_GRPC_URL_MESSAGE = (
    f"{GRPC_URL_ENV} must contain a valid Spark Connect URL"
)

reset_spark_session 

reset_spark_session(*, stop=False)

Clear the context-local active Spark session and its connection metadata.

Parameters:

  • stop (bool, default: False ) –

    Also stop the active session. Defaults to clearing only the reference.

Sessions created inside a decorated execution remain owned by that execution and are stopped on exit, even after a reset. New threads require an explicit session; work inheriting a context must finish before the owning execution exits. Reset after manually stopping a session. If stopping fails, the reference is still cleared, but scope ownership is retained for cleanup to retry on exit; the stopping exception is propagated to the caller.

get_spark_session 

get_spark_session(
    schema=None,
    celonis_url=None,
    client_id=None,
    client_secret=None,
    scope=None,
    *,
    reuse=True,
)

Return the active SparkSession, or create one connected to Storium Spark Proxy.

Authentication

If client_id and client_secret are both provided, they are used to obtain an OAuth2 token directly. Otherwise a dummy token is used - this is the expected path when running as a triggered Python Task, where the platform exchanges the token for the correctly scoped one at execution time.

In ML Workbench, Spark Connect traffic is routed through the local gRPC sidecar URL supplied in the GRPC_URL environment variable. Outside ML Workbench, the Spark endpoint falls back to port 443 on the resolved Celonis host.

Callers may supply only schema and rely on this environment-based URL and authentication flow. Explicit OAuth credentials are also supported.

Parameters:

  • schema (Optional[str], default: None ) –

    Schema ID to use for creation. May be omitted when reusing an active session.

  • celonis_url (Optional[str], default: None ) –

    Celonis cluster URL (e.g., 'https://my-team.eu-1.celonis.cloud'). Falls back to CELONIS_URL env var.

  • client_id (Optional[str], default: None ) –

    Optional Celonis OAuth2 client ID. Authentication falls back to a dummy token that is replaced when executing from ML Workbench or as a Python Task.

  • client_secret (Optional[str], default: None ) –

    Optional Celonis OAuth2 client secret. Authentication falls back to a dummy token that is replaced when executing from ML Workbench or as a Python Task.

  • scope (Optional[str], default: None ) –

    OAuth scope to request. Defaults to 'spark_connect'.

  • reuse (bool, default: True ) –

    Return the active context-local session if available. Omitted arguments inherit its configuration; conflicting explicit values raise ValueError. Defaults to True. Pass False to create and register a fresh session. Reuse performs no network health check. Outside a decorated execution, callers must stop sessions they deliberately replace; replacing the active reference does not stop the previous session.

Returns:

  • SparkSession –

    SparkSession connected to Storium Spark Proxy

Raises:

  • ValueError –

    If creation requires a missing URL/schema, or explicit arguments conflict with the active session.

  • ImportError –

    If PySpark is not installed.

Examples:

Using explicit credentials: 

spark = get_spark_session(
    celonis_url="https://my-team.eu-1.celonis.cloud",
    schema="my-schema-id",
    client_id="my-client-id",
    client_secret="my-client-secret",
)

Using the CELONIS_URL environment variable (ML Workbench or external): 

# With CELONIS_URL set
spark = get_spark_session(
    schema="my-schema-id"
)