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