fix(superset): get_database

preset-io · May 30, 2023 · 24ffa69 · 24ffa69
1 parent 2dd046d
commit 24ffa69
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 2 deletions.
diff --git a/src/preset_cli/api/clients/superset.py b/src/preset_cli/api/clients/superset.py
@@ -477,7 +477,17 @@ def get_database(self, database_id: int) -> Any:
         """
         Return a single database.
         """
-        return self.get_resource("database", database_id)
+        database = self.get_resource("database", database_id)
+        if "sqlalchemy_uri" in database:
+            return database
+
+        url = self.baseurl / "api/v1/database" / str(database_id) / "connection"
+        response = self.session.get(url)
+        validate_response(response)
+
+        resource = response.json()["result"]
+
+        return resource
 
     def get_databases(self, **kwargs: str) -> List[Any]:
         """

diff --git a/tests/api/clients/superset_test.py b/tests/api/clients/superset_test.py
@@ -889,7 +889,9 @@ def test_get_database(mocker: MockerFixture) -> None:
     """
     auth = Auth()
     client = SupersetClient("https://superset.example.org/", auth)
-    get_resource = mocker.patch.object(client, "get_resource")
+    get_resource = mocker.patch.object(
+        client, "get_resource", return_value={"sqlalchemy_uri": "preset://"}
+    )
 
     client.get_database(1)
     get_resource.assert_called_with("database", 1)
@@ -2957,3 +2959,84 @@ def test_create_virtual_dataset(requests_mock: Mocker) -> None:
             },
         ],
     }
+
+
+def test_get_database_connection(requests_mock: Mocker) -> None:
+    """
+    Test that the client uses the new database endpoint for DB info.
+    """
+    requests_mock.get(
+        "https://superset.example.org/api/v1/database/1",
+        json={
+            "id": 1,
+            "result": {
+                "allow_ctas": False,
+                "allow_cvas": False,
+                "allow_dml": True,
+                "allow_file_upload": True,
+                "allow_run_async": False,
+                "backend": "preset",
+                "cache_timeout": None,
+                "configuration_method": "sqlalchemy_form",
+                "database_name": "File Uploads",
+                "driver": "apsw",
+                "engine_information": {
+                    "disable_ssh_tunneling": False,
+                    "supports_file_upload": True,
+                },
+                "expose_in_sqllab": True,
+                "force_ctas_schema": None,
+                "id": 1,
+                "impersonate_user": False,
+                "is_managed_externally": False,
+                "uuid": "d21c03ce-86ec-461f-ac82-0add92a8ee07",
+            },
+        },
+    )
+    requests_mock.get(
+        "https://superset.example.org/api/v1/database/1/connection",
+        json={
+            "id": 1,
+            "result": {
+                "allow_ctas": False,
+                "allow_cvas": False,
+                "allow_dml": True,
+                "allow_file_upload": True,
+                "allow_run_async": False,
+                "backend": "preset",
+                "cache_timeout": None,
+                "configuration_method": "sqlalchemy_form",
+                "database_name": "File Uploads",
+                "driver": "apsw",
+                "engine_information": {
+                    "disable_ssh_tunneling": False,
+                    "supports_file_upload": True,
+                },
+                "expose_in_sqllab": True,
+                "extra": json.dumps(
+                    {
+                        "allows_virtual_table_explore": True,
+                        "metadata_params": {},
+                        "engine_params": {},
+                        "schemas_allowed_for_file_upload": ["main"],
+                    },
+                ),
+                "force_ctas_schema": None,
+                "id": 1,
+                "impersonate_user": False,
+                "is_managed_externally": False,
+                "masked_encrypted_extra": None,
+                "parameters": {},
+                "parameters_schema": {},
+                "server_cert": None,
+                "sqlalchemy_uri": "preset://",
+                "uuid": "d21c03ce-86ec-461f-ac82-0add92a8ee07",
+            },
+        },
+    )
+
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+
+    response = client.get_database(1)
+    assert response["sqlalchemy_uri"] == "preset://"