Introduce dataspace awareness

Author: ppandit-sfdc

README.md

@@ -180,6 +180,7 @@ Options:
 - `--client-id TEXT`: Connected App Client ID
 - `--client-secret TEXT`: Connected App Client Secret
 - `--login-url TEXT`: Salesforce login URL
+- `--dataspace TEXT`: Dataspace name (optional, for non-default dataspaces)
 
 
 #### `datacustomcode init`
@@ -308,6 +309,16 @@ You can read more about Jupyter Notebooks here: https://jupyter.org/
 
 You now have all fields necessary for the `datacustomcode configure` command.
 
+### Working with Dataspaces
+
+If you're working with a non-default dataspace in Salesforce Data Cloud, you can specify the dataspace during configuration:
+
+```bash
+datacustomcode configure --dataspace my-dataspace
+```
+
+**For default dataspaces**, you can omit the `--dataspace` parameter entirely - the SDK will connect to the default dataspace automatically.
+
 ## Other docs
 
 - [Troubleshooting](./docs/troubleshooting.md)

src/datacustomcode/cli.py

@@ -50,13 +50,15 @@ def version():
 @click.option("--client-id", prompt=True)
 @click.option("--client-secret", prompt=True)
 @click.option("--login-url", prompt=True)
+@click.option("--dataspace", default="default", help="Dataspace name (optional, for non-default dataspaces)")
 def configure(
     username: str,
     password: str,
     client_id: str,
     client_secret: str,
     login_url: str,
     profile: str,
+    dataspace: str | None,
 ) -> None:
     from datacustomcode.credentials import Credentials
 
@@ -66,6 +68,7 @@ def configure(
         client_id=client_id,
         client_secret=client_secret,
         login_url=login_url,
+        dataspace=dataspace,
     ).update_ini(profile=profile)
 
 

src/datacustomcode/credentials.py

@@ -26,6 +26,7 @@
     "client_id": "SFDC_CLIENT_ID",
     "client_secret": "SFDC_CLIENT_SECRET",
     "login_url": "SFDC_LOGIN_URL",
+    "dataspace": "SFDC_DATASPACE",
 }
 INI_FILE = os.path.expanduser("~/.datacustomcode/credentials.ini")
 
@@ -37,6 +38,7 @@ class Credentials:
     client_id: str
     client_secret: str
     login_url: str
+    dataspace: str | None = None
 
     @classmethod
     def from_ini(
@@ -47,21 +49,30 @@ def from_ini(
         config = configparser.ConfigParser()
         logger.debug(f"Reading {ini_file} for profile {profile}")
         config.read(ini_file)
+        dataspace = config[profile].get("dataspace")
         return cls(
             username=config[profile]["username"],
             password=config[profile]["password"],
             client_id=config[profile]["client_id"],
             client_secret=config[profile]["client_secret"],
             login_url=config[profile]["login_url"],
+            dataspace=dataspace,
         )
 
     @classmethod
     def from_env(cls) -> Credentials:
         try:
-            return cls(**{k: os.environ[v] for k, v in ENV_CREDENTIALS.items()})
+            credentials_data = {}
+            for k, v in ENV_CREDENTIALS.items():
+                if k == "dataspace":
+                    credentials_data[k] = os.environ.get(v)
+                else:
+                    credentials_data[k] = os.environ[v]
+            return cls(**credentials_data)
         except KeyError as exc:
+            required_vars = [v for k, v in ENV_CREDENTIALS.items() if k != "dataspace"]
             raise ValueError(
-                f"All of {ENV_CREDENTIALS.values()} must be set in environment."
+                f"All of {required_vars} must be set in environment. "
             ) from exc
 
     @classmethod
@@ -92,6 +103,9 @@ def update_ini(self, profile: str = "default", ini_file: str = INI_FILE):
         config[profile]["client_id"] = self.client_id
         config[profile]["client_secret"] = self.client_secret
         config[profile]["login_url"] = self.login_url
+        
+        if self.dataspace is not None:
+            config[profile]["dataspace"] = self.dataspace
 
         with open(expanded_ini_file, "w") as f:
             config.write(f)

src/datacustomcode/io/reader/query_api.py

@@ -81,13 +81,19 @@ def __init__(
         self.spark = spark
         credentials = Credentials.from_available(profile=credentials_profile)
 
-        self._conn = SalesforceCDPConnection(
+        connection_args = [
             credentials.login_url,
             credentials.username,
             credentials.password,
             credentials.client_id,
             credentials.client_secret,
-        )
+        ]
+        
+        connection_kwargs = {}
+        if credentials.dataspace is not None:
+            connection_kwargs["dataspace"] = credentials.dataspace
+
+        self._conn = SalesforceCDPConnection(*connection_args, **connection_kwargs)
 
     def read_dlo(
         self,