diff --git a/README.md b/README.md index e30665d..aa7f3d2 100644 --- a/README.md +++ b/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) diff --git a/src/datacustomcode/cli.py b/src/datacustomcode/cli.py index 0876ce3..c5e195f 100644 --- a/src/datacustomcode/cli.py +++ b/src/datacustomcode/cli.py @@ -50,6 +50,11 @@ 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, @@ -57,6 +62,7 @@ def configure( 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) diff --git a/src/datacustomcode/credentials.py b/src/datacustomcode/credentials.py index d7db5e2..21b2a27 100644 --- a/src/datacustomcode/credentials.py +++ b/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) diff --git a/src/datacustomcode/io/reader/query_api.py b/src/datacustomcode/io/reader/query_api.py index b6b87fa..2a33d0f 100644 --- a/src/datacustomcode/io/reader/query_api.py +++ b/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, diff --git a/tests/test_credentials.py b/tests/test_credentials.py index e2f1bcc..70d4512 100644 --- a/tests/test_credentials.py +++ b/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,6 +74,7 @@ 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 @@ -54,6 +82,7 @@ def test_from_ini(self): 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.""" diff --git a/tests/test_credentials_profile_integration.py b/tests/test_credentials_profile_integration.py index 92a1538..862edb7 100644 --- a/tests/test_credentials_profile_integration.py +++ b/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",