iop
/
superset
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

This commit is contained in:
Jack Fragassi 2023-02-15 02:38:51 -08:00 committed by GitHub
parent 7d7e3e4ade
commit 3fae7dbc34
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 237 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
superset-frontend/src/SqlLab/actions/sqlLab.js
View File

@ -1513,13 +1513,13 @@ export function createCtasDatasource(vizOptions) {
return dispatch => {
dispatch(createDatasourceStarted());
return SupersetClient.post({
endpoint: '/superset/get_or_create_table/',
postPayload: { data: vizOptions },
endpoint: '/api/v1/dataset/get_or_create/',
jsonPayload: vizOptions,
})
.then(({ json }) => {
dispatch(createDatasourceSuccess(json));
dispatch(createDatasourceSuccess(json.result));
return json;
return json.result;
})
.catch(() => {
const errorMsg = t('An error occurred while creating the data source');

6
superset-frontend/src/SqlLab/components/ExploreCtasResultsButton/index.tsx
View File

@ -48,10 +48,10 @@ const ExploreCtasResultsButton = ({
const dispatch = useDispatch<(dispatch: any) => Promise<JsonObject>>();
const buildVizOptions = {
datasourceName: table,
table_name: table,
schema,
dbId,
templateParams,
database_id: dbId,
template_params: templateParams,
};
const visualize = () => {

2
superset/connectors/sqla/views.py
View File

@ -35,7 +35,6 @@ from superset.constants import MODEL_VIEW_RW_METHOD_PERMISSION_MAP, RouteMethod
from superset.superset_typing import FlaskResponse
from superset.utils import core as utils
from superset.views.base import (
create_table_permissions,
DatasourceFilter,
DeleteMixin,
ListWidgetWithCheckboxes,
@ -511,7 +510,6 @@ class TableModelView( # pylint: disable=too-many-ancestors
) -> None:
if fetch_metadata:
item.fetch_metadata()
create_table_permissions(item)
if flash_message:
flash(
_(

70
superset/datasets/api.py
View File

@ -61,6 +61,7 @@ from superset.datasets.schemas import (
DatasetRelatedObjectsResponse,
get_delete_ids_schema,
get_export_ids_schema,
GetOrCreateDatasetSchema,
)
from superset.utils.core import parse_boolean_string
from superset.views.base import DatasourceFilter, generate_download_headers
@ -93,6 +94,7 @@ class DatasetRestApi(BaseSupersetModelRestApi):
"refresh",
"related_objects",
"duplicate",
"get_or_create_dataset",
}
list_columns = [
"id",
@ -240,6 +242,7 @@ class DatasetRestApi(BaseSupersetModelRestApi):
openapi_spec_component_schemas = (
DatasetRelatedObjectsResponse,
DatasetDuplicateSchema,
GetOrCreateDatasetSchema,
)
list_outer_default_load = True
@ -877,3 +880,70 @@ class DatasetRestApi(BaseSupersetModelRestApi):
)
command.run()
return self.response(200, message="OK")
@expose("/get_or_create/", methods=["POST"])
@protect()
@safe
@statsd_metrics
@event_logger.log_this_with_context(
action=lambda self, *args, **kwargs: f"{self.__class__.__name__}"
f".get_or_create_dataset",
log_to_statsd=False,
)
def get_or_create_dataset(self) -> Response:
"""Retrieve a dataset by name, or create it if it does not exist
---
post:
summary: Retrieve a table by name, or create it if it does not exist
requestBody:
required: true
content:
application/json:
schema:
$ref: '#/components/schemas/GetOrCreateDatasetSchema'
responses:
200:
description: The ID of the table
content:
application/json:
schema:
type: object
properties:
result:
type: object
properties:
table_id:
type: integer
400:
$ref: '#/components/responses/400'
401:
$ref: '#/components/responses/401'
422:
$ref: '#/components/responses/422'
500:
$ref: '#/components/responses/500'
"""
try:
body = GetOrCreateDatasetSchema().load(request.json)
except ValidationError as ex:
return self.response(400, message=ex.messages)
table_name = body["table_name"]
database_id = body["database_id"]
table = DatasetDAO.get_table_by_name(database_id, table_name)
if table:
return self.response(200, result={"table_id": table.id})
body["database"] = database_id
try:
tbl = CreateDatasetCommand(body).run()
return self.response(200, result={"table_id": tbl.id})
except DatasetInvalidError as ex:
return self.response_422(message=ex.normalized_messages())
except DatasetCreateFailedError as ex:
logger.error(
"Error creating model %s: %s",
self.__class__.__name__,
str(ex),
exc_info=True,
)
return self.response_422(message=ex.message)

8
superset/datasets/dao.py
View File

@ -388,6 +388,14 @@ class DatasetDAO(BaseDAO): # pylint: disable=too-many-public-methods
db.session.rollback()
raise ex
@staticmethod
def get_table_by_name(database_id: int, table_name: str) -> Optional[SqlaTable]:
return (
db.session.query(SqlaTable)
.filter_by(database_id=database_id, table_name=table_name)
.one_or_none()
)
class DatasetColumnDAO(BaseDAO):
model_cls = TableColumn

11
superset/datasets/schemas.py
View File

@ -228,6 +228,17 @@ class ImportV1DatasetSchema(Schema):
external_url = fields.String(allow_none=True)
class GetOrCreateDatasetSchema(Schema):
table_name = fields.String(required=True, description="Name of table")
database_id = fields.Integer(
required=True, description="ID of database table belongs to"
)
schema = fields.String(
description="The schema the table belongs to", allow_none=True
)
template_params = fields.String(description="Template params for the table")
class DatasetSchema(SQLAlchemyAutoSchema):
"""
Schema for the ``Dataset`` model.

6
superset/views/base.py
View File

@ -299,12 +299,6 @@ def validate_sqlatable(table: models.SqlaTable) -> None:
) from ex
def create_table_permissions(table: models.SqlaTable) -> None:
security_manager.add_permission_view_menu("datasource_access", table.get_perm())
if table.schema:
security_manager.add_permission_view_menu("schema_access", table.schema_perm)
class BaseSupersetView(BaseView):
@staticmethod
def json_response(obj: Any, status: int = 200) -> FlaskResponse:

4
superset/views/core.py
View File

@ -142,7 +142,6 @@ from superset.views.base import (
api,
BaseSupersetView,
common_bootstrap_payload,
create_table_permissions,
CsvResponse,
data_payload_response,
deprecated,
@ -1927,6 +1926,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
@has_access
@expose("/get_or_create_table/", methods=["POST"])
@event_logger.log_this
@deprecated()
def sqllab_table_viz(self) -> FlaskResponse: # pylint: disable=no-self-use
"""Gets or creates a table object with attributes passed to the API.
@ -1956,11 +1956,11 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
table.schema = data.get("schema")
table.template_params = data.get("templateParams")
# needed for the table validation.
# fn can be deleted when this endpoint is removed
validate_sqlatable(table)
db.session.add(table)
table.fetch_metadata()
create_table_permissions(table)
db.session.commit()
return json_success(json.dumps({"table_id": table.id}))

89
tests/integration_tests/datasets/api_tests.py
View File

@ -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()

51
tests/integration_tests/datasets/commands_tests.py
View File

@ -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()