Support artifact upload (#11419)

* wip

* reorganize

* changie

* retry

* nits

* nits

* improve retry, adjust error, adjust host name

* adjust logic

* pr_feedback

* Update .changes/unreleased/Features-20250323-151625.yaml

Co-authored-by: Emily Rockman <[email protected]>

---------

Co-authored-by: Emily Rockman <[email protected]>

Author: ChenyuLInx

.changes/unreleased/Features-20250323-151625.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Support artifact upload to dbt Cloud
+time: 2025-03-23T15:16:25.160263-07:00
+custom:
+  Author: ChenyuLInx
+  Issue: "11418"

core/dbt/cli/artifact_upload.py

@@ -0,0 +1,162 @@
+import os
+import time
+import uuid
+import zipfile
+
+import requests
+from pydantic import BaseSettings
+
+import dbt.tracking
+from dbt.config.runtime import UnsetProfile, load_project
+from dbt.constants import (
+    CATALOG_FILENAME,
+    MANIFEST_FILE_NAME,
+    RUN_RESULTS_FILE_NAME,
+    SOURCE_RESULT_FILE_NAME,
+)
+from dbt.events.types import ArtifactUploadSkipped, ArtifactUploadSuccess
+from dbt.exceptions import DbtProjectError
+from dbt_common.events.functions import fire_event
+from dbt_common.exceptions import DbtBaseException as DbtException
+
+MAX_RETRIES = 3
+
+EXECUTION_ARTIFACTS = [MANIFEST_FILE_NAME, RUN_RESULTS_FILE_NAME]
+
+ARTIFACTS_TO_UPLOAD = {
+    "retry": EXECUTION_ARTIFACTS,
+    "clone": EXECUTION_ARTIFACTS,
+    "build": EXECUTION_ARTIFACTS,
+    "run": EXECUTION_ARTIFACTS,
+    "run-operation": EXECUTION_ARTIFACTS,
+    "seed": EXECUTION_ARTIFACTS,
+    "snapshot": EXECUTION_ARTIFACTS,
+    "test": EXECUTION_ARTIFACTS,
+    "freshness": [MANIFEST_FILE_NAME, SOURCE_RESULT_FILE_NAME],
+    "generate": [MANIFEST_FILE_NAME, CATALOG_FILENAME],
+}
+
+
+class ArtifactUploadConfig(BaseSettings):
+    tenant_hostname: str
+    DBT_CLOUD_TOKEN: str
+    DBT_CLOUD_ACCOUNT_ID: str
+    DBT_CLOUD_ENVIRONMENT_ID: str
+
+    def get_ingest_url(self):
+        return f"https://{self.tenant_hostname}/api/private/accounts/{self.DBT_CLOUD_ACCOUNT_ID}/environments/{self.DBT_CLOUD_ENVIRONMENT_ID}/ingests/"
+
+    def get_complete_url(self, ingest_id):
+        return f"{self.get_ingest_url()}{ingest_id}/"
+
+    def get_headers(self, invocation_id=None):
+        if invocation_id is None:
+            invocation_id = str(uuid.uuid4())
+        return {
+            "Accept": "application/json",
+            "X-Invocation-Id": invocation_id,
+            "Authorization": f"Token {self.DBT_CLOUD_TOKEN}",
+        }
+
+
+def _retry_with_backoff(operation_name, func, max_retries=MAX_RETRIES, retry_codes=None):
+    """Execute a function with exponential backoff retry logic.
+
+    Args:
+        operation_name: Name of the operation for error messages
+        func: Function to execute that returns (success, result)
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        The result from the function if successful
+
+    Raises:
+        DbtException: If all retry attempts fail
+    """
+    if retry_codes is None:
+        retry_codes = [500, 502, 503, 504]
+    retry_delay = 1
+    for attempt in range(max_retries):
+        try:
+            success, result = func()
+            if success:
+                return result
+
+            if result.status_code not in retry_codes:
+                raise DbtException(f"Error {operation_name}: {result}")
+            if attempt == max_retries - 1:  # Last attempt
+                raise DbtException(f"Error {operation_name}: {result}")
+        except requests.RequestException as e:
+            if attempt == max_retries - 1:  # Last attempt
+                raise DbtException(f"Error {operation_name}: {str(e)}")
+
+        time.sleep(retry_delay)
+        retry_delay *= 2  # exponential backoff
+
+
+def upload_artifacts(project_dir, target_path, command):
+    # Check if there are artifacts to upload for this command
+    if command not in ARTIFACTS_TO_UPLOAD:
+        fire_event(ArtifactUploadSkipped(msg=f"No artifacts to upload for command {command}"))
+        return
+
+    # read configurations
+    try:
+        project = load_project(
+            project_dir, version_check=False, profile=UnsetProfile(), cli_vars=None
+        )
+        if not project.dbt_cloud or "tenant_hostname" not in project.dbt_cloud:
+            raise DbtProjectError("dbt_cloud.tenant_hostname not found in dbt_project.yml")
+        tenant_hostname = project.dbt_cloud["tenant_hostname"]
+        if not tenant_hostname:
+            raise DbtProjectError("dbt_cloud.tenant_hostname is empty in dbt_project.yml")
+    except Exception as e:
+        raise DbtProjectError(
+            f"Error reading dbt_cloud.tenant_hostname from dbt_project.yml: {str(e)}"
+        )
+
+    config = ArtifactUploadConfig(tenant_hostname=tenant_hostname)
+
+    if not target_path:
+        target_path = "target"
+
+    # Create zip file with artifacts
+    zip_file_name = "target.zip"
+    with zipfile.ZipFile(zip_file_name, "w") as z:
+        for artifact in ARTIFACTS_TO_UPLOAD[command]:
+            z.write(os.path.join(target_path, artifact), artifact)
+
+    # Step 1: Create ingest request with retry
+    def create_ingest():
+        response = requests.post(url=config.get_ingest_url(), headers=config.get_headers())
+        return response.status_code == 200, response
+
+    response = _retry_with_backoff("creating ingest request", create_ingest)
+    response_data = response.json()
+    ingest_id = response_data["data"]["id"]
+    upload_url = response_data["data"]["upload_url"]
+
+    # Step 2: Upload the zip file to the provided URL with retry
+    with open(zip_file_name, "rb") as f:
+        file_data = f.read()
+
+        def upload_file():
+            upload_response = requests.put(url=upload_url, data=file_data)
+            return upload_response.status_code in (200, 204), upload_response
+
+        _retry_with_backoff("uploading artifacts", upload_file)
+
+    # Step 3: Mark the ingest as successful with retry
+    def complete_ingest():
+        complete_response = requests.patch(
+            url=config.get_complete_url(ingest_id),
+            headers=config.get_headers(),
+            json={"upload_status": "SUCCESS"},
+        )
+        return complete_response.status_code == 204, complete_response
+
+    _retry_with_backoff("completing ingest", complete_ingest)
+
+    fire_event(ArtifactUploadSuccess(msg=f"command {command} completed successfully"))
+    if dbt.tracking.active_user is not None:
+        dbt.tracking.track_artifact_upload({"command": command})

core/dbt/cli/main.py

@@ -141,6 +141,7 @@ def global_flags(func):
     @p.warn_error_options
     @p.write_json
     @p.use_fast_test_edges
+    @p.upload_artifacts
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)

core/dbt/cli/params.py

@@ -761,3 +761,10 @@ def _version_callback(ctx, _param, value):
     help="Whether or not to write the manifest.json and run_results.json files to the target directory",
     default=True,
 )
+
+upload_artifacts = click.option(
+    "--upload-to-artifacts-ingest-api/--no-upload-to-artifacts-ingest-api",
+    envvar="DBT_UPLOAD_TO_ARTIFACTS_INGEST_API",
+    help="Whether or not to upload the artifacts to the dbt Cloud API",
+    default=False,
+)

core/dbt/cli/requires.py

@@ -9,6 +9,7 @@
 
 import dbt.tracking
 from dbt.adapters.factory import adapter_management, get_adapter, register_adapter
+from dbt.cli.artifact_upload import upload_artifacts
 from dbt.cli.exceptions import ExceptionExit, ResultExit
 from dbt.cli.flags import Flags
 from dbt.config import RuntimeConfig
@@ -17,6 +18,7 @@
 from dbt.context.query_header import generate_query_header_context
 from dbt.events.logging import setup_event_logger
 from dbt.events.types import (
+    ArtifactUploadError,
     CommandCompleted,
     MainEncounteredError,
     MainReportArgs,
@@ -26,7 +28,7 @@
     ResourceReport,
 )
 from dbt.exceptions import DbtProjectError, FailFastError
-from dbt.flags import get_flag_dict, set_flags
+from dbt.flags import get_flag_dict, get_flags, set_flags
 from dbt.mp_context import get_mp_context
 from dbt.parser.manifest import parse_manifest
 from dbt.plugins import set_up_plugin_manager
@@ -164,6 +166,15 @@ def wrapper(*args, **kwargs):
         finally:
             # Fire ResourceReport, but only on systems which support the resource
             # module. (Skip it on Windows).
+            try:
+                if get_flags().upload_to_artifacts_ingest_api:
+                    upload_artifacts(
+                        get_flags().project_dir, get_flags().target_path, ctx.command.name
+                    )
+
+            except Exception as e:
+                fire_event(ArtifactUploadError(msg=str(e)))
+
             if importlib.util.find_spec("resource") is not None:
                 import resource
 

core/dbt/constants.py

@@ -22,3 +22,6 @@
 MINIMUM_REQUIRED_TIME_SPINE_GRANULARITY = TimeGranularity.DAY
 PARTIAL_PARSE_FILE_NAME = "partial_parse.msgpack"
 PACKAGE_LOCK_HASH_KEY = "sha1_hash"
+RUN_RESULTS_FILE_NAME = "run_results.json"
+CATALOG_FILENAME = "catalog.json"
+SOURCE_RESULT_FILE_NAME = "sources.json"

core/dbt/contracts/state.py

@@ -5,6 +5,7 @@
 from dbt.artifacts.schemas.freshness import FreshnessExecutionResultArtifact
 from dbt.artifacts.schemas.manifest import WritableManifest
 from dbt.artifacts.schemas.run import RunResultsArtifact
+from dbt.constants import RUN_RESULTS_FILE_NAME
 from dbt.contracts.graph.manifest import Manifest
 from dbt.events.types import WarnStateTargetEqual
 from dbt_common.events.functions import fire_event
@@ -43,7 +44,7 @@ def __init__(self, state_path: Path, target_path: Path, project_root: Path) -> N
                 exc.add_filename(str(manifest_path))
                 raise
 
-        results_path = self.project_root / self.state_path / "run_results.json"
+        results_path = self.project_root / self.state_path / RUN_RESULTS_FILE_NAME
         self.results = load_result_state(results_path)
 
         sources_path = self.project_root / self.state_path / "sources.json"

core/dbt/events/core_types.proto

@@ -2231,3 +2231,34 @@ message ResourceReportMsg {
     CoreEventInfo info = 1;
     ResourceReport data = 2;
 }
+
+
+// Z061
+message ArtifactUploadError {
+    string msg = 1;
+}
+
+message ArtifactUploadErrorMsg {
+    CoreEventInfo info = 1;
+    ArtifactUploadError data = 2;
+}
+
+// Z062
+message ArtifactUploadSuccess {
+    string msg = 1;
+}
+
+message ArtifactUploadSuccessMsg {
+    CoreEventInfo info = 1;
+    ArtifactUploadSuccess data = 2;
+}
+
+// Z063
+message ArtifactUploadSkipped {
+    string msg = 1;
+}
+
+message ArtifactUploadSkippedMsg {
+    CoreEventInfo info = 1;
+    ArtifactUploadSkipped data = 2;
+}

core/dbt/events/core_types_pb2.py

@@ -2185,3 +2185,30 @@ def code(self) -> str:
 
     def message(self) -> str:
         return f"Resource report: {self.to_json()}"
+
+
+# Artifact Upload Events #
+
+
+class ArtifactUploadError(ErrorLevel):
+    def code(self) -> str:
+        return "Z061"
+
+    def message(self) -> str:
+        return f"Error uploading artifacts to artifact ingestion API: {self.msg}"
+
+
+class ArtifactUploadSuccess(InfoLevel):
+    def code(self) -> str:
+        return "Z062"
+
+    def message(self) -> str:
+        return f"Artifacts uploaded successfully to artifact ingestion API: {self.msg}"
+
+
+class ArtifactUploadSkipped(DebugLevel):
+    def code(self) -> str:
+        return "Z063"
+
+    def message(self) -> str:
+        return f"Artifacts skipped for command : {self.msg}"