Merge branch 'main' into ilongin/12010-concurrent-insert-buffer

Author: ilongin

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ ci:
   skip: [mypy]
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: check-added-large-files
         exclude: '^tests/examples/data/'
@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.12.7'
+    rev: 'v0.12.8'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

docs/commands/job/run.md

@@ -32,6 +32,7 @@ This command runs a job in Studio using the specified query file. You can config
 * `--env-file ENV_FILE` - File with environment variables for the job
 * `--env ENV` - Environment variables in KEY=VALUE format
 * `--cluster CLUSTER` - Compute cluster to run the job on
+* `--credentials-name CREDENTIALS_NAME` - Name of the credentials to use for the job
 * `--workers WORKERS` - Number of workers for the job
 * `--files FILES` - Additional files to include in the job
 * `--python-version PYTHON_VERSION` - Python version for the job (e.g., 3.9, 3.10, 3.11)
@@ -97,7 +98,12 @@ datachain job clusters
 datachain job run --cluster 1 query.py
 ```
 
-9. Schedule a job to run once at a specific time
+9. Run a job with specific credentials
+```bash
+datachain job run --credentials-name my-aws-credentials query.py
+```
+
+10. Schedule a job to run once at a specific time
 ```bash
 # Run job tomorrow at 3pm
 datachain job run --start-time "tomorrow 3pm" query.py
@@ -112,7 +118,7 @@ datachain job run --start-time "monday 9am" query.py
 datachain job run --start-time "2024-01-15 14:30:00" query.py
 ```
 
-10. Schedule a recurring job using cron expression
+11. Schedule a recurring job using cron expression
 ```bash
 # Run job daily at midnight
 datachain job run --cron "0 0 * * *" query.py
@@ -127,13 +133,13 @@ datachain job run --cron "0 * * * *" query.py
 datachain job run --cron "@monthly" query.py
 ```
 
-11. Schedule a recurring job with a start time
+12. Schedule a recurring job with a start time
 ```bash
 # Start the cron job after tomorrow 3pm
 datachain job run --start-time "tomorrow 3pm" --cron "0 0 * * *" query.py
 ```
 
-12. Start the job and do not wait for the job to complete
+13. Start the job and do not wait for the job to complete
 ```bash
 # Do not follow or tail the logs from Studio.
 datachain job run query.py --no-wait

src/datachain/cli/__init__.py

@@ -67,7 +67,7 @@ def main(argv: Optional[list[str]] = None) -> int:
     except BrokenPipeError as exc:
         error, return_code = handle_broken_pipe_error(exc)
         return return_code
-    except (KeyboardInterrupt, Exception) as exc:
+    except (KeyboardInterrupt, Exception) as exc:  # noqa: BLE001
         error, return_code = handle_general_exception(exc, args, logging_level)
         return return_code
     finally:

src/datachain/cli/parser/job.py

@@ -63,6 +63,13 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         help="Compute cluster to run the job on",
     )
 
+    studio_run_parser.add_argument(
+        "-c",
+        "--credentials-name",
+        action="store",
+        help="Name of the credentials to use for the job",
+    )
+
     studio_run_parser.add_argument(
         "--workers",
         type=int,

src/datachain/remote/studio.py

@@ -431,6 +431,7 @@ def create_job(
         cluster: Optional[str] = None,
         start_time: Optional[str] = None,
         cron: Optional[str] = None,
+        credentials_name: Optional[str] = None,
     ) -> Response[JobData]:
         data = {
             "query": query,
@@ -446,6 +447,7 @@ def create_job(
             "compute_cluster_name": cluster,
             "start_after": start_time,
             "cron_expression": cron,
+            "credentials_name": credentials_name,
         }
         return self._send_request("datachain/job", data)
 

src/datachain/studio.py

@@ -50,6 +50,7 @@ def process_jobs_args(args: "Namespace"):
             args.start_time,
             args.cron,
             args.no_wait,
+            args.credentials_name,
         )
 
     if args.cmd == "cancel":
@@ -356,6 +357,7 @@ def create_job(
     start_time: Optional[str] = None,
     cron: Optional[str] = None,
     no_wait: Optional[bool] = False,
+    credentials_name: Optional[str] = None,
 ):
     query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
     with open(query_file) as f:
@@ -393,6 +395,7 @@ def create_job(
         cluster=cluster,
         start_time=parsed_start_time,
         cron=cron,
+        credentials_name=credentials_name,
     )
     if not response.ok:
         raise DataChainError(response.message)

tests/test_cli_studio.py

@@ -401,6 +401,8 @@ def test_studio_run(capsys, mocker, tmp_dir):
                     "https://github.com/iterative/datachain",
                     "--cluster",
                     "default",
+                    "--credentials-name",
+                    "my-credentials",
                 ]
             )
             == 0
@@ -442,6 +444,7 @@ def test_studio_run(capsys, mocker, tmp_dir):
         "compute_cluster_name": "default",
         "start_after": None,
         "cron_expression": None,
+        "credentials_name": "my-credentials",
     }