def export_db(
archive,
include_db=True,
include_logs=True,
task_version=utcnow().strftime("%Y%m%d_%H%M%S"),
):
# type: (Path, bool, bool, str)-> None
from dbnd._core.current import get_databand_context
logger.info("Compressing files to %s..." % archive)
with tarfile.open(str(archive), "w:gz") as tar:
if include_db:
dbnd_context = get_databand_context()
conn_string = dbnd_context.config.get("webserver", "sql_alchemy_conn")
if conn_string.startswith("postgresql"):
with tempfile.NamedTemporaryFile(prefix="dbdump.", suffix=".sql") as tf:
dump_postgres(conn_string, tf.name)
tar.add(tf.name, arcname="postgres-dbnd.sql")
else:
raise DatabandRuntimeError(
"Can not export db! "
"Currently, we support only sqlite and postgres db in automatic export"
)
if include_logs:
context = get_databand_context()
local_env = context.settings.get_env_config(CloudType.local)
logs_folder = local_env.dbnd_local_root.folder("logs").path
if os.path.exists(logs_folder):
logger.info("Adding run folder from '%s'", logs_folder)
tar.add(logs_folder, "run")
else:
logger.warning("Logs dir '%s' is not found", logs_folder)
def build_dbnd_task(self, task_name, task_kwargs=None, expected_type=None):
task_kwargs = task_kwargs or dict()
task_kwargs.setdefault("task_name", task_name)
task_cls = self.get_task_cls(task_name) # type: Type[Task]
if is_airflow_enabled():
from dbnd_airflow.dbnd_task_executor.airflow_operator_as_dbnd import (
AirflowDagAsDbndTask, )
if issubclass(task_cls, AirflowDagAsDbndTask):
# we are running old style dag
dag = self._get_aiflow_dag(task_name)
airflow_task = AirflowDagAsDbndTask.build_dbnd_task_from_dag(
dag=dag)
return airflow_task
try:
logger.debug("Building %s task",
task_cls.task_definition.full_task_family)
obj = task_cls(**task_kwargs)
except Exception:
exc = get_databand_context().settings.log.format_exception_as_str(
sys.exc_info(), isolate=True)
logger.error("Failed to build %s: \n\n%s",
task_cls.get_task_family(), exc)
raise
if expected_type and not issubclass(task_cls, expected_type):
raise friendly_error.task_registry.wrong_type_for_task(
task_name, task_cls, expected_type)
return obj
def __init__(
self,
task_name,
task_definition, # type: TaskDefinition
task_signature_obj,
task_params, # type: Parameters
):
super(_BaseTask, self).__init__()
self.task_definition = task_definition
self.task_params = task_params
self.task_name = task_name # type: str
# miscellaneous
self.task_id = "{}__{}".format(task_name, task_signature_obj.signature)
self.task_type = self._conf__task_type_name
self.task_signature_obj = task_signature_obj
self.task_outputs_signature_obj = None
# define it at the time of creation
# we can remove the strong reference, the moment we have global cache for instances
# otherwise if we moved to another databand_context, task_id relationships will not be found
self.dbnd_context = get_databand_context()
# we count task meta creation
# even if cached task is going to be used we will increase creation id
# if t2 created after t1, t2.task_creation_id > t1.task_creation_id
_BaseTask._current_task_creation_id += 1
self.task_creation_id = _BaseTask._current_task_creation_id
def get_job_by_name(job_name):
query_params = build_query_api_params(
filters=build_job_name_filter(job_name))
job_result = _get_job(get_databand_context().databand_api_client,
query_params)
return first(job_result, None)
def __init__(self, config, task_cls, task_args, task_kwargs):
# type:(DbndConfig, Type[_BaseTask], Any, Any)->None
self.task_cls = task_cls
self.task_definition = task_cls.task_definition # type: TaskDefinition
# keep copy of user inputs
self.task_kwargs__ctor = task_kwargs.copy()
self.task_args__ctor = list(task_args)
self.parent_task = try_get_current_task()
# let find if we are running this constructor withing another Databand Task
self.dbnd_context = get_databand_context()
self.task_call_source = [
self.dbnd_context.user_code_detector.find_user_side_frame(2)
]
if self.task_call_source and self.parent_task:
self.task_call_source.extend(
self.parent_task.task_meta.task_call_source)
self.task_family = self.task_definition.task_family
self.task_name = self.task_family
self.multi_sec_conf = MultiSectionConfig(config, [])
self._task_params = self.task_definition._task_params.clone()
self.ctor_kwargs = {}
self._exc_desc = self.task_family
self.task_errors = []
def get_scheduled_jobs(name_pattern=None,
from_file_only=False,
include_deleted=False,
ctx=None):
ctx = ctx or get_databand_context()
schema = ScheduledJobSchemaV2(strict=False)
query_filter = []
if from_file_only:
query_filter.append({"name": "from_file", "op": "eq", "val": True})
if not include_deleted:
query_filter.append({
"name": "deleted_from_file",
"op": "eq",
"val": False
})
if name_pattern:
query_filter.append({
"name": "name",
"op": "like",
"val": name_pattern
})
query = {"filter": json.dumps(query_filter)}
res = ctx.databand_api_client.api_request("/api/v1/scheduled_jobs",
None,
method="GET",
query=query,
no_prefix=True)
return schema.load(data=res["data"], many=True).data
def create_alert(
job_name,
job_id,
task_name,
uid,
alert_class,
severity,
operator,
value,
user_metric,
):
"""add alert for existing job or scheduled job"""
alert = _build_alert(
job_name,
job_id,
task_name,
uid,
alert_class,
severity,
operator,
value,
user_metric,
)
alert_def_uid = _post_alert(get_databand_context().databand_api_client,
alert)
return alert_def_uid
def _sig(task):
with new_databand_run(context=get_databand_context(), job_name=task.task_name):
for child in [task] + list(task.descendants.get_children()):
name = "signature %s" % child.task_name
logger.info(child.ctrl.banner(name))
return task.task_signature
def log_exception_to_server(exception=None):
try:
from dbnd._core.current import get_databand_context
client = get_databand_context().databand_api_client
if client is None or not client.is_configured():
return
e_type, e_value, e_traceback = sys.exc_info()
if exception:
e_type, e_value, e_traceback = (
type(exception),
exception,
exception.__traceback__,
)
trace = _format_exception(e_type, e_value, e_traceback)
data = {
"dbnd_version": dbnd.__version__,
"source": "tracking-sdk",
"stack_trace": trace,
"timestamp": utcnow().isoformat(),
}
return client.api_request(endpoint="log_exception",
method="POST",
data=data)
except Exception: # noqa
logger.debug("Error sending monitoring exception message",
exc_info=True)
def edit_airflow_instance(
tracking_source_uid,
url,
external_url,
fetcher,
env,
dag_ids,
last_seen_dag_run_id,
last_seen_log_id,
name,
system_alert_definitions,
monitor_config,
):
client = get_databand_context().databand_api_client
endpoint = "airflow_monitor/edit"
request_data = {
"tracking_source_uid": tracking_source_uid,
"base_url": url,
"external_url": external_url,
"fetcher": fetcher,
"env": env,
"monitor_config": monitor_config,
"dag_ids": dag_ids,
"last_seen_dag_run_id": last_seen_dag_run_id,
"last_seen_log_id": last_seen_log_id,
"name": name,
"system_alert_definitions": system_alert_definitions,
}
client.api_request(endpoint, request_data, method="POST")
def send_heartbeat(run_uid, databand_url, heartbeat_interval, driver_pid,
tracker, tracker_api):
from dbnd import config
from dbnd._core.settings import CoreConfig
from dbnd._core.task_executor.heartbeat_sender import send_heartbeat_continuously
with config({
"core": {
"tracker": tracker.split(","),
"tracker_api": tracker_api,
"databand_url": databand_url,
}
}):
requred_context = []
if tracker_api == "db":
from dbnd import new_dbnd_context
requred_context.append(
new_dbnd_context(name="send_heartbeat",
autoload_modules=False))
with nested_context.nested(*requred_context):
tracking_store = get_databand_context().tracking_store
send_heartbeat_continuously(run_uid, tracking_store,
heartbeat_interval, driver_pid)
def post_scheduled_job(scheduled_job_dict, ctx=None):
ctx = ctx or get_databand_context()
schema = ScheduledJobSchemaV2(strict=False)
data, _ = schema.dump({"DbndScheduledJob": scheduled_job_dict})
response = ctx.databand_api_client.api_request("/api/v1/scheduled_jobs",
data,
method="POST",
no_prefix=True)
return schema.load(data=response).data
def delete_scheduled_job(scheduled_job_name, revert=False, ctx=None):
ctx = ctx or get_databand_context()
ctx.databand_api_client.api_request(
"/api/v1/scheduled_jobs?name=%s&revert=%s" %
(scheduled_job_name, str(revert).lower()),
None,
method="DELETE",
no_prefix=True,
)
def set_scheduled_job_active(name, value, ctx=None):
ctx = ctx or get_databand_context()
ctx.databand_api_client.api_request(
"/api/v1/scheduled_jobs/set_active?name=%s&value=%s" %
(name, str(value).lower()),
None,
method="PUT",
no_prefix=True,
)
def dbnd_run(self):
# type: (...)-> DatabandRun
"""Run task via Databand execution system."""
# this code should be executed under context!
from dbnd._core.current import get_databand_context
ctx = get_databand_context()
run = ctx.dbnd_run_task(self)
return run
def _sig(task):
name = "signature %s" % task.task_name
with new_databand_run(context=get_databand_context(),
task_or_task_name=task):
logger.info(task.ctrl.banner(name))
for child in task.task_meta.get_children():
logger.info(child.ctrl.banner(name))
return task.task_signature
def test_prod_immutable_output_dict_prod(self):
env = get_databand_context().env
prod_env = env.clone(production=True)
task = TProdImmutbaleOutputs(task_env=prod_env)
assert task.task_enabled_in_prod
assert task.task_signature[:5] not in str(task.splits)
config.log_current_config()
actual = assert_run_task(task)
print(actual)
def get_scheduled_job_by_name(job_name):
query_params = build_query_api_params(filters=build_scheduled_job_filter(
job_name=job_name))
job_result = _get_scheduled_jobs(
get_databand_context().databand_api_client, query_params)
try:
return first(job_result)["DbndScheduledJob"]
except ValueError:
return None
def run_dbnd_subprocess(args,
retcode=255,
clean_env=False,
blocking=True,
**kwargs):
# implement runner with https://docs.pytest.org/en/latest/capture.html
# do not run in subprocess
# main.main(['run', '--module', str(factories.__name__), ] + args)
# return
kwargs = kwargs.copy()
cmd_args = list(map(str, args))
env = kwargs.pop("env", os.environ).copy()
if clean_env:
for key in list(env.keys()):
if key.startswith("DBND") or key.startswith("AIRFLOW"):
del env[key]
# env['PYTHONUNBUFFERED'] = 'false'
# env['PYTHONPATH'] = env.get('PYTHONPATH', '') + ':.:test'
from dbnd._core.current import get_databand_context
env["DBND__CORE__SQL_ALCHEMY_CONN"] = get_databand_context(
).settings.core.get_sql_alchemy_conn()
cmd_line = list2cmdline(cmd_args)
logger.info("Running at %s: %s", kwargs.get("cwd", "current dir"),
cmd_line) # To simplify rerunning failing tests
if blocking:
try:
output = fast_subprocess.check_output(cmd_args,
stderr=subprocess.STDOUT,
env=env,
**kwargs)
# we don't decode ascii //.decode("ascii")
output = output.decode("utf-8")
logger.info("Cmd line %s output:\n %s", cmd_line, output)
return output
except subprocess.CalledProcessError as ex:
logger.error(
"Failed to run %s :\n\n\n -= Output =-\n%s\n\n\n -= See output above =-",
cmd_line,
ex.output.decode("utf-8", errors="ignore"),
)
if ex.returncode == retcode:
return ex.output.decode("utf-8")
raise ex
else:
return subprocess.Popen(cmd_args,
stderr=subprocess.STDOUT,
env=env,
**kwargs)
def create_airflow_instance(
url,
external_url,
fetcher,
env,
dag_ids,
last_seen_dag_run_id,
last_seen_log_id,
name,
generate_token,
system_alert_definitions,
monitor_config,
):
client = get_databand_context().databand_api_client
endpoint = "airflow_monitor/add"
request_data = {
"base_url": url,
"external_url": external_url,
"fetcher": fetcher,
"env": env,
"monitor_config": monitor_config,
"dag_ids": dag_ids,
"last_seen_dag_run_id": last_seen_dag_run_id,
"last_seen_log_id": last_seen_log_id,
"name": name,
"system_alert_definitions": system_alert_definitions,
}
if monitor_config:
request_data["monitor_config"] = monitor_config
resp = client.api_request(endpoint, request_data, method="POST")
config_json = resp["server_info_dict"]
config_json["core"]["databand_url"] = get_databand_context(
).settings.core.databand_url
if generate_token:
token_resp = generate_access_token(name, generate_token)
config_json["core"]["databand_access_token"] = token_resp["token"]
return config_json
def generate_access_token(name, lifespan):
client = get_databand_context().databand_api_client
resp = client.api_request(
"/api/v1/auth/personal_access_token",
{
"label": name,
"lifespan": lifespan
},
method="POST",
)
return resp
def get_alerts_filtered(alert_def_uid=None, job_name=None, custom_name=None):
"""get alerts by filters"""
params = {
"uids": alert_def_uid,
"job_name": job_name,
"custom_name": custom_name
}
query_params = "&".join(f"{key}={value}" for key, value in params.items()
if value is not None)
return _get_alerts(get_databand_context().databand_api_client,
query_params)
def create_airflow_instance(url, external_url, fetcher, api_mode, composer_client_id):
client = get_databand_context().databand_api_client
endpoint = "airflow_monitor/add"
request_data = {
"base_url": url,
"ext_url": external_url,
"fetcher": fetcher,
"api_mode": api_mode,
"composer_client_id": composer_client_id,
}
client.api_request(endpoint, request_data, method="POST")
def get_scheduled_jobs(name_pattern=None,
from_file_only=None,
include_deleted=False,
ctx=None):
ctx = ctx or get_databand_context()
query_params = build_query_api_params(filters=build_scheduled_job_filter(
name_pattern=name_pattern,
from_file_flag=from_file_only,
include_deleted_flag=include_deleted,
))
return _get_scheduled_jobs(ctx.databand_api_client, query_params)
def from_databand_context(cls, name=None):
"""Syntax sugar for accessing the current config instance."""
from dbnd._core.current import get_databand_context
if not name:
if cls._conf__task_family:
# using the current cls section name to get the current instance of a class
name = cls._conf__task_family
else:
raise ConfigLookupError(
"name is required for retrieving a config instance"
)
return get_databand_context().settings.get_config(name)
def parse_from_str(self, s):
s = super(VersionValueType, self).parse_from_str(s)
if not s:
return s
if s.lower() == VersionAlias.now:
return utcnow().strftime("%Y%m%d_%H%M%S")
if s.lower() == VersionAlias.context_uid:
return get_databand_context().current_context_uid
if s.lower() == VersionAlias.git:
return str(get_project_git())
return s
def _create_task(cls, args, kwargs):
task_definition = cls.task_definition
# we need to have context initialized before we start to run all logic in config() scope
# update config with current class defaults
# we apply them to config only if there are no values (this is defaults)
with config(
config_values=task_definition.task_defaults_config_store,
source=task_definition.task_passport.format_source_name(
"defaults"),
merge_settings=ConfigMergeSettings.on_non_exists_only,
) as task_config:
tracking_mode = TaskEssence.TRACKING.is_included(cls)
# create task meta first
task_meta_factory = (TrackedTaskMetaFactory
if tracking_mode else TaskMetaFactory)
factory = task_meta_factory(config=task_config,
task_cls=cls,
task_args=args,
task_kwargs=kwargs)
task_meta = factory.create_dbnd_task_meta()
# If a Task has already been instantiated with the same parameters,
# the previous instance is returned to reduce number of object instances.
tic = get_databand_context().task_instance_cache
task = tic.get_task_obj_by_id(task_meta.obj_key.id)
if not task or tracking_mode or hasattr(task, "_dbnd_no_cache"):
task = cls._build_task_obj(task_meta)
tic.register_task_obj_instance(task)
# now the task is created - all nested constructors will see it as parent
with task_context(task, TaskContextPhase.BUILD):
task._initialize()
task._validate()
task.task_meta.config_layer = config.config_layer
tic.register_task_instance(task)
parent_task = try_get_current_task()
if (parent_task and hasattr(task, "task_id")
and (task.task_essence != TaskEssence.CONFIG)):
parent_task.descendants.add_child(task.task_id)
return task
def create_alert(
job_name,
job_id,
task_name,
task_repr,
uid,
alert_class,
severity,
operator,
value,
user_metric,
alert_on_historical_runs=True,
is_str_value=False,
**optional_fields,
):
alert = {
"type": alert_class,
"job_name": job_name,
"job_id": job_id,
"task_name": task_name,
"task_repr": task_repr,
"uid": uid,
"operator": operator,
"value": value,
"severity": severity,
"user_metric": user_metric,
"alert_on_historical_runs": alert_on_historical_runs,
"is_str_value": is_str_value,
}
if operator == AlertDefOperator.ANOMALY.value:
alert["ml_alert"] = {
"look_back": optional_fields["look_back"],
"sensitivity": optional_fields["sensitivity"],
}
if operator == AlertDefOperator.RANGE.value:
alert["operator"] = "RangeAlert"
alert["value"] = json.dumps({
"baseline": optional_fields["baseline"],
"range": optional_fields["range"]
})
alert = {key: value for key, value in alert.items() if value is not None}
alert_def_uid = _post_alert(get_databand_context().databand_api_client,
alert)
return alert_def_uid
def create_dbnd_task(config, new_task_factory, task_cls, task_args, task_kwargs):
# type:(DbndConfig, Any, Type[_BaseTask], Any, Any, bool)->None
tracking_mode = task_cls.is_tracking_mode
task_meta_factory = TrackedTaskMetaFactory if tracking_mode else TaskMetaFactory
factory = task_meta_factory(
config=config, task_cls=task_cls, task_args=task_args, task_kwargs=task_kwargs,
)
task_meta = factory.create_dbnd_task_meta()
# If a Task has already been instantiated with the same parameters,
# the previous instance is returned to reduce number of object instances.
tic = get_databand_context().task_instance_cache
task = tic.get_task_obj_by_id(task_meta.obj_key.id)
if not task or tracking_mode or hasattr(task, "_dbnd_no_cache"):
task = new_task_factory(task_meta)
tic.register_task_obj_instance(task)
# now the task is created - all nested constructors will see it as parent
with task_context(task, TaskContextPhase.BUILD):
task._initialize()
task._validate()
task.task_meta.config_layer = config.config_layer
tic.register_task_instance(task)
parent_task = try_get_current_task()
if (
parent_task
and hasattr(task, "task_id")
and isinstance(task, _TaskParamContainer)
):
parent_task.task_meta.add_child(task.task_id)
return task
def delete_alerts(uids):
"""delete alerts by uids"""
get_databand_context().databand_api_client.api_request(
endpoint="alert_defs/delete", data=uids, method="POST"
)
