1. # coding: utf8
    2. import os
    3. import re
    4. from typing import Any, Dict, Optional
    6. from airflow.configuration import conf
    7. from airflow.exceptions import AirflowException
    8. from airflow.models import BaseOperator
    9. from custom_hooks._impala import ImpalaHook
    10. from airflow.utils import operator_helpers
    11. from airflow.utils.operator_helpers import context_to_airflow_vars
    14. class ImpalaOperator(BaseOperator):
    16.     template_fields = ('hql', 'impala_conn_id')
    17.     template_ext = ('.hql', '.sql')
    18.     ui_color = '#f0e4ec'
    21.     def __init__(
    22.         self,
    23.         *,
    24.         hql: str = None,
    25.         impala_conn_id: str = 'impala_default',
    26.         db: str = 'default',
    27.         hql_file: Optional[str] = None,
    28.         hql_name: Optional[str] = None,
    29.         impalaconf_jinja_translate: bool = False,
    30.         run_as_owner: bool = False,
    31.         **kwargs: Any,
    32.         ) -> None:
    33.         super().__init__(**kwargs)
    34.         self.hql = hql
    35.         self.impala_conn_id = impala_conn_id
    36.         self.db = db
    37.         self.hql_file = hql_file
    38.         self.hql_name = hql_name
    39.         self.impalaconf_jinja_translate = impalaconf_jinja_translate
    40.         self.run_as = None
    42.         if run_as_owner:
    43.             self.run_as = self.dag.owner
    45.         # assigned lazily - just for consistency we can create the attribute with a
    46.         # `None` initial value, later it will be populated by the execute method.
    47.         # This also makes `on_kill` implementation consistent since it assumes `self.hook`
    48.         # is defined.
    49.         self.hook: Optional[ImpalaHook] = None
    52.     def get_hook(self) -> ImpalaHook:
    53.         """Get Impala hook"""
    54.         return ImpalaHook()
    57.     def _parse_sqlfile(self):
    58.         with open(self.hql_file, 'r', encoding="utf8") as f:
    60.             if self.impalaconf_jinja_translate:
    61.                 file_content = re.sub(r"(\$\{([ batch_date|run_date|BATCH_DATE|RUN_DATE]*)\})", 
    62.                                       r"{{ ds }}", f.read())
    63.             else:
    64.                 file_content = f.read()
    66.             pattern = r"--\[(.*?)\](.*?)\n--\[end\].*?"
    67.             sqls_dict = dict([(k, v) for k, v in re.findall(pattern, file_content, re.S) if v != ""])
    69.             if self.hql_name:
    70.                 return sqls_dict[self.hql_name]
    71.             elif len(sqls_dict.keys()) == 1:
    72.                 key = list(sqls_dict.keys())[0]
    73.                 return sqls_dict[key]
    74.             else:
    75.                 raise AirflowException("You must specify `hql_name` when `hql_file` is defined!")
    78.     def prepare_template(self) -> None:
    79.         if self.impalaconf_jinja_translate:            
    80.             if self.hql_file is not None:
    81.                 self.hql = self._parse_sqlfile()
    82.             else:
    83.                 self.hql = re.sub(r"(\$\{([ a-zA-Z0-9_|batch_date|run_date|BATCH_DATE|RUN_DATE|]*)\})", 
    84.                                   r"{{ ds }}", self.hql)
    85.         else:
    86.             if self.hql_file is not None:
    87.                 self.hql = self._parse_sqlfile()
    90.     def execute(self, context: Dict[str, Any]) -> None:
    91.         # self.log.info('Executing: %s', self.hql)
    92.         self.hook = self.get_hook()
    93.         self.hook.run_hql(self.hql)
    94.         self.log.info("Executed successfully!")
    97.     def dry_run(self) -> None:
    98.         # Reset airflow environment variables to prevent
    99.         # existing env vars from impacting behavior.
    100.         self.clear_airflow_vars()
    102.         self.hook = self.get_hook()
    103.         self.hook.test_hql(hql=self.hql)
    106.     def on_kill(self) -> None:
    107.         if self.hook:
    108.             self.hook.kill()
    111.     def clear_airflow_vars(self) -> None:
    112.         """Reset airflow environment variables to prevent existing ones from impacting behavior."""
    113.         blank_env_vars = {
    114.             value['env_var_format']: '' for value in operator_helpers.AIRFLOW_VAR_NAME_FORMAT_MAPPING.values()
    115.         }
    116.         os.environ.update(blank_env_vars)