From 235ebf640080bc784002e42b856c49af4db406c0 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Wed, 28 Feb 2024 13:29:03 -0500
Subject: [PATCH] fix(dbt Cloud): allow passing job ID (#264)

---
 .../cli/superset/sync/dbt/command.py           | 18 +++++++++++++-----
 tests/cli/superset/sync/dbt/command_test.py    | 18 ++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/preset_cli/cli/superset/sync/dbt/command.py b/src/preset_cli/cli/superset/sync/dbt/command.py
index 9ee93305..f9b7a1e8 100644
--- a/src/preset_cli/cli/superset/sync/dbt/command.py
+++ b/src/preset_cli/cli/superset/sync/dbt/command.py
@@ -249,10 +249,14 @@ def get_account_id(client: DBTClient) -> int:
         click.echo(click.style("No accounts available", fg="bright_red"))
         sys.exit(1)
     if len(accounts) == 1:
-        return accounts[0]["id"]
+        account = accounts[0]
+        click.echo(
+            f'Using account {account["name"]} [id={account["id"]}] since it\'s the only one',
+        )
+        return account["id"]
     click.echo("Choose an account:")
     for i, account in enumerate(accounts):
-        click.echo(f'({i+1}) {account["name"]}')
+        click.echo(f'({i+1}) {account["name"]} [id={account["id"]}]')
 
     while True:
         try:
@@ -279,7 +283,7 @@ def get_project_id(client: DBTClient, account_id: Optional[int] = None) -> int:
         return projects[0]["id"]
     click.echo("Choose a project:")
     for i, project in enumerate(projects):
-        click.echo(f'({i+1}) {project["name"]}')
+        click.echo(f'({i+1}) {project["name"]} [id={project["id"]}]')
 
     while True:
         try:
@@ -316,7 +320,7 @@ def get_job(
 
         click.echo("Choose a job:")
         for i, job in enumerate(jobs):
-            click.echo(f'({i+1}) {job["name"]}')
+            click.echo(f'({i+1}) {job["name"]} [id={job["id"]}]')
 
         while True:
             try:
@@ -373,6 +377,8 @@ def process_sl_metrics(
 
 @click.command()
 @click.argument("token")
+@click.argument("account_id", type=click.INT, required=False, default=None)
+@click.argument("project_id", type=click.INT, required=False, default=None)
 @click.argument("job_id", type=click.INT, required=False, default=None)
 @click.option(
     "--disallow-edits",
@@ -433,6 +439,8 @@ def dbt_cloud(  # pylint: disable=too-many-arguments, too-many-locals
     select: Tuple[str, ...],
     exclude: Tuple[str, ...],
     exposures: Optional[str] = None,
+    account_id: Optional[int] = None,
+    project_id: Optional[int] = None,
     job_id: Optional[int] = None,
     disallow_edits: bool = False,
     external_url_prefix: str = "",
@@ -468,7 +476,7 @@ def dbt_cloud(  # pylint: disable=too-many-arguments, too-many-locals
     preserve_metadata = preserve_columns if preserve_columns else preserve_metadata
 
     try:
-        job = get_job(dbt_client, job_id=job_id)
+        job = get_job(dbt_client, account_id, project_id, job_id)
     except ValueError:
         click.echo(click.style(f"Job {job_id} not available", fg="bright_red"))
         sys.exit(2)
diff --git a/tests/cli/superset/sync/dbt/command_test.py b/tests/cli/superset/sync/dbt/command_test.py
index 05cb613d..04714e5d 100644
--- a/tests/cli/superset/sync/dbt/command_test.py
+++ b/tests/cli/superset/sync/dbt/command_test.py
@@ -998,6 +998,8 @@ def test_dbt_cloud(mocker: MockerFixture) -> None:
             "sync",
             "dbt-cloud",
             "XXX",
+            "1",
+            "2",
             "123",
         ],
         catch_exceptions=False,
@@ -1068,6 +1070,8 @@ def test_dbt_cloud_preserve_metadata(mocker: MockerFixture) -> None:
             "sync",
             "dbt-cloud",
             "XXX",
+            "1",
+            "2",
             "123",
             "--preserve-metadata",
         ],
@@ -1121,6 +1125,8 @@ def test_dbt_cloud_preserve_columns(mocker: MockerFixture) -> None:
             "sync",
             "dbt-cloud",
             "XXX",
+            "1",
+            "2",
             "123",
             "--preserve-columns",
         ],
@@ -1174,6 +1180,8 @@ def test_dbt_cloud_merge_metadata(mocker: MockerFixture) -> None:
             "sync",
             "dbt-cloud",
             "XXX",
+            "1",
+            "2",
             "123",
             "--merge-metadata",
         ],
@@ -1207,6 +1215,8 @@ def test_dbt_cloud_preserve_and_merge(mocker: MockerFixture) -> None:
             "sync",
             "dbt-cloud",
             "XXX",
+            "1",
+            "2",
             "123",
             "--preserve-metadata",
             "--merge-metadata",
@@ -1424,6 +1434,8 @@ def test_dbt_cloud_no_database(mocker: MockerFixture) -> None:
             "sync",
             "dbt-cloud",
             "XXX",
+            "1",
+            "2",
             "123",
         ],
         catch_exceptions=False,
@@ -1460,6 +1472,8 @@ def test_dbt_cloud_invalid_job_id(mocker: MockerFixture) -> None:
             "sync",
             "dbt-cloud",
             "XXX",
+            "1",
+            "2",
             "123",
         ],
         catch_exceptions=False,
@@ -1503,6 +1517,8 @@ def test_dbt_cloud_multiple_databases(mocker: MockerFixture) -> None:
                 "sync",
                 "dbt-cloud",
                 "XXX",
+                "1",
+                "2",
                 "123",
             ],
             catch_exceptions=False,
@@ -1618,6 +1634,8 @@ def test_dbt_cloud_exposures_only(mocker: MockerFixture, fs: FakeFilesystem) ->
             "sync",
             "dbt-cloud",
             "XXX",
+            "1",
+            "2",
             "123",
             "--exposures",
             str(exposures),