Introduce dataspace awareness

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,19 @@ 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 +72,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,32 @@ 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":
+                    dataspace_value = os.environ.get(v)
+                    if dataspace_value is not None:
+                        credentials_data[k] = dataspace_value
+                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
@@ -93,5 +106,8 @@ def update_ini(self, profile: str = "default", ini_file: str = INI_FILE):
         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,

tests/test_credentials.py

@@ -18,6 +18,7 @@ def test_from_env(self):
             "client_id": "test_client_id",
             "client_secret": "test_secret",
             "login_url": "https://test.login.url",
+            "dataspace": "test_dataspace",
         }
 
         with patch.dict(
@@ -30,6 +31,32 @@ def test_from_env(self):
             assert creds.client_id == test_creds["client_id"]
             assert creds.client_secret == test_creds["client_secret"]
             assert creds.login_url == test_creds["login_url"]
+            assert creds.dataspace == test_creds["dataspace"]
+
+    def test_from_env_without_dataspace(self):
+        """Test loading credentials from environment variables without dataspace."""
+        test_creds = {
+            "username": "test_user",
+            "password": "test_pass",
+            "client_id": "test_client_id",
+            "client_secret": "test_secret",
+            "login_url": "https://test.login.url",
+        }
+
+        # Create env dict without dataspace
+        env_dict = {
+            v: test_creds[k] for k, v in ENV_CREDENTIALS.items() if k != "dataspace"
+        }
+
+        with patch.dict(os.environ, env_dict, clear=True):
+            creds = Credentials.from_env()
+
+            assert creds.username == test_creds["username"]
+            assert creds.password == test_creds["password"]
+            assert creds.client_id == test_creds["client_id"]
+            assert creds.client_secret == test_creds["client_secret"]
+            assert creds.login_url == test_creds["login_url"]
+            assert creds.dataspace is None  # Should default to None
 
     def test_from_env_missing_vars(self):
         """Test that missing environment variables raise appropriate error."""
@@ -47,13 +74,15 @@ def test_from_ini(self):
         client_id = ini_client_id
         client_secret = ini_secret
         login_url = https://ini.login.url
+        dataspace = ini_dataspace
 
         [other_profile]
         username = other_user
         password = other_pass
         client_id = other_client_id
         client_secret = other_secret
         login_url = https://other.login.url
+        dataspace = other_dataspace
         """
 
         with (
@@ -73,6 +102,7 @@ def test_from_ini(self):
                 assert creds.client_id == "ini_client_id"
                 assert creds.client_secret == "ini_secret"
                 assert creds.login_url == "https://ini.login.url"
+                assert creds.dataspace == "ini_dataspace"
 
                 # Test other profile
                 creds = Credentials.from_ini(
@@ -83,6 +113,7 @@ def test_from_ini(self):
                 assert creds.client_id == "other_client_id"
                 assert creds.client_secret == "other_secret"
                 assert creds.login_url == "https://other.login.url"
+                assert creds.dataspace == "other_dataspace"
 
     def test_from_available_env(self):
         """Test that from_available uses environment variables when available."""
@@ -92,6 +123,7 @@ def test_from_available_env(self):
             "client_id": "test_client_id",
             "client_secret": "test_secret",
             "login_url": "https://test.login.url",
+            "dataspace": "test_dataspace",
         }
 
         with (
@@ -107,6 +139,7 @@ def test_from_available_env(self):
             assert creds.client_id == test_creds["client_id"]
             assert creds.client_secret == test_creds["client_secret"]
             assert creds.login_url == test_creds["login_url"]
+            assert creds.dataspace == test_creds["dataspace"]
 
     def test_from_available_ini(self):
         """Test that from_available uses INI file when env vars not available."""
@@ -117,6 +150,7 @@ def test_from_available_ini(self):
         client_id = ini_client_id
         client_secret = ini_secret
         login_url = https://ini.login.url
+        dataspace = ini_dataspace
         """
 
         with (
@@ -137,6 +171,7 @@ def test_from_available_ini(self):
                 assert creds.client_id == "ini_client_id"
                 assert creds.client_secret == "ini_secret"
                 assert creds.login_url == "https://ini.login.url"
+                assert creds.dataspace == "ini_dataspace"
 
     def test_from_available_no_creds(self):
         """Test that from_available raises error when no credentials are found."""

tests/test_credentials_profile_integration.py

@@ -32,6 +32,7 @@ def test_query_api_reader_with_custom_profile(self):
             mock_credentials.password = "custom_password"
             mock_credentials.client_id = "custom_client_id"
             mock_credentials.client_secret = "custom_secret"
+            mock_credentials.dataspace = "custom_dataspace"
             mock_from_available.return_value = mock_credentials
 
             # Mock the SalesforceCDPConnection
@@ -49,7 +50,51 @@ def test_query_api_reader_with_custom_profile(self):
                 # Verify the correct profile was used
                 mock_from_available.assert_called_with(profile="custom_profile")
 
-                # Verify the connection was created with the custom credentials
+                # Verify the connection was created
+                # with the custom credentials including dataspace
+                mock_conn_class.assert_called_once_with(
+                    "https://custom.salesforce.com",
+                    "custom@example.com",
+                    "custom_password",
+                    "custom_client_id",
+                    "custom_secret",
+                    dataspace="custom_dataspace",
+                )
+
+    def test_query_api_reader_without_dataspace(self):
+        """Test QueryAPIDataCloudReader works when dataspace is None."""
+        mock_spark = MagicMock()
+
+        with patch(
+            "datacustomcode.credentials.Credentials.from_available"
+        ) as mock_from_available:
+            # Mock credentials without dataspace
+            mock_credentials = MagicMock()
+            mock_credentials.login_url = "https://custom.salesforce.com"
+            mock_credentials.username = "custom@example.com"
+            mock_credentials.password = "custom_password"
+            mock_credentials.client_id = "custom_client_id"
+            mock_credentials.client_secret = "custom_secret"
+            mock_credentials.dataspace = None  # No dataspace
+            mock_from_available.return_value = mock_credentials
+
+            # Mock the SalesforceCDPConnection
+            with patch(
+                "datacustomcode.io.reader.query_api.SalesforceCDPConnection"
+            ) as mock_conn_class:
+                mock_conn = MagicMock()
+                mock_conn_class.return_value = mock_conn
+
+                # Test with custom profile
+                QueryAPIDataCloudReader(
+                    mock_spark, credentials_profile="custom_profile"
+                )
+
+                # Verify the correct profile was used
+                mock_from_available.assert_called_with(profile="custom_profile")
+
+                # Verify the connection was created
+                # WITHOUT dataspace parameter when None
                 mock_conn_class.assert_called_once_with(
                     "https://custom.salesforce.com",
                     "custom@example.com",