chore: Migrate get_or_create_table endpoint to api v1 (#22931)

tests/integration_tests/datasets/api_tests.py

@@ -34,6 +34,7 @@ from superset.dao.exceptions import (
     DAODeleteFailedError,
     DAOUpdateFailedError,
 )
+from superset.datasets.commands.exceptions import DatasetCreateFailedError
 from superset.datasets.models import Dataset
 from superset.extensions import db, security_manager
 from superset.models.core import Database
@@ -474,6 +475,7 @@ class TestDatasetApi(SupersetTestCase):
             "can_write",
             "can_export",
             "can_duplicate",
+            "can_get_or_create_dataset",
         }
 
     def test_create_dataset_item(self):
@@ -2302,3 +2304,90 @@ class TestDatasetApi(SupersetTestCase):
         }
         rv = self.post_assert_metric(uri, table_data, "duplicate")
         assert rv.status_code == 422
+
+    @pytest.mark.usefixtures("app_context", "virtual_dataset")
+    def test_get_or_create_dataset_already_exists(self):
+        """
+        Dataset API: Test get or create endpoint when table already exists
+        """
+        self.login(username="admin")
+        rv = self.client.post(
+            "api/v1/dataset/get_or_create/",
+            json={
+                "table_name": "virtual_dataset",
+                "database_id": get_example_database().id,
+            },
+        )
+        self.assertEqual(rv.status_code, 200)
+        response = json.loads(rv.data.decode("utf-8"))
+        dataset = (
+            db.session.query(SqlaTable)
+            .filter(SqlaTable.table_name == "virtual_dataset")
+            .one()
+        )
+        self.assertEqual(response["result"], {"table_id": dataset.id})
+
+    def test_get_or_create_dataset_database_not_found(self):
+        """
+        Dataset API: Test get or create endpoint when database doesn't exist
+        """
+        self.login(username="admin")
+        rv = self.client.post(
+            "api/v1/dataset/get_or_create/",
+            json={"table_name": "virtual_dataset", "database_id": 999},
+        )
+        self.assertEqual(rv.status_code, 422)
+        response = json.loads(rv.data.decode("utf-8"))
+        self.assertEqual(response["message"], {"database": ["Database does not exist"]})
+
+    @patch("superset.datasets.commands.create.CreateDatasetCommand.run")
+    def test_get_or_create_dataset_create_fails(self, command_run_mock):
+        """
+        Dataset API: Test get or create endpoint when create fails
+        """
+        command_run_mock.side_effect = DatasetCreateFailedError
+        self.login(username="admin")
+        rv = self.client.post(
+            "api/v1/dataset/get_or_create/",
+            json={
+                "table_name": "virtual_dataset",
+                "database_id": get_example_database().id,
+            },
+        )
+        self.assertEqual(rv.status_code, 422)
+        response = json.loads(rv.data.decode("utf-8"))
+        self.assertEqual(response["message"], "Dataset could not be created.")
+
+    def test_get_or_create_dataset_creates_table(self):
+        """
+        Dataset API: Test get or create endpoint when table is created
+        """
+        self.login(username="admin")
+
+        examples_db = get_example_database()
+        with examples_db.get_sqla_engine_with_context() as engine:
+            engine.execute("DROP TABLE IF EXISTS test_create_sqla_table_api")
+            engine.execute("CREATE TABLE test_create_sqla_table_api AS SELECT 2 as col")
+
+        rv = self.client.post(
+            "api/v1/dataset/get_or_create/",
+            json={
+                "table_name": "test_create_sqla_table_api",
+                "database_id": examples_db.id,
+                "template_params": '{"param": 1}',
+            },
+        )
+        self.assertEqual(rv.status_code, 200)
+        response = json.loads(rv.data.decode("utf-8"))
+        table = (
+            db.session.query(SqlaTable)
+            .filter_by(table_name="test_create_sqla_table_api")
+            .one()
+        )
+        self.assertEqual(response["result"], {"table_id": table.id})
+        self.assertEqual(table.template_params, '{"param": 1}')
+
+        db.session.delete(table)
+        with examples_db.get_sqla_engine_with_context() as engine:
+            engine.execute("DROP TABLE test_create_sqla_table_api")
+        db.session.commit()

tests/integration_tests/datasets/commands_tests.py

@@ -20,13 +20,18 @@ from unittest.mock import patch
 
 import pytest
 import yaml
+from sqlalchemy.exc import SQLAlchemyError
 
 from superset import db, security_manager
 from superset.commands.exceptions import CommandInvalidError
 from superset.commands.importers.exceptions import IncorrectVersionError
 from superset.connectors.sqla.models import SqlaTable
 from superset.databases.commands.importers.v1 import ImportDatabasesCommand
-from superset.datasets.commands.exceptions import DatasetNotFoundError
+from superset.datasets.commands.create import CreateDatasetCommand
+from superset.datasets.commands.exceptions import (
+    DatasetInvalidError,
+    DatasetNotFoundError,
+)
 from superset.datasets.commands.export import ExportDatasetsCommand
 from superset.datasets.commands.importers import v0, v1
 from superset.models.core import Database
@@ -519,3 +524,47 @@ def _get_table_from_list_by_name(name: str, tables: List[Any]):
         if table.table_name == name:
             return table
     raise ValueError(f"Table {name} does not exists in database")
+
+
+class TestCreateDatasetCommand(SupersetTestCase):
+    def test_database_not_found(self):
+        self.login(username="admin")
+        with self.assertRaises(DatasetInvalidError):
+            CreateDatasetCommand({"table_name": "table", "database": 9999}).run()
+
+    @patch("superset.models.core.Database.get_table")
+    def test_get_table_from_database_error(self, get_table_mock):
+        self.login(username="admin")
+        get_table_mock.side_effect = SQLAlchemyError
+        with self.assertRaises(DatasetInvalidError):
+            CreateDatasetCommand(
+                {"table_name": "table", "database": get_example_database().id}
+            ).run()
+
+    @patch("superset.security.manager.g")
+    @patch("superset.commands.utils.g")
+    def test_create_dataset_command(self, mock_g, mock_g2):
+        mock_g.user = security_manager.find_user("admin")
+        mock_g2.user = mock_g.user
+        examples_db = get_example_database()
+        with examples_db.get_sqla_engine_with_context() as engine:
+            engine.execute("DROP TABLE IF EXISTS test_create_dataset_command")
+            engine.execute(
+                "CREATE TABLE test_create_dataset_command AS SELECT 2 as col"
+            )
+
+        table = CreateDatasetCommand(
+            {"table_name": "test_create_dataset_command", "database": examples_db.id}
+        ).run()
+        fetched_table = (
+            db.session.query(SqlaTable)
+            .filter_by(table_name="test_create_dataset_command")
+            .one()
+        )
+        self.assertEqual(table, fetched_table)
+        self.assertEqual([owner.username for owner in table.owners], ["admin"])
+
+        db.session.delete(table)
+        with examples_db.get_sqla_engine_with_context() as engine:
+            engine.execute("DROP TABLE test_create_dataset_command")
+        db.session.commit()