def log_snowflake_table(
table_name, # type: str
connection_string, # type: str
database, # type: str
schema, # type: str
with_preview=True,
with_schema=True,
):
try:
if not is_plugin_enabled("dbnd-snowflake",
module_import="dbnd_snowflake"):
return
from dbnd_snowflake import snowflake_values
with log_duration("log_snowflake_table__time_seconds",
source="system"):
conn_params = snowflake_values.conn_str_to_conn_params(
connection_string)
account = conn_params["account"]
user = conn_params["user"]
password = conn_params["password"]
snowflake_table = snowflake_values.SnowflakeTable(
account, user, password, database, schema, table_name)
log_data(
table_name,
snowflake_table,
with_preview=with_preview,
with_schema=with_schema,
with_histograms=False,
)
except Exception as ex:
logger.exception("Failed to log_snowflake_table")
def signal_handler(signum, frame):
logger.info("Task runner received signal. Exiting...")
if is_plugin_enabled("dbnd-docker") and is_plugin_enabled(
"dbnd-airflow"):
from dbnd_docker.kubernetes.kubernetes_engine_config import (
ENV_DBND_POD_NAME, )
if ENV_DBND_POD_NAME in os.environ:
# We are running inside cluster
# We should log all events on sigterm for debugging when running inside cluster
from dbnd_airflow_contrib.kubernetes_metrics_logger import (
log_pod_events_on_sigterm, )
log_pod_events_on_sigterm(frame)
raise DatabandSigTermError("Task received signal",
help_msg="Probably the job was canceled")
def should_log_pg_histogram(luigi_task):
# type: (luigi.Task) -> bool
if not is_plugin_enabled("dbnd-postgres", module_import="dbnd_postgres"):
return False
try:
from luigi.contrib.postgres import PostgresQuery
except ImportError:
return False
return isinstance(luigi_task, PostgresQuery)
def dbnd_setup_plugin():
# register configs
from dbnd_aws.env import AwsEnvConfig
register_config_cls(AwsEnvConfig)
if is_plugin_enabled("dbnd-spark"):
from dbnd_aws.emr.emr_config import EmrConfig
register_config_cls(EmrConfig)
if is_plugin_enabled("dbnd-docker"):
from dbnd_aws.batch.aws_batch_ctrl import AwsBatchConfig
register_config_cls(AwsBatchConfig)
from dbnd_aws.fs import build_s3_fs_client
register_file_system(FileSystems.s3, build_s3_fs_client)
def dbnd_setup_plugin():
from dbnd_gcp.dataflow.dataflow_config import DataflowConfig
from dbnd_gcp.env import GcpEnvConfig
register_config_cls(GcpEnvConfig)
register_config_cls(DataflowConfig)
if is_plugin_enabled("dbnd-spark"):
from dbnd_gcp.dataproc.dataproc_config import DataprocConfig
register_config_cls(DataprocConfig)
register_file_system(FileSystems.gcs, build_gcs_client)
def log_snowflake_table(
table_name: str,
connection_string: Union[str, SnowflakeConnection],
database: str,
schema: str,
key: Optional[str] = None,
with_preview: Optional[bool] = None,
with_schema: Optional[bool] = None,
raise_on_error: bool = False,
):
"""
:param table_name: table name
:param connection_string: either connection_string or actual connection
:param database:
:param schema:
:param key:
:param with_preview:
:param with_schema:
:param raise_on_error:
:return:
"""
if not is_plugin_enabled("dbnd-snowflake", module_import="dbnd_snowflake"):
return
from dbnd_snowflake import snowflake_values
with log_duration("log_snowflake_table__time_seconds",
source="system"), SnowflakeController(
connection_string) as snowflake_ctrl:
config = SnowflakeConfig()
snowflake_table = snowflake_values.SnowflakeTable(
snowflake_ctrl,
database,
schema,
table_name,
config.table_preview_rows,
)
log_data(
key or table_name,
snowflake_table,
with_preview=with_preview,
with_schema=with_schema,
with_size=with_schema,
with_histograms=False,
raise_on_error=raise_on_error,
)
def log_snowflake_table(
table_name: str,
connection_string: str,
database: str,
schema: str,
key: Optional[str] = None,
with_preview: Optional[bool] = None,
with_schema: Optional[bool] = None,
raise_on_error: bool = False,
):
if not is_plugin_enabled("dbnd-snowflake", module_import="dbnd_snowflake"):
return
from dbnd_snowflake import snowflake_values
with log_duration("log_snowflake_table__time_seconds", source="system"):
conn_params = snowflake_values.conn_str_to_conn_params(connection_string)
account = conn_params["account"]
user = conn_params["user"]
password = conn_params["password"]
config = SnowflakeConfig()
snowflake_table = snowflake_values.SnowflakeTable(
account,
user,
password,
database,
schema,
table_name,
config.table_preview_rows,
)
log_data(
key or table_name,
snowflake_table,
with_preview=with_preview,
with_schema=with_schema,
with_size=with_schema,
with_histograms=False,
raise_on_error=raise_on_error,
)
def log_snowflake_table_targets(
table_op: TableTargetOperation,
connection_string: Union[str, SnowflakeConnection],
with_preview: Optional[bool] = None,
with_schema: Optional[bool] = None,
):
if not is_plugin_enabled("dbnd-snowflake", module_import="dbnd_snowflake"):
return
from dbnd_snowflake.snowflake_values import SnowflakeTable
with SnowflakeController(connection_string) as snowflake_ctrl:
snowflake_table = SnowflakeTable.from_table(snowflake_ctrl,
table_op.name)
log_dataset_op(
op_path=table_op.path,
op_type=table_op.operation,
success=table_op.success,
data=snowflake_table,
with_preview=with_preview,
with_schema=with_schema,
)
def log_pg_table(
table_name,
connection_string,
with_preview=None, # type: Optional[bool]
with_schema=None, # type: Optional[bool]
with_histograms=None, # type: Union[LogDataRequest, bool, str, List[str]]
):
try:
if not is_plugin_enabled("dbnd-postgres", module_import="dbnd_postgres"):
return
from dbnd_postgres import postgres_values
pg_table = postgres_values.PostgresTable(table_name, connection_string)
log_data(
table_name,
pg_table,
with_preview=with_preview,
with_schema=with_schema,
with_histograms=with_histograms,
)
except Exception:
logger.exception("Failed to log_pg_table")
def log_pg_table(
table_name,
connection_string,
with_preview=None, # type: Optional[bool]
with_schema=None, # type: Optional[bool]
with_histograms=None, # type: Union[LogDataRequest, bool, str, List[str]]
):
"""
Log the data of postgres table to dbnd.
@param table_name: name of the table to log
@param connection_string: a connection string used to reach the table.
@param with_preview: True if should log a preview of the table.
@param with_schema: True if should log the schema of the table.
@param with_histograms: True if should calculate and log histogram of the table data.
"""
try:
if not is_plugin_enabled("dbnd-postgres",
module_import="dbnd_postgres"):
logger.warning(
"Can't log postgres table: dbnd-postgres package is not installed\n"
"Help: pip install dbnd-postgres")
return
from dbnd_postgres import postgres_values
pg_table = postgres_values.PostgresTable(table_name, connection_string)
log_data(
table_name,
pg_table,
with_preview=with_preview,
with_schema=with_schema,
with_histograms=with_histograms,
)
except Exception:
logger.exception("Failed to log_pg_table")
def calculate_task_executor_type(submit_tasks, remote_engine, settings):
run_config = settings.run
parallel = run_config.parallel
task_executor_type = run_config.task_executor_type
if task_executor_type is None:
if is_airflow_enabled():
from dbnd_airflow.executors import AirflowTaskExecutorType
task_executor_type = AirflowTaskExecutorType.airflow_inprocess
else:
task_executor_type = TaskExecutorType.local
if is_airflow_enabled():
from dbnd_airflow.executors import AirflowTaskExecutorType
if parallel:
if task_executor_type == TaskExecutorType.local:
logger.warning(
"Auto switching to engine type '%s' due to parallel mode.",
AirflowTaskExecutorType.airflow_multiprocess_local,
)
task_executor_type = AirflowTaskExecutorType.airflow_multiprocess_local
if task_executor_type == AirflowTaskExecutorType.airflow_inprocess:
logger.warning(
"Auto switching to engine type '%s' due to parallel mode.",
AirflowTaskExecutorType.airflow_multiprocess_local,
)
task_executor_type = AirflowTaskExecutorType.airflow_multiprocess_local
if (task_executor_type
== AirflowTaskExecutorType.airflow_multiprocess_local
or task_executor_type
== AirflowTaskExecutorType.airflow_kubernetes):
if "sqlite" in settings.core.sql_alchemy_conn:
if settings.run.enable_concurent_sqlite:
logger.warning(
"You are running parallel execution on top of sqlite database! (see run.enable_concurent_sqlite)"
)
else:
# in theory sqlite can support a decent amount of parallelism, but in practice
# the way airflow works each process holds the db exlusively locked which leads
# to sqlite DB is locked exceptions
raise friendly_error.execute_engine.parallel_or_remote_sqlite(
"parallel")
if is_plugin_enabled("dbnd-docker"):
from dbnd_docker.kubernetes.kubernetes_engine_config import (
KubernetesEngineConfig, )
if (submit_tasks
and isinstance(remote_engine, KubernetesEngineConfig)
and run_config.enable_airflow_kubernetes):
if task_executor_type != AirflowTaskExecutorType.airflow_kubernetes:
logger.info(
"Using dedicated kubernetes executor for this run")
task_executor_type = AirflowTaskExecutorType.airflow_kubernetes
parallel = True
else:
if parallel:
logger.warning(
"Airflow is not installed, parallel mode is not supported")
all_executor_types = [TaskExecutorType.local]
if is_airflow_enabled():
from dbnd_airflow.executors import AirflowTaskExecutorType
all_executor_types.extend(AirflowTaskExecutorType.all())
if task_executor_type not in all_executor_types:
raise DatabandConfigError("Unsupported engine type %s" %
task_executor_type)
return task_executor_type, parallel