salesforce-data-customcode 3.0.2__tar.gz → 4.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (94) hide show
  1. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/PKG-INFO +2 -1
  2. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/pyproject.toml +2 -1
  3. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/__init__.py +0 -4
  4. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/auth.py +1 -1
  5. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/cli.py +102 -32
  6. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/client.py +0 -8
  7. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/cmd.py +1 -1
  8. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/config.py +7 -22
  9. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/config.yaml +6 -5
  10. salesforce_data_customcode-4.0.0/src/datacustomcode/constants.py +45 -0
  11. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/deploy.py +36 -88
  12. salesforce_data_customcode-4.0.0/src/datacustomcode/einstein_platform_client.py +89 -0
  13. salesforce_data_customcode-4.0.0/src/datacustomcode/einstein_platform_config.py +41 -0
  14. salesforce_data_customcode-4.0.0/src/datacustomcode/einstein_predictions/impl/default.py +79 -0
  15. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/einstein_predictions_config.py +4 -12
  16. salesforce_data_customcode-4.0.0/src/datacustomcode/function/feature_types/chunking.py +197 -0
  17. salesforce_data_customcode-4.0.0/src/datacustomcode/function_utils.py +367 -0
  18. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/reader/sf_cli.py +3 -58
  19. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway/default.py +19 -7
  20. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway_config.py +4 -12
  21. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/run.py +90 -10
  22. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/scan.py +2 -0
  23. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/template.py +30 -1
  24. salesforce_data_customcode-4.0.0/src/datacustomcode/templates/function/chunking/payload/entrypoint.py +145 -0
  25. salesforce_data_customcode-4.0.0/src/datacustomcode/templates/function/chunking/requirements.txt +1 -0
  26. salesforce_data_customcode-4.0.0/src/datacustomcode/templates/function/payload/config.json +1 -0
  27. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/function/payload/entrypoint.py +25 -13
  28. salesforce_data_customcode-4.0.0/src/datacustomcode/templates/script/README.md +0 -0
  29. salesforce_data_customcode-4.0.0/src/datacustomcode/token_provider.py +149 -0
  30. salesforce_data_customcode-3.0.2/src/datacustomcode/einstein_predictions/impl/default.py +0 -35
  31. salesforce_data_customcode-3.0.2/src/datacustomcode/function/feature_types/chunking.py +0 -89
  32. salesforce_data_customcode-3.0.2/src/datacustomcode/proxy/client/LocalProxyClientProvider.py +0 -34
  33. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/LICENSE.txt +0 -0
  34. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/README.md +0 -0
  35. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/common_config.py +0 -0
  36. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/credentials.py +0 -0
  37. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/einstein_predictions/__init__.py +1 -1
  38. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/einstein_predictions/base.py +0 -0
  39. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/einstein_predictions/types.py +0 -0
  40. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/file/__init__.py +0 -0
  41. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/file/base.py +0 -0
  42. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/file/path/__init__.py +0 -0
  43. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/file/path/default.py +0 -0
  44. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/function/__init__.py +0 -0
  45. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/function/base.py +0 -0
  46. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/function/feature_types/__init__.py +0 -0
  47. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/function/runtime.py +0 -0
  48. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/__init__.py +0 -0
  49. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/base.py +0 -0
  50. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/reader/__init__.py +0 -0
  51. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/reader/base.py +0 -0
  52. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/reader/query_api.py +0 -0
  53. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/reader/utils.py +0 -0
  54. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/writer/__init__.py +0 -0
  55. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/writer/base.py +0 -0
  56. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/writer/csv.py +0 -0
  57. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/io/writer/print.py +0 -0
  58. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway/__init__.py +0 -0
  59. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway/base.py +0 -0
  60. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway/types/__init__.py +0 -0
  61. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway/types/generate_text_request.py +0 -0
  62. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway/types/generate_text_request_builder.py +0 -0
  63. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway/types/generate_text_response.py +0 -0
  64. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/llm_gateway/types/generate_text_response_builder.py +0 -0
  65. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/mixin.py +0 -0
  66. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/proxy/__init__.py +0 -0
  67. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/proxy/base.py +0 -0
  68. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/proxy/client/__init__.py +0 -0
  69. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/proxy/client/base.py +0 -0
  70. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/py.typed +0 -0
  71. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/spark/__init__.py +0 -0
  72. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/spark/base.py +0 -0
  73. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/spark/default.py +0 -0
  74. /salesforce_data_customcode-3.0.2/src/datacustomcode/templates/function/README.md → /salesforce_data_customcode-4.0.0/src/datacustomcode/templates/__init__.py +0 -0
  75. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/function/.devcontainer/devcontainer.json +0 -0
  76. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/function/Dockerfile.dependencies +0 -0
  77. {salesforce_data_customcode-3.0.2/src/datacustomcode/templates/script → salesforce_data_customcode-4.0.0/src/datacustomcode/templates/function}/README.md +0 -0
  78. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/function/build_native_dependencies.sh +0 -0
  79. {salesforce_data_customcode-3.0.2/src/datacustomcode/templates/function → salesforce_data_customcode-4.0.0/src/datacustomcode/templates/function/chunking}/payload/config.json +0 -0
  80. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/function/requirements-dev.txt +0 -0
  81. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/function/requirements.txt +0 -0
  82. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/.devcontainer/devcontainer.json +0 -0
  83. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/Dockerfile +0 -0
  84. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/Dockerfile.dependencies +0 -0
  85. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/account.ipynb +0 -0
  86. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/build_native_dependencies.sh +0 -0
  87. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/examples/employee_hierarchy/employee_data.csv +0 -0
  88. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/examples/employee_hierarchy/entrypoint.py +0 -0
  89. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/jupyterlab.sh +0 -0
  90. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/payload/config.json +0 -0
  91. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/payload/entrypoint.py +0 -0
  92. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/requirements-dev.txt +0 -0
  93. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/templates/script/requirements.txt +0 -0
  94. {salesforce_data_customcode-3.0.2 → salesforce_data_customcode-4.0.0}/src/datacustomcode/version.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: salesforce-data-customcode
3
- Version: 3.0.2
3
+ Version: 4.0.0
4
4
  Summary: Data Cloud Custom Code SDK
5
5
  License-Expression: Apache-2.0
6
6
  License-File: LICENSE.txt
@@ -18,6 +18,7 @@ Requires-Dist: pandas
18
18
  Requires-Dist: pydantic (==2.13.1)
19
19
  Requires-Dist: pyspark (==3.5.1)
20
20
  Requires-Dist: pyyaml (>=6.0,<7.0)
21
+ Requires-Dist: requests (==2.33.1)
21
22
  Requires-Dist: salesforce-cdp-connector (>=1.0.19)
22
23
  Requires-Dist: setuptools_scm (>=7.1.0,<8.0.0)
23
24
  Description-Content-Type: text/markdown
@@ -18,7 +18,7 @@ license = "Apache-2.0"
18
18
  name = "salesforce-data-customcode"
19
19
  readme = "README.md"
20
20
  requires-python = ">=3.10,<3.12"
21
- version = "3.0.2"
21
+ version = "4.0.0"
22
22
 
23
23
  [tool.black]
24
24
  exclude = '''
@@ -105,6 +105,7 @@ pydantic = "2.13.1"
105
105
  pyspark = "3.5.1"
106
106
  python = ">=3.10,<3.12"
107
107
  pyyaml = "^6.0"
108
+ requests = "2.33.1"
108
109
  salesforce-cdp-connector = ">=1.0.19"
109
110
  setuptools_scm = "^7.1.0"
110
111
 
@@ -17,15 +17,11 @@ from datacustomcode.client import Client
17
17
  from datacustomcode.credentials import AuthType, Credentials
18
18
  from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
19
19
  from datacustomcode.io.writer.print import PrintDataCloudWriter
20
- from datacustomcode.proxy.client.LocalProxyClientProvider import (
21
- LocalProxyClientProvider,
22
- )
23
20
 
24
21
  __all__ = [
25
22
  "AuthType",
26
23
  "Client",
27
24
  "Credentials",
28
- "LocalProxyClientProvider",
29
25
  "PrintDataCloudWriter",
30
26
  "QueryAPIDataCloudReader",
31
27
  ]
@@ -170,7 +170,7 @@ def do_oauth_browser_flow(
170
170
 
171
171
  # Start callback server
172
172
  click.echo(f"\nStarting local callback server on {redirect_uri}...")
173
- server, actual_port = _run_oauth_callback_server(redirect_uri, auth_code_queue)
173
+ server, _actual_port = _run_oauth_callback_server(redirect_uri, auth_code_queue)
174
174
 
175
175
  # Build authorization URL with final redirect_uri
176
176
  auth_url = (
@@ -27,6 +27,13 @@ from loguru import logger
27
27
 
28
28
  from datacustomcode import AuthType
29
29
  from datacustomcode.auth import configure_oauth_tokens
30
+ from datacustomcode.constants import (
31
+ CONFIG_FILE,
32
+ ENTRYPOINT_FILE,
33
+ PAYLOAD_DIR,
34
+ TEST_FILE,
35
+ TESTS_DIR,
36
+ )
30
37
  from datacustomcode.scan import find_base_directory, get_package_type
31
38
 
32
39
 
@@ -74,6 +81,30 @@ def _configure_client_credentials(
74
81
  )
75
82
 
76
83
 
84
+ def _generate_function_test_file(entrypoint_path: str) -> Optional[str]:
85
+ """Generate test.json file for a function.
86
+
87
+ Args:
88
+ entrypoint_path: Path to the function's entrypoint.py
89
+
90
+ Returns:
91
+ Path to generated test.json, or None if generation failed
92
+ """
93
+ from datacustomcode.function_utils import generate_test_json
94
+
95
+ tests_dir = os.path.join(os.path.dirname(entrypoint_path), TESTS_DIR)
96
+ os.makedirs(tests_dir, exist_ok=True)
97
+ test_json_path = os.path.join(tests_dir, TEST_FILE)
98
+
99
+ try:
100
+ generate_test_json(entrypoint_path, test_json_path)
101
+ logger.debug(f"Generated test JSON at {test_json_path}")
102
+ return test_json_path
103
+ except Exception as e:
104
+ logger.warning(f"Could not generate test.json: {e}")
105
+ return None
106
+
107
+
77
108
  @cli.command()
78
109
  @click.option("--profile", default="default", help="Credential profile name")
79
110
  @click.option(
@@ -162,7 +193,6 @@ def zip(path: str, network: str):
162
193
 
163
194
  Choose based on your workload requirements.""",
164
195
  )
165
- @click.option("--function-invoke-opt")
166
196
  @click.option(
167
197
  "--sf-cli-org",
168
198
  default=None,
@@ -176,16 +206,18 @@ def deploy(
176
206
  cpu_size: str,
177
207
  profile: str,
178
208
  network: str,
179
- function_invoke_opt: str,
180
209
  sf_cli_org: Optional[str],
181
210
  ):
182
- from datacustomcode.credentials import Credentials
211
+ from datacustomcode.constants import USE_IN_FEATURE_MAPPING_FOR_CONNECT_API
183
212
  from datacustomcode.deploy import (
184
213
  COMPUTE_TYPES,
185
- AccessTokenResponse,
186
214
  CodeExtensionMetadata,
187
- _retrieve_access_token_from_sf_cli,
188
215
  deploy_full,
216
+ infer_use_in_feature,
217
+ )
218
+ from datacustomcode.token_provider import (
219
+ CredentialsTokenProvider,
220
+ SFCLITokenProvider,
189
221
  )
190
222
 
191
223
  logger.debug("Deploying project")
@@ -210,32 +242,34 @@ def deploy(
210
242
  )
211
243
 
212
244
  if package_type == "function":
213
- if not function_invoke_opt:
245
+ # Infer use_in_feature from function signature
246
+ entrypoint_path = os.path.join(path, ENTRYPOINT_FILE)
247
+ use_in_feature = infer_use_in_feature(entrypoint_path)
248
+ if use_in_feature:
249
+ logger.info(f"Inferred use_in_feature: {use_in_feature}")
250
+ else:
214
251
  click.secho(
215
- "Error: Function invoke options are required for function package type",
252
+ "Error: Could not infer function invoke options. "
253
+ "Please provide --use-in-feature",
216
254
  fg="red",
217
255
  )
218
256
  raise click.Abort()
257
+
258
+ # Map user-provided feature names to API names
259
+ mapped_feature = USE_IN_FEATURE_MAPPING_FOR_CONNECT_API.get(
260
+ use_in_feature, use_in_feature
261
+ )
262
+ metadata.functionInvokeOptions = [mapped_feature]
263
+
264
+ try:
265
+ if sf_cli_org:
266
+ auth = SFCLITokenProvider(sf_cli_org).get_token()
219
267
  else:
220
- function_invoke_options = function_invoke_opt.split(",")
221
- metadata.functionInvokeOptions = function_invoke_options
222
-
223
- auth: Union[Credentials, AccessTokenResponse]
224
- if sf_cli_org:
225
- try:
226
- auth = _retrieve_access_token_from_sf_cli(sf_cli_org)
227
- except RuntimeError as e:
228
- click.secho(f"Error: {e}", fg="red")
229
- raise click.Abort() from None
230
- else:
231
- try:
232
- auth = Credentials.from_available(profile=profile)
233
- except ValueError as e:
234
- click.secho(
235
- f"Error: {e}",
236
- fg="red",
237
- )
238
- raise click.Abort() from None
268
+ auth = CredentialsTokenProvider(profile).get_token()
269
+ except RuntimeError as e:
270
+ click.secho(f"Error: {e}", fg="red")
271
+ raise click.Abort() from None
272
+
239
273
  deploy_full(path, metadata, auth, network)
240
274
 
241
275
 
@@ -244,7 +278,12 @@ def deploy(
244
278
  @click.option(
245
279
  "--code-type", default="script", type=click.Choice(["script", "function"])
246
280
  )
247
- def init(directory: str, code_type: str):
281
+ @click.option(
282
+ "--use-in-feature",
283
+ default="SearchIndexChunking",
284
+ help="Feature where this function will be used (only applicable for function).",
285
+ )
286
+ def init(directory: str, code_type: str, use_in_feature: Optional[str]):
248
287
  from datacustomcode.scan import (
249
288
  dc_config_json_from_file,
250
289
  update_config,
@@ -256,9 +295,9 @@ def init(directory: str, code_type: str):
256
295
  if code_type == "script":
257
296
  copy_script_template(directory)
258
297
  elif code_type == "function":
259
- copy_function_template(directory)
260
- entrypoint_path = os.path.join(directory, "payload", "entrypoint.py")
261
- config_location = os.path.join(os.path.dirname(entrypoint_path), "config.json")
298
+ copy_function_template(directory, use_in_feature)
299
+ entrypoint_path = os.path.join(directory, PAYLOAD_DIR, ENTRYPOINT_FILE)
300
+ config_location = os.path.join(os.path.dirname(entrypoint_path), CONFIG_FILE)
262
301
 
263
302
  # Write package type to SDK-specific config
264
303
  sdk_config = {"type": code_type}
@@ -271,6 +310,7 @@ def init(directory: str, code_type: str):
271
310
  updated_config_json = update_config(entrypoint_path)
272
311
  with open(config_location, "w") as f:
273
312
  json.dump(updated_config_json, f, indent=2)
313
+
274
314
  click.echo(
275
315
  "Start developing by updating the code in "
276
316
  + click.style(entrypoint_path, fg="blue", bold=True)
@@ -281,6 +321,24 @@ def init(directory: str, code_type: str):
281
321
  + " to automatically update config.json when you make changes to your code"
282
322
  )
283
323
 
324
+ # Generate test.json for functions
325
+ if code_type == "function":
326
+ test_json_path = _generate_function_test_file(entrypoint_path)
327
+ if test_json_path:
328
+ click.echo(
329
+ "Generated test file at "
330
+ + click.style(test_json_path, fg="blue", bold=True)
331
+ )
332
+ click.echo(
333
+ "Test your function locally with "
334
+ + click.style(
335
+ f"datacustomcode run {entrypoint_path} "
336
+ f"--test-with {test_json_path}",
337
+ fg="blue",
338
+ bold=True,
339
+ )
340
+ )
341
+
284
342
 
285
343
  @cli.command()
286
344
  @click.argument("filename")
@@ -292,7 +350,7 @@ def init(directory: str, code_type: str):
292
350
  def scan(filename: str, config: str, dry_run: bool, no_requirements: bool):
293
351
  from datacustomcode.scan import update_config, write_requirements_file
294
352
 
295
- config_location = config or os.path.join(os.path.dirname(filename), "config.json")
353
+ config_location = config or os.path.join(os.path.dirname(filename), CONFIG_FILE)
296
354
  click.echo(
297
355
  "Dumping scan results to config file: "
298
356
  + click.style(config_location, fg="blue", bold=True)
@@ -318,6 +376,12 @@ def scan(filename: str, config: str, dry_run: bool, no_requirements: bool):
318
376
  @click.option("--config-file", default=None)
319
377
  @click.option("--dependencies", default=[], multiple=True)
320
378
  @click.option("--profile", default="default")
379
+ @click.option(
380
+ "--test-with",
381
+ default=None,
382
+ type=click.Path(exists=True),
383
+ help="Path to test JSON file for function testing",
384
+ )
321
385
  @click.option(
322
386
  "--sf-cli-org",
323
387
  default=None,
@@ -328,10 +392,16 @@ def run(
328
392
  config_file: Union[str, None],
329
393
  dependencies: List[str],
330
394
  profile: str,
395
+ test_with: Optional[str],
331
396
  sf_cli_org: Optional[str],
332
397
  ):
333
398
  from datacustomcode.run import run_entrypoint
334
399
 
335
400
  run_entrypoint(
336
- entrypoint, config_file, dependencies, profile, sf_cli_org=sf_cli_org
401
+ entrypoint,
402
+ config_file,
403
+ dependencies,
404
+ profile,
405
+ test_file=test_with,
406
+ sf_cli_org=sf_cli_org,
337
407
  )
@@ -33,7 +33,6 @@ if TYPE_CHECKING:
33
33
 
34
34
  from datacustomcode.io.reader.base import BaseDataCloudReader
35
35
  from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
36
- from datacustomcode.proxy.client.base import BaseProxyClient
37
36
  from datacustomcode.spark.base import BaseSparkSessionProvider
38
37
 
39
38
 
@@ -107,7 +106,6 @@ class Client:
107
106
  _reader: BaseDataCloudReader
108
107
  _writer: BaseDataCloudWriter
109
108
  _file: DefaultFindFilePath
110
- _proxy: Optional[BaseProxyClient]
111
109
  _data_layer_history: dict[DataCloudObjectType, set[str]]
112
110
  _code_type: str
113
111
 
@@ -115,7 +113,6 @@ class Client:
115
113
  cls,
116
114
  reader: Optional[BaseDataCloudReader] = None,
117
115
  writer: Optional["BaseDataCloudWriter"] = None,
118
- proxy: Optional[BaseProxyClient] = None,
119
116
  spark_provider: Optional["BaseSparkSessionProvider"] = None,
120
117
  code_type: str = "script",
121
118
  ) -> Client:
@@ -223,11 +220,6 @@ class Client:
223
220
  self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
224
221
  return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
225
222
 
226
- def call_llm_gateway(self, LLM_MODEL_ID: str, prompt: str, maxTokens: int) -> str:
227
- if self._proxy is None:
228
- raise ValueError("No proxy configured; set proxy or proxy_config")
229
- return self._proxy.call_llm_gateway(LLM_MODEL_ID, prompt, maxTokens) # type: ignore[no-any-return]
230
-
231
223
  def find_file_path(self, file_name: str) -> Path:
232
224
  """Return a file path"""
233
225
 
@@ -104,6 +104,6 @@ def _cmd_output(
104
104
 
105
105
 
106
106
  def cmd_output(*cmd: str, **kwargs: Any) -> Union[str, None]:
107
- returncode, stdout_b, stderr_b = _cmd_output(*cmd, **kwargs)
107
+ _returncode, stdout_b, _stderr_b = _cmd_output(*cmd, **kwargs)
108
108
  stdout = stdout_b.decode() if stdout_b is not None else None
109
109
  return stdout
@@ -37,15 +37,14 @@ from datacustomcode.common_config import (
37
37
  # This lets all readers and writers to be findable via config
38
38
  from datacustomcode.io import * # noqa: F403
39
39
  from datacustomcode.io.base import BaseDataAccessLayer
40
- from datacustomcode.io.reader.base import BaseDataCloudReader # noqa: TCH002
41
- from datacustomcode.io.writer.base import BaseDataCloudWriter # noqa: TCH002
42
- from datacustomcode.proxy.base import BaseProxyAccessLayer
43
- from datacustomcode.proxy.client.base import BaseProxyClient # noqa: TCH002
44
40
  from datacustomcode.spark.base import BaseSparkSessionProvider
45
41
 
46
42
  if TYPE_CHECKING:
47
43
  from pyspark.sql import SparkSession
48
44
 
45
+ from datacustomcode.io.reader.base import BaseDataCloudReader
46
+ from datacustomcode.io.writer.base import BaseDataCloudWriter
47
+
49
48
 
50
49
  _T = TypeVar("_T", bound="BaseDataAccessLayer")
51
50
 
@@ -55,7 +54,7 @@ class AccessLayerObjectConfig(BaseObjectConfig, Generic[_T]):
55
54
 
56
55
  def to_object(self, spark: SparkSession) -> _T:
57
56
  type_ = self.type_base.subclass_from_config_name(self.type_config_name)
58
- return cast(_T, type_(spark=spark, **self.options))
57
+ return cast("_T", type_(spark=spark, **self.options))
59
58
 
60
59
 
61
60
  class SparkConfig(ForceableConfig):
@@ -74,31 +73,18 @@ class SparkConfig(ForceableConfig):
74
73
 
75
74
  _P = TypeVar("_P", bound=BaseSparkSessionProvider)
76
75
 
77
- _PX = TypeVar("_PX", bound=BaseProxyAccessLayer)
78
-
79
-
80
- class ProxyAccessLayerObjectConfig(BaseObjectConfig, Generic[_PX]):
81
- """Config for proxy clients that take no constructor args (e.g. no spark)."""
82
-
83
- type_base: ClassVar[Type[BaseProxyAccessLayer]] = BaseProxyAccessLayer
84
-
85
- def to_object(self) -> _PX:
86
- type_ = self.type_base.subclass_from_config_name(self.type_config_name)
87
- return cast(_PX, type_(**self.options))
88
-
89
76
 
90
77
  class SparkProviderConfig(BaseObjectConfig, Generic[_P]):
91
78
  type_base: ClassVar[Type[BaseSparkSessionProvider]] = BaseSparkSessionProvider
92
79
 
93
80
  def to_object(self) -> _P:
94
81
  type_ = self.type_base.subclass_from_config_name(self.type_config_name)
95
- return cast(_P, type_(**self.options))
82
+ return cast("_P", type_(**self.options))
96
83
 
97
84
 
98
85
  class ClientConfig(BaseConfig):
99
- reader_config: Union[AccessLayerObjectConfig[BaseDataCloudReader], None] = None
100
- writer_config: Union[AccessLayerObjectConfig[BaseDataCloudWriter], None] = None
101
- proxy_config: Union[ProxyAccessLayerObjectConfig[BaseProxyClient], None] = None
86
+ reader_config: Union[AccessLayerObjectConfig["BaseDataCloudReader"], None] = None
87
+ writer_config: Union[AccessLayerObjectConfig["BaseDataCloudWriter"], None] = None
102
88
  spark_config: Union[SparkConfig, None] = None
103
89
  spark_provider_config: Union[
104
90
  SparkProviderConfig[BaseSparkSessionProvider], None
@@ -126,7 +112,6 @@ class ClientConfig(BaseConfig):
126
112
 
127
113
  self.reader_config = merge(self.reader_config, other.reader_config)
128
114
  self.writer_config = merge(self.writer_config, other.writer_config)
129
- self.proxy_config = merge(self.proxy_config, other.proxy_config)
130
115
  self.spark_config = merge(self.spark_config, other.spark_config)
131
116
  self.spark_provider_config = merge(
132
117
  self.spark_provider_config, other.spark_provider_config
@@ -19,11 +19,12 @@ spark_config:
19
19
  spark.sql.execution.arrow.pyspark.enabled: 'true'
20
20
  spark.driver.extraJavaOptions: -Djava.security.manager=allow
21
21
 
22
- proxy_config:
23
- type_config_name: LocalProxyClientProvider
22
+ einstein_predictions_config:
23
+ type_config_name: DefaultEinsteinPredictions
24
24
  options:
25
25
  credentials_profile: default
26
26
 
27
- einstein_predictions_config:
28
- type_config_name: DefaultEinsteinPredictions
29
- options: {}
27
+ llm_gateway_config:
28
+ type_config_name: DefaultLLMGateway
29
+ options:
30
+ credentials_profile: default
@@ -0,0 +1,45 @@
1
+ # Copyright (c) 2025, Salesforce, Inc.
2
+ # SPDX-License-Identifier: Apache-2
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ """Constants used throughout the datacustomcode package."""
17
+
18
+ # File and directory names
19
+ ENTRYPOINT_FILE = "entrypoint.py"
20
+ CONFIG_FILE = "config.json"
21
+ PAYLOAD_DIR = "payload"
22
+ TESTS_DIR = "tests"
23
+ TEST_FILE = "test.json"
24
+ REQUIREMENTS_FILE = "requirements.txt"
25
+
26
+ # Default values
27
+ DEFAULT_PROFILE = "default"
28
+ DEFAULT_NETWORK = "default"
29
+ DEFAULT_CPU_SIZE = "CPU_2XL"
30
+
31
+ # Feature to template folder mapping
32
+ FEATURE_TEMPLATE_MAPPING = {
33
+ "SearchIndexChunking": "chunking",
34
+ }
35
+
36
+ # Feature name to Connect API name mapping
37
+ USE_IN_FEATURE_MAPPING_FOR_CONNECT_API = {
38
+ "SearchIndexChunking": "UnstructuredChunking",
39
+ }
40
+
41
+ # Pydantic request/response type names to feature names
42
+ REQUEST_TYPE_TO_FEATURE = {
43
+ "SearchIndexChunkingV1Request": "SearchIndexChunking",
44
+ "SearchIndexChunkingV1Response": "SearchIndexChunking",
45
+ }
@@ -19,11 +19,9 @@ import json
19
19
  import os
20
20
  import re
21
21
  import shutil
22
- import subprocess
23
22
  import tempfile
24
23
  import time
25
24
  from typing import (
26
- TYPE_CHECKING,
27
25
  Any,
28
26
  Callable,
29
27
  Dict,
@@ -37,15 +35,11 @@ from pydantic import BaseModel
37
35
  import requests
38
36
 
39
37
  from datacustomcode.cmd import cmd_output
40
- from datacustomcode.credentials import AuthType
38
+ from datacustomcode.constants import REQUEST_TYPE_TO_FEATURE
41
39
  from datacustomcode.scan import find_base_directory, get_package_type
42
40
 
43
- if TYPE_CHECKING:
44
- from datacustomcode.credentials import Credentials
45
-
46
41
  DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code"
47
42
  DATA_TRANSFORMS_PATH = "services/data/v63.0/ssot/data-transforms"
48
- AUTH_PATH = "services/oauth2/token"
49
43
  WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000
50
44
 
51
45
  # Available compute types for Data Cloud deployments.
@@ -72,6 +66,40 @@ def _sanitize_api_name(name: str) -> str:
72
66
  return sanitized
73
67
 
74
68
 
69
+ def infer_use_in_feature(entrypoint_path: str) -> Union[str, None]:
70
+ """Infer the use_in_feature from function signature.
71
+
72
+ Checks both the request parameter type and return type annotation.
73
+ Both must map to the same feature for a valid inference.
74
+
75
+ Uses static AST parsing to avoid importing dependencies.
76
+
77
+ Args:
78
+ entrypoint_path: Path to the entrypoint.py file
79
+
80
+ Returns:
81
+ The feature name if both request and response match, None otherwise
82
+ """
83
+ from datacustomcode.function_utils import inspect_function_types_static
84
+
85
+ request_type_name, response_type_name = inspect_function_types_static(
86
+ entrypoint_path
87
+ )
88
+
89
+ if not request_type_name or not response_type_name:
90
+ return None
91
+
92
+ # Look up features for both types
93
+ request_feature = REQUEST_TYPE_TO_FEATURE.get(request_type_name)
94
+ response_feature = REQUEST_TYPE_TO_FEATURE.get(response_type_name)
95
+
96
+ # Both must be present and must match
97
+ if request_feature and response_feature and request_feature == response_feature:
98
+ return request_feature
99
+
100
+ return None
101
+
102
+
75
103
  class CodeExtensionMetadata(BaseModel):
76
104
  name: str
77
105
  version: str
@@ -163,80 +191,6 @@ class AccessTokenResponse(BaseModel):
163
191
  instance_url: str
164
192
 
165
193
 
166
- def _retrieve_access_token(credentials: Credentials) -> AccessTokenResponse:
167
- """Get an access token for the Salesforce API."""
168
- logger.debug("Getting oauth token...")
169
-
170
- url = f"{credentials.login_url.rstrip('/')}/{AUTH_PATH.lstrip('/')}"
171
-
172
- if credentials.auth_type == AuthType.OAUTH_TOKENS:
173
- data = {
174
- "grant_type": "refresh_token",
175
- "refresh_token": credentials.refresh_token,
176
- "client_id": credentials.client_id,
177
- "client_secret": credentials.client_secret,
178
- }
179
- elif credentials.auth_type == AuthType.CLIENT_CREDENTIALS:
180
- data = {
181
- "grant_type": "client_credentials",
182
- "client_id": credentials.client_id,
183
- "client_secret": credentials.client_secret,
184
- }
185
- else:
186
- raise ValueError(f"Unsupported auth_type: {credentials.auth_type}")
187
-
188
- response = _make_api_call(url, "POST", data=data)
189
- return AccessTokenResponse(**response)
190
-
191
-
192
- def _retrieve_access_token_from_sf_cli(sf_cli_org: str) -> AccessTokenResponse:
193
- """Get an access token from the Salesforce CLI."""
194
- try:
195
- result = subprocess.run(
196
- ["sf", "org", "display", "--target-org", sf_cli_org, "--json"],
197
- capture_output=True,
198
- text=True,
199
- check=True,
200
- timeout=30,
201
- )
202
- except FileNotFoundError as exc:
203
- raise RuntimeError(
204
- "The 'sf' command was not found. "
205
- "Please install Salesforce CLI: https://developer.salesforce.com/tools/salesforcecli"
206
- ) from exc
207
- except subprocess.TimeoutExpired as exc:
208
- raise RuntimeError(
209
- f"'sf org display' timed out for org '{sf_cli_org}'"
210
- ) from exc
211
- except subprocess.CalledProcessError as exc:
212
- raise RuntimeError(
213
- f"'sf org display' failed for org '{sf_cli_org}'.\n"
214
- f"Ensure the org is authenticated via 'sf org login web'.\n"
215
- f"stderr: {exc.stderr.strip()}"
216
- ) from exc
217
-
218
- try:
219
- data = json.loads(result.stdout)
220
- except json.JSONDecodeError as exc:
221
- raise RuntimeError(f"Failed to parse 'sf org display' output: {exc}") from exc
222
-
223
- if data.get("status") != 0:
224
- raise RuntimeError(
225
- f"SF CLI error for org '{sf_cli_org}': "
226
- f"{data.get('message', 'unknown error')}"
227
- )
228
-
229
- org_result = data.get("result", {})
230
- access_token = org_result.get("accessToken")
231
- instance_url = org_result.get("instanceUrl")
232
- if not access_token or not instance_url:
233
- raise RuntimeError(
234
- f"'sf org display' did not return an access token or instance URL "
235
- f"for org '{sf_cli_org}'"
236
- )
237
- return AccessTokenResponse(access_token=access_token, instance_url=instance_url)
238
-
239
-
240
194
  class CreateDeploymentResponse(BaseModel):
241
195
  fileUploadUrl: str
242
196
 
@@ -567,16 +521,11 @@ def zip(
567
521
  def deploy_full(
568
522
  directory: str,
569
523
  metadata: CodeExtensionMetadata,
570
- credentials: Union["Credentials", AccessTokenResponse],
524
+ access_token: AccessTokenResponse,
571
525
  docker_network: str,
572
526
  callback=None,
573
527
  ) -> AccessTokenResponse:
574
528
  """Deploy a data transform in the DataCloud."""
575
- if isinstance(credentials, AccessTokenResponse):
576
- access_token = credentials
577
- else:
578
- access_token = _retrieve_access_token(credentials)
579
-
580
529
  # prepare payload
581
530
  config = get_config(directory)
582
531
 
@@ -587,7 +536,6 @@ def deploy_full(
587
536
  wait_for_deployment(access_token, metadata, callback)
588
537
 
589
538
  # create data transform
590
-
591
539
  if isinstance(config, DataTransformConfig):
592
540
  create_data_transform(directory, access_token, metadata, config)
593
541
  return access_token