Stop with savepoint if query has been removed (#19)

Author: grzegorz8

README.md

@@ -179,9 +179,7 @@ _Decision_: Flink jobs are deployed on a shared Flink Session Cluster.
 - All necessary Flink dependencies are available on the classpath on the EMR master node.
 - Python virtualenv with all necessary python dependencies is created on each EMR node.
   ```bash
-  python3 -m pip install \
-      -r flink-sql-runner/deployment-scripts/jobs-deployment/requirements.txt \
-      -r flink-sql-runner/python/requirements.txt
+  python3 -m pip install flink-sql-runner
   ```
 - An S3 bucket for storing YAML manifests file is created. See [Job manifest](#job-manifest) section.
 

flink_sql_runner/deploy.py

@@ -2,14 +2,16 @@
 import argparse
 import logging
 import os
-import tempfile
 from typing import List
 
 import yaml
 
-from flink_sql_runner.deploy_job import EmrJobRunner, JinjaTemplateResolver
-from flink_sql_runner.flink_clients import (FlinkStandaloneClusterRunner,
+from flink_sql_runner.deploy_job import FlinkJobRunner, JinjaTemplateResolver
+from flink_sql_runner.flink_clients import (FlinkCli,
+                                            FlinkStandaloneClusterRunner,
                                             FlinkYarnRunner)
+from flink_sql_runner.job_configuration import JobConfiguration
+from flink_sql_runner.manifest import ManifestManager
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
 
@@ -58,58 +60,117 @@ def parse_args():
     return parser.parse_known_args()
 
 
-def list_query_files(base_path: str) -> List[str]:
-    result = []
-    for root, dirs, files in os.walk(base_path):
-        for f in files:
-            if f.endswith(".yaml"):
-                result.append(os.path.abspath(os.path.join(root, f)))
-    return result
-
-
-def read_config(query_file, template_file):
-    # FIXME: refactor variables resolutions and yaml merge
-    with open(query_file) as qf:
-        query_specification = yaml.load(qf, yaml.FullLoader)
-        if "sql" in query_specification:
-            query_specification["sql"] = query_specification["sql"].replace("\n", " ")
-        with open(template_file) as tf:
-            raw_defaults = tf.read().format(job_name=query_specification["name"])
-            default_config = yaml.safe_load(raw_defaults)
-            final_flink_props = {
-                **default_config["flinkProperties"],
-                **query_specification["flinkProperties"],
-            }
-            final_config = {**default_config, **query_specification}
-            final_config["flinkProperties"] = final_flink_props
-            final_config["flinkProperties"]["pipeline.name"] = query_specification["name"]
-            logging.info(f"Final configuration:\n{final_config}")
-            return final_config
+class FlinkRunner:
+    def __init__(self,
+                 queries_base_path: str,
+                 table_definition_path: str,
+                 pyflink_runner_dir: str,
+                 template_file: str,
+                 pyexec_path: str,
+                 flink_cli_runner: FlinkCli,
+                 manifest_manager: ManifestManager,
+                 jinja_template_resolver: JinjaTemplateResolver,
+                 passthrough_args,
+                 ):
+        self.queries_base_path = queries_base_path
+        self.table_definition_path = table_definition_path
+        self.pyflink_runner_dir = pyflink_runner_dir
+        self.template_file = template_file
+        self.pyexec_path = pyexec_path
+        self.flink_cli_runner = flink_cli_runner
+        self.manifest_manager = manifest_manager
+        self.jinja_template_resolver = jinja_template_resolver
+        self.passthrough_args = passthrough_args
+
+    def run(self):
+        query_files = self.__list_query_files(self.queries_base_path)
+        existing_query_manifests = self.manifest_manager.list_manifests()
+        queries_to_cancel = self.__get_queries_to_cancel(query_files, existing_query_manifests)
+
+        logging.info("Queries to start or update: {}", query_files)
+        logging.info("Queries to cancel: {}", queries_to_cancel)
+
+        self.__add_or_update_jobs(query_files)
+        self.__cancel_deleted_jobs(queries_to_cancel)
+
+    def __get_queries_to_cancel(self, query_files: List[str], existing_query_manifests: List[str]) -> List[str]:
+        running_query_names = [n.split("/")[-1] for n in existing_query_manifests]
+        current_query_names = [n.split("/")[-1] for n in query_files]
+        difference = set(running_query_names).difference(set(current_query_names))
+        return [q.replace(".yaml", "") for q in difference]
+
+    def __add_or_update_jobs(self, query_files):
+        for query_file in query_files:
+            final_config = self.__read_config(query_file, self.template_file)
+            new_job_conf = JobConfiguration(final_config)
+            FlinkJobRunner(
+                job_name=new_job_conf.get_name(),
+                new_job_conf=new_job_conf,
+                pyflink_runner_dir=self.pyflink_runner_dir,
+                table_definition_paths=self.table_definition_path,
+                pyexec_path=self.pyexec_path,
+                flink_cli_runner=self.flink_cli_runner,
+                jinja_template_resolver=self.jinja_template_resolver,
+                manifest_manager=self.manifest_manager,
+                passthrough_args=self.passthrough_args,
+            ).run()
+
+    def __cancel_deleted_jobs(self, queries_to_cancel):
+        for query_to_remove in queries_to_cancel:
+            FlinkJobRunner(
+                job_name=query_to_remove,
+                new_job_conf=None,
+                pyflink_runner_dir=self.pyflink_runner_dir,
+                table_definition_paths=self.table_definition_path,
+                pyexec_path=self.pyexec_path,
+                flink_cli_runner=self.flink_cli_runner,
+                jinja_template_resolver=self.jinja_template_resolver,
+                manifest_manager=self.manifest_manager,
+                passthrough_args=self.passthrough_args,
+            ).run()
+
+    def __list_query_files(self, base_path: str) -> List[str]:
+        result = []
+        for root, dirs, files in os.walk(base_path):
+            for f in files:
+                if f.endswith(".yaml"):
+                    result.append(os.path.abspath(os.path.join(root, f)))
+        return result
+
+    def __read_config(self, query_file, template_file):
+        # FIXME: refactor variables resolutions and yaml merge
+        with open(query_file) as qf:
+            query_specification = yaml.load(qf, yaml.FullLoader)
+            if "sql" in query_specification:
+                query_specification["sql"] = query_specification["sql"].replace("\n", " ")
+            with open(template_file) as tf:
+                raw_defaults = tf.read().format(job_name=query_specification["name"])
+                default_config = yaml.safe_load(raw_defaults)
+                final_flink_props = {
+                    **default_config["flinkProperties"],
+                    **query_specification["flinkProperties"],
+                }
+                final_config = {**default_config, **query_specification}
+                final_config["flinkProperties"] = final_flink_props
+                logging.info(f"Final configuration:\n{final_config}")
+                return final_config
 
 
 if __name__ == "__main__":
     args, passthrough_args = parse_args()
-    query_files = list_query_files(args.path)
-
-    for query_file in query_files:
-        final_config = read_config(query_file, args.template_file)
-        query_name = final_config["name"]
-        with tempfile.NamedTemporaryFile(mode="w+t", prefix=query_name, suffix=".yaml") as tmp:
-            yaml.dump(final_config, tmp)
-            flink_cli_runner = (
-                FlinkYarnRunner()
-                if args.deployment_target == "yarn"
-                else FlinkStandaloneClusterRunner(args.jobmanager_address)
-            )
-            jinja_template_resolver = JinjaTemplateResolver()
-            EmrJobRunner(
-                tmp.name,
-                args.pyflink_runner_dir,
-                args.external_job_config_bucket,
-                args.external_job_config_prefix,
-                args.table_definition_path,
-                args.pyexec_path,
-                flink_cli_runner,
-                jinja_template_resolver,
-                passthrough_args,
-            ).run()
+
+    FlinkRunner(
+        queries_base_path=args.path,
+        table_definition_path=args.table_definition_path,
+        pyflink_runner_dir=args.pyflink_runner_dir,
+        template_file=args.template_file,
+        pyexec_path=args.pyexec_path,
+        flink_cli_runner=(
+            FlinkYarnRunner()
+            if args.deployment_target == "yarn"
+            else FlinkStandaloneClusterRunner(args.jobmanager_address)
+        ),
+        manifest_manager=ManifestManager(args.external_job_config_bucket, args.external_job_config_prefix),
+        jinja_template_resolver=JinjaTemplateResolver(),
+        passthrough_args=passthrough_args,
+    ).run()

flink_sql_runner/deploy_job.py

@@ -9,13 +9,14 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import yaml
-from jinja2 import Environment, FileSystemLoader
 
 from flink_sql_runner.flink_clients import (FlinkCli,
                                             FlinkStandaloneClusterRunner,
                                             FlinkYarnRunner)
+from flink_sql_runner.jinja import JinjaTemplateResolver
 from flink_sql_runner.job_configuration import JobConfiguration
-from flink_sql_runner.s3 import get_content, get_latest_object, upload_content
+from flink_sql_runner.manifest import ManifestManager
+from flink_sql_runner.s3 import get_latest_object
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
 
@@ -66,87 +67,69 @@ def parse_args():
     return parser.parse_known_args()
 
 
-class JinjaTemplateResolver(object):
-    def resolve(
-            self,
-            template_dir: str,
-            template_file: str,
-            vars: Dict[str, str],
-            output_file_path: str,
-    ) -> None:
-        environment = Environment(loader=FileSystemLoader(template_dir))
-        template = environment.get_template(template_file)
-        content = template.render(**vars)
-        with open(output_file_path, mode="w", encoding="utf-8") as run_file:
-            run_file.truncate()
-            run_file.write(content)
-
-
-class EmrJobRunner(object):
+class FlinkJobRunner(object):
     def __init__(
             self,
-            job_config_path: str,
+            job_name: str,
+            new_job_conf: Optional[JobConfiguration],
             pyflink_runner_dir: str,
-            external_job_config_bucket: str,
-            external_job_config_prefix: str,
             table_definition_paths: str,
             pyexec_path: str,
             flink_cli_runner: FlinkCli,
             jinja_template_resolver: JinjaTemplateResolver,
+            manifest_manager: ManifestManager,
             passthrough_args: List[str],
     ):
-        self.job_config_path = job_config_path
+        self.job_name = job_name
+        self.new_job_conf = new_job_conf
         self.pyflink_runner_dir = pyflink_runner_dir
-        self.external_job_config_bucket = external_job_config_bucket
-        self.external_job_config_prefix = external_job_config_prefix
         self.table_definition_paths = table_definition_paths
         self.pyexec_path = pyexec_path
         self.pyclientexec_path = pyexec_path
         self.flink_cli_runner = flink_cli_runner
         self.jinja_template_resolver = jinja_template_resolver
+        self.manifest_manager = manifest_manager
         self.passthrough_args = passthrough_args
-        self.new_job_conf = JobConfiguration(self.__read_config(job_config_path))
 
     def run(self) -> None:
-        logging.info(f"Deploying '{self.new_job_conf.get_name()}'.")
+        if self.new_job_conf is None:
+            logging.info(f"Deleting job '{self.job_name}'.")
+            job_manifest = self.manifest_manager.fetch_job_manifest(self.job_name)
+            if job_manifest is None:
+                raise ValueError(f"Job manifest for {self.job_name} not found.")
+            self.__stop_with_savepoint(job_manifest)
+            return
+
+        logging.info(f"Deploying '{self.job_name}'.")
         if self.new_job_conf.is_sql():
             logging.info(f"Deploying query: |{self.new_job_conf.get_sql()}|")
         else:
             logging.info(f"Deploying code:\n{self.new_job_conf.get_code()}")
 
-        external_config = self.__fetch_job_manifest(
-            self.external_job_config_bucket,
-            self.external_job_config_prefix,
-            self.new_job_conf.get_name(),
-        )
+        external_config = self.manifest_manager.fetch_job_manifest(self.job_name)
         logging.info(f"External config:\n{external_config}")
 
         if not external_config:
             # The job manifest did not exist. Starting a newly created job.
             self.__start_new_job(self.new_job_conf)
-            self.__upload_job_manifest(self.new_job_conf)
+            self.manifest_manager.upload_job_manifest(self.new_job_conf)
         elif external_config and not self.__has_job_manifest_changed(external_config, self.new_job_conf):
             # The job manifest has not been modified. There is no need to restart the job. Just ensure it's running.
-            if self.__is_job_running(self.new_job_conf.get_name()):
+            if self.__is_job_running(self.job_name):
                 logging.info("Job manifest has not changed. Skipping job restart.")
             else:
                 self.__start_job_with_unchanged_query(external_config, self.new_job_conf)
         else:
             # The job manifest has been modified. Job needs to be restarted.
-            if self.__is_job_running(self.new_job_conf.get_name()):
+            if self.__is_job_running(self.job_name):
                 # Stop the job using the old config (query-version in particular).
                 self.__stop_with_savepoint(external_config)
 
             if external_config and not self.__has_job_definition_changed(external_config, self.new_job_conf):
                 self.__start_job_with_unchanged_query(external_config, self.new_job_conf)
             else:
                 self.__start_new_job_with_changed_query(external_config, self.new_job_conf)
-            self.__upload_job_manifest(self.new_job_conf)
-
-    @staticmethod
-    def __read_config(config_file: str):
-        with open(config_file) as qf:
-            return yaml.load(qf, yaml.FullLoader)
+            self.manifest_manager.upload_job_manifest(self.new_job_conf)
 
     def __is_job_running(self, job_name: str) -> bool:
         return self.flink_cli_runner.is_job_running(job_name)
@@ -169,12 +152,6 @@ def __start_new_job_with_changed_query(self, external_config, job_conf):
         job_conf.set_meta_query_create_timestamp(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
         self.__start_with_clean_state(job_conf)
 
-    def __upload_job_manifest(self, job_conf):
-        upload_path = os.path.join(self.external_job_config_prefix, f"{job_conf.get_name()}.yaml")
-        logging.info(f"Uploading the new config file to 's3://{self.external_job_config_bucket}/{upload_path}'.")
-        upload_content(yaml.dump(job_conf.to_dict()), self.external_job_config_bucket, upload_path)
-        logging.info("The config file has been uploaded.")
-
     def __stop_with_savepoint(self, job_conf: JobConfiguration) -> None:
         job_id = self.flink_cli_runner.get_job_id(job_conf.get_name())
         savepoint_path = os.path.join(job_conf.get_flink_savepoints_dir(), job_conf.get_meta_query_version_str())
@@ -316,12 +293,6 @@ def __find_latest_state_internal(
             logging.info(f"State found at '{state_path}'.")
             return state_path, last_created_ts
 
-    def __fetch_job_manifest(self, bucket_name: str, prefix: str, job_name: str) -> Optional[JobConfiguration]:
-        object_key = os.path.join(prefix, f"{job_name}.yaml")
-        logging.info(f"Looking for config at s3://{bucket_name}/{object_key}.")
-        raw_manifest = get_content(bucket_name, object_key)
-        return JobConfiguration(yaml.safe_load(raw_manifest)) if raw_manifest else None
-
     def __has_job_manifest_changed(self, old_job_conf: JobConfiguration, new_job_conf: JobConfiguration) -> bool:
         return self.__has_job_definition_changed(old_job_conf, new_job_conf) or self.__have_flink_properties_changed(
             old_job_conf, new_job_conf
@@ -352,21 +323,26 @@ def __escape_query(query: str) -> str:
         return query.replace("`", "\\`")
 
 
+def read_config(config_file: str):
+    with open(config_file) as qf:
+        return yaml.load(qf, yaml.FullLoader)
+
+
 if __name__ == "__main__":
     args, passthrough_args = parse_args()
-    flink_cli_runner = (
-        FlinkYarnRunner() if args.deployment_target == "yarn" else FlinkStandaloneClusterRunner(args.jobmanager_address)
-    )
-    jinja_template_resolver = JinjaTemplateResolver()
-
-    EmrJobRunner(
-        args.job_config_path,
-        args.pyflink_runner_dir,
-        args.external_job_config_bucket,
-        args.external_job_config_prefix,
-        args.base_output_path,
-        args.pyexec_path,
-        flink_cli_runner,
-        jinja_template_resolver,
-        passthrough_args,
+    configuration = JobConfiguration(read_config(args.job_config_path))
+    FlinkJobRunner(
+        job_name=configuration.get_name(),
+        new_job_conf=configuration,
+        pyflink_runner_dir=args.pyflink_runner_dir,
+        table_definition_paths=args.base_output_path,
+        pyexec_path=args.pyexec_path,
+        flink_cli_runner=(
+            FlinkYarnRunner()
+            if args.deployment_target == "yarn"
+            else FlinkStandaloneClusterRunner(args.jobmanager_address)
+        ),
+        jinja_template_resolver=JinjaTemplateResolver(),
+        manifest_manager=ManifestManager(args.external_job_config_bucket, args.external_job_config_prefix),
+        passthrough_args=passthrough_args,
     ).run()