def _add_table_metrics(datasource: SqlaTable) -> None:
# By accessing the attribute first, we make sure `datasource.columns` and
# `datasource.metrics` are already loaded. Otherwise accessing them later
# may trigger an unnecessary and unexpected `after_update` event.
columns, metrics = datasource.columns, datasource.metrics
if not any(col.column_name == "num_california" for col in columns):
col_state = str(column("state").compile(db.engine))
col_num = str(column("num").compile(db.engine))
columns.append(
TableColumn(
column_name="num_california",
expression=f"CASE WHEN {col_state} = 'CA' THEN {col_num} ELSE 0 END",
)
)
if not any(col.metric_name == "sum__num" for col in metrics):
col = str(column("num").compile(db.engine))
metrics.append(SqlMetric(metric_name="sum__num", expression=f"SUM({col})"))
for col in columns:
if col.column_name == "ds":
col.is_dttm = True
break
datasource.columns = columns
datasource.metrics = metrics
def test_update_sqlatable(mocker: MockFixture, app_context: None,
session: Session) -> None:
"""
Test that updating a ``SqlaTable`` also updates the corresponding ``Dataset``.
"""
# patch session
mocker.patch("superset.security.SupersetSecurityManager.get_session",
return_value=session)
from superset.columns.models import Column
from superset.connectors.sqla.models import SqlaTable, TableColumn
from superset.datasets.models import Dataset
from superset.models.core import Database
from superset.tables.models import Table
engine = session.get_bind()
Dataset.metadata.create_all(engine) # pylint: disable=no-member
columns = [
TableColumn(column_name="ds", is_dttm=1, type="TIMESTAMP"),
]
sqla_table = SqlaTable(
table_name="old_dataset",
columns=columns,
metrics=[],
database=Database(database_name="my_database",
sqlalchemy_uri="sqlite://"),
)
session.add(sqla_table)
session.flush()
dataset = session.query(Dataset).one()
assert len(dataset.columns) == 1
# add a column to the original ``SqlaTable`` instance
sqla_table.columns.append(
TableColumn(column_name="user_id", type="INTEGER"))
session.flush()
# check that the column was added to the dataset
dataset = session.query(Dataset).one()
assert len(dataset.columns) == 2
# delete the column in the original instance
sqla_table.columns = sqla_table.columns[1:]
session.flush()
# check that the column was also removed from the dataset
dataset = session.query(Dataset).one()
assert len(dataset.columns) == 1
# modify the attribute in a column
sqla_table.columns[0].is_dttm = True
session.flush()
# check that the dataset column was modified
dataset = session.query(Dataset).one()
assert dataset.columns[0].is_temporal is True
def load_dataset_with_columns() -> Generator[SqlaTable, None, None]:
with app.app_context():
engine = create_engine(app.config["SQLALCHEMY_DATABASE_URI"],
echo=True)
meta = MetaData()
session = db.session
students = Table(
"students",
meta,
Column("id", Integer, primary_key=True),
Column("name", String(255)),
Column("lastname", String(255)),
Column("ds", Date),
)
meta.create_all(engine)
students.insert().values(name="George", ds="2021-01-01")
dataset = SqlaTable(database_id=db.session.query(Database).first().id,
table_name="students")
column = TableColumn(table_id=dataset.id, column_name="name")
dataset.columns = [column]
session.add(dataset)
session.commit()
yield dataset
# cleanup
students_table = meta.tables.get("students")
if students_table is not None:
base = declarative_base()
# needed for sqlite
session.commit()
base.metadata.drop_all(engine, [students_table], checkfirst=True)
session.delete(dataset)
session.delete(column)
session.commit()
