Enable dataspaces for customcode run

Author: ppandit-sfdc

src/datacustomcode/io/reader/query_api.py

@@ -18,6 +18,7 @@
 from typing import (
     TYPE_CHECKING,
     Final,
+    Optional,
     Union,
 )
 
@@ -76,18 +77,31 @@ class QueryAPIDataCloudReader(BaseDataCloudReader):
     CONFIG_NAME = "QueryAPIDataCloudReader"
 
     def __init__(
-        self, spark: SparkSession, credentials_profile: str = "default"
+        self,
+        spark: SparkSession,
+        credentials_profile: str = "default",
+        dataspace: Optional[str] = None,
     ) -> None:
         self.spark = spark
         credentials = Credentials.from_available(profile=credentials_profile)
 
-        self._conn = SalesforceCDPConnection(
-            credentials.login_url,
-            credentials.username,
-            credentials.password,
-            credentials.client_id,
-            credentials.client_secret,
-        )
+        if dataspace is not None:
+            self._conn = SalesforceCDPConnection(
+                credentials.login_url,
+                credentials.username,
+                credentials.password,
+                credentials.client_id,
+                credentials.client_secret,
+                dataspace=dataspace,
+            )
+        else:
+            self._conn = SalesforceCDPConnection(
+                credentials.login_url,
+                credentials.username,
+                credentials.password,
+                credentials.client_id,
+                credentials.client_secret,
+            )
 
     def read_dlo(
         self,

src/datacustomcode/io/writer/print.py

@@ -30,13 +30,20 @@ def __init__(
         spark: SparkSession,
         reader: Optional[QueryAPIDataCloudReader] = None,
         credentials_profile: str = "default",
+        dataspace: Optional[str] = None,
     ) -> None:
         super().__init__(spark)
-        self.reader = (
-            QueryAPIDataCloudReader(self.spark, credentials_profile)
-            if reader is None
-            else reader
-        )
+        if reader is None:
+            if dataspace is not None:
+                self.reader = QueryAPIDataCloudReader(
+                    self.spark, credentials_profile, dataspace=dataspace
+                )
+            else:
+                self.reader = QueryAPIDataCloudReader(
+                    self.spark, credentials_profile
+                )
+        else:
+            self.reader = reader
 
     def validate_dataframe_columns_against_dlo(
         self,

src/datacustomcode/run.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import importlib
+import json
+import os
 from pathlib import Path
 import runpy
 import sys
@@ -36,6 +38,37 @@ def run_entrypoint(
         profile: The credentials profile to use.
     """
     add_py_folder(entrypoint)
+    
+    # Read dataspace from config.json if it exists
+    entrypoint_dir = os.path.dirname(entrypoint)
+    config_json_path = os.path.join(entrypoint_dir, "config.json")
+    if os.path.exists(config_json_path):
+        try:
+            with open(config_json_path, "r") as f:
+                config_json = json.load(f)
+                dataspace = config_json.get("dataspace")
+                if dataspace:
+                    # Add dataspace to reader config options
+                    if (
+                        config.reader_config
+                        and hasattr(config.reader_config, "options")
+                    ):
+                        config.reader_config.options["dataspace"] = dataspace
+                    # Add dataspace to writer config options (for PrintDataCloudWriter)
+                    if (
+                        config.writer_config
+                        and hasattr(config.writer_config, "options")
+                    ):
+                        config.writer_config.options["dataspace"] = dataspace
+        except json.JSONDecodeError as err:
+            raise ValueError(
+                f"config.json at {config_json_path} is not valid JSON"
+            ) from err
+        except FileNotFoundError as err:
+            raise FileNotFoundError(
+                f"config.json not found at {config_json_path}"
+            ) from err
+    
     if profile != "default":
         if config.reader_config and hasattr(config.reader_config, "options"):
             config.reader_config.options["credentials_profile"] = profile