betodealmeida commented on a change in pull request #13561:
URL: https://github.com/apache/superset/pull/13561#discussion_r613694105
##########
File path: superset/utils/data.py
##########
@@ -161,5 +227,145 @@ def generate_data(columns: List[ColumnInfo], num_rows:
int) -> List[Dict[str, An
def generate_column_data(column: ColumnInfo, num_rows: int) -> List[Any]:
- func = get_type_generator(column["type"])
- return [func() for _ in range(num_rows)]
+ gen = get_type_generator(column["type"])
+ return [gen() for _ in range(num_rows)]
+
+
+def add_to_model(session: Session, model: Type[Model], count: int) ->
List[Model]:
+ """
+ Add entities of a given model.
+
+ :param Model model: a Superset/FAB model
+ :param int count: how many entities to generate and insert
+ """
+ inspector = inspect(model)
+
+ # select samples to copy relationship values
+ relationships = inspector.relationships.items()
+ samples = session.query(model).limit(count).all() if relationships else []
+
+ entities: List[Model] = []
+ max_primary_key: Optional[int] = None
+ for i in range(count):
+ sample = samples[i % len(samples)] if samples else None
+ kwargs = {}
+ for column in inspector.columns.values():
+ # for primary keys, keep incrementing
+ if column.primary_key:
+ if max_primary_key is None:
+ max_primary_key = (
+ session.query(func.max(getattr(model,
column.name))).scalar()
+ or 0
+ )
+ max_primary_key += 1
+ kwargs[column.name] = max_primary_key
+
+ # if the column has a foreign key, copy the value from an existing
entity
+ elif column.foreign_keys:
+ if sample:
+ kwargs[column.name] = getattr(sample, column.name)
+ else:
+ kwargs[column.name] = get_valid_foreign_key(column)
+
+ # should be an enum but it's not
+ elif column.name == "datasource_type":
Review comment:
I'm not expecting this to be common, which is why I left it inline. If
it happens again for other models I'll refactor and move it to a separate part.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
