Lakehouse Admin
86 lines
3.2 KiB
Python
86 lines
3.2 KiB
Python
"""Generate and save query_context for API-created Superset charts."""
|
|
import json
|
|
|
|
app = __import__("superset.app", fromlist=["create_app"]).create_app()
|
|
|
|
with app.app_context():
|
|
from flask import g
|
|
from superset.extensions import db
|
|
from superset.models.slice import Slice
|
|
from superset.models.core import Database
|
|
from superset.charts.schemas import ChartDataQueryContextSchema
|
|
from superset import security_manager
|
|
|
|
admin = security_manager.find_user(username="admin")
|
|
g.user = admin
|
|
|
|
charts = db.session.query(Slice).order_by(Slice.id).all()
|
|
for sl in charts:
|
|
try:
|
|
fd = sl.form_data
|
|
metric = fd.get("metric")
|
|
metrics = fd.get("metrics") or ([metric] if metric else [])
|
|
if not metrics:
|
|
metrics = [
|
|
{
|
|
"expressionType": "SQL",
|
|
"sqlExpression": "COUNT(*)",
|
|
"label": "COUNT(*)",
|
|
}
|
|
]
|
|
|
|
groupby = fd.get("groupby") or []
|
|
payload = {
|
|
"datasource": {
|
|
"id": sl.datasource_id,
|
|
"type": sl.datasource_type,
|
|
},
|
|
"force": False,
|
|
"queries": [
|
|
{
|
|
"filters": [],
|
|
"extras": {"having": "", "where": ""},
|
|
"applied_time_extras": {},
|
|
"columns": groupby if isinstance(groupby, list) else [],
|
|
"metrics": metrics,
|
|
"orderby": [],
|
|
"annotation_layers": [],
|
|
"row_limit": int(fd.get("row_limit") or 1000),
|
|
"series_limit": 0,
|
|
"order_desc": True,
|
|
"url_params": {},
|
|
"custom_params": {},
|
|
"custom_form_data": {},
|
|
}
|
|
],
|
|
"form_data": fd,
|
|
"result_format": "json",
|
|
"result_type": "full",
|
|
}
|
|
qc = ChartDataQueryContextSchema().load(payload)
|
|
ctx = qc.cache_values if hasattr(qc, "cache_values") else None
|
|
if ctx is None:
|
|
# fallback: store factory input dict
|
|
from superset.common.query_context_factory import QueryContextFactory
|
|
|
|
factory = QueryContextFactory()
|
|
ctx = {
|
|
"datasource": {
|
|
"id": sl.datasource_id,
|
|
"type": sl.datasource_type,
|
|
},
|
|
"force": False,
|
|
"queries": payload["queries"],
|
|
"form_data": fd,
|
|
"result_format": "json",
|
|
"result_type": "full",
|
|
}
|
|
sl.query_context = json.dumps(ctx) if isinstance(ctx, dict) else json.dumps(payload)
|
|
sl.query_context_generation = True
|
|
db.session.add(sl)
|
|
print("OK", sl.id, sl.slice_name[:50])
|
|
except Exception as e:
|
|
print("ERR", sl.id, sl.slice_name[:40], e)
|
|
db.session.commit()
|
|
print("committed")
|