Add data-path, extend tests, update readme

Author: Daniel Molnar

README.md

@@ -139,6 +139,7 @@ composer-dev create \
   --port WEB_SERVER_PORT \
   --dags-path LOCAL_DAGS_PATH \
   --plugins-path LOCAL_PLUGINS_PATH \
+  --data-path LOCAL_DATA_PATH \
   --database DATABASE_ENGINE \
   LOCAL_ENVIRONMENT_NAME
 ```
@@ -152,6 +153,9 @@ Replace:
     located.
 - `LOCAL_PLUGINS_PATH` with the path to a local directory where the plugins
     files are located.
+- `LOCAL_DATA_PATH` with the path to a local directory where the data files are
+    located. If omitted, a `data` directory inside the environment directory is
+    used.
 - `DATABASE_ENGINE` with the database engine you wanted to use. You can use
     `sqlite` or `postgresql` (default).
 - `LOCAL_ENVIRONMENT_NAME` with the name of this local Airflow environment.
@@ -200,7 +204,8 @@ composer-dev create LOCAL_ENVIRONMENT_NAME \
     --project PROJECT_ID \
     --port WEB_SERVER_PORT \
     --dags-path LOCAL_DAGS_PATH \
-    --plugins-path LOCAL_PLUGINS_PATH
+    --plugins-path LOCAL_PLUGINS_PATH \
+    --data-path LOCAL_DATA_PATH
 ```
 
 Replace:
@@ -215,6 +220,9 @@ Replace:
     located.
 - `LOCAL_PLUGINS_PATH` with a path to a local directory where the plugins are
     located.
+- `LOCAL_DATA_PATH` with a path to a local directory where the data files are
+    located. If omitted, a `data` directory inside the environment directory is
+    used.
 
 Example:
 
@@ -302,17 +310,20 @@ composer-dev stop LOCAL_ENVIRONMENT_NAME
 
 ## Add and update DAGs
 
-DAGs and plugins are stored in the directories that you specified in the
-`--dags-path` and `--plugins-path` parameters respectively when you created
-your local Airflow environment. By default, these directories are
-`./composer/<local_environment_name>/dags` and
-`./composer/<local_environment_name>/plugins`.
+DAGs, plugins, and data are stored in the directories that you specified in the
+`--dags-path`, `--plugins-path`, and `--data-path` parameters respectively when
+you created your local Airflow environment. By default, these directories are
+`./composer/<local_environment_name>/dags`,
+`./composer/<local_environment_name>/plugins`, and
+`./composer/<local_environment_name>/data`. The data directory is mounted into
+the container at `/home/airflow/gcs/data`.
 
 You can get the directories used by your environment with the
 [`describe` command](#get-a-list-and-status-of-local-airflow-environments).
 
-To add and update DAGs and plugins, change files in these directories. You do
-not need to restart your local Airflow environment for changes to take effect.
+To add and update DAGs, plugins, and data, change files in these directories.
+You do not need to restart your local Airflow environment for changes to take
+effect.
 
 ## View local Airflow environment logs
 

composer_local_dev/cli.py

@@ -48,7 +48,12 @@
         },
         {
             "name": "Environment options",
-            "options": ["--web-server-port", "--dags-path", "--plugins-path"],
+            "options": [
+                "--web-server-port",
+                "--dags-path",
+                "--plugins-path",
+                "--data-path",
+            ],
         },
         {
             "name": "Container Memory and CPUs limit",
@@ -260,6 +265,13 @@ def _complete_environment(ctx, param, incomplete):
     metavar="PATH",
     type=click.Path(file_okay=False),
 )
+@click.option(
+    "--data-path",
+    help="Path to data folder. If it does not exist, it will be created.",
+    show_default="'data' directory in the environment directory",
+    metavar="PATH",
+    type=click.Path(file_okay=False),
+)
 @click.option(
     "--database-engine",
     "--database",
@@ -285,6 +297,7 @@ def create(
     database_engine: str,
     dags_path: Optional[pathlib.Path] = None,
     plugins_path: Optional[pathlib.Path] = None,
+    data_path: Optional[pathlib.Path] = None,
     container_memory_limit: Optional[str] = None,
     container_cpu_limit: Optional[str] = None,
 ):
@@ -341,6 +354,7 @@ def create(
             web_server_port=web_server_port,
             dags_path=dags_path,
             plugins_path=plugins_path,
+            data_path=data_path,
             database_engine=database_engine,
             memory_limit=container_memory_limit,
             cpu_count=container_cpu_limit,
@@ -354,6 +368,7 @@ def create(
             port=web_server_port,
             dags_path=dags_path,
             plugins_path=plugins_path,
+            data_path=data_path,
             database_engine=database_engine,
             memory_limit=container_memory_limit,
             cpu_count=container_cpu_limit,

composer_local_dev/constants.py

@@ -94,7 +94,8 @@ def choices(cls):
 
 1. You can put your DAGs in {dags_path}
 2. You can put your plugins in {plugins_path}
-3. Access Airflow at http://localhost:{port}
+3. You can put your data in {data_path}
+4. Access Airflow at http://localhost:{port}
 """
 # TODO: Fill source environment info
 DESCRIBE_ENV_MESSAGE = """
@@ -103,6 +104,7 @@ def choices(cls):
 Image version: {image_version}
 Dags directory: {dags_path}.
 Plugins directory: {plugins_path}.
+Data directory: {data_path}.
 The environment is using credentials from gcloud located at {gcloud_path}.
 """
 KUBECONFIG_PATH_MESSAGE = """
@@ -142,6 +144,15 @@ def choices(cls):
 PLUGINS_PATH_NOT_EXISTS_ERROR = (
     "Plugins path does not exist or is not a directory: {plugins_path}"
 )
+CREATING_DATA_PATH_WARN = (
+    "Data path '{data_path}' does not exist. It will be created."
+)
+DATA_PATH_NOT_PROVIDED_WARN = (
+    "No data directory provided, using default data directory."
+)
+DATA_PATH_NOT_EXISTS_ERROR = (
+    "Data path does not exist or is not a directory: {data_path}"
+)
 FAILED_TO_GET_DOCKER_PORT_WARN = (
     "Failed to retrieve used port from the Docker daemon, "
     "using port from the environment configuration."

composer_local_dev/environment.py

@@ -52,6 +52,7 @@ def get_image_mounts(
     kube_config_path: Optional[str],
     requirements: pathlib.Path,
     database_mounts: Dict[pathlib.Path, str],
+    data_path: Optional[str] = None,
 ) -> List[docker.types.Mount]:
     """
     Return list of docker volumes to be mounted inside container.
@@ -63,11 +64,14 @@ def get_image_mounts(
      - environment airflow sqlite db file location
      - database_mounts which contains the path for database mounts
     """
+    # Backwards compatibility: when no explicit data_path is given, the data
+    # directory lives inside the environment directory (the original behaviour).
+    data_path = data_path if data_path is not None else env_path / "data"
     mount_paths = {
         requirements: "composer_requirements.txt",
         dags_path: "gcs/dags/",
         plugins_path: "gcs/plugins/",
-        env_path / "data": "gcs/data/",
+        data_path: "gcs/data/",
         gcloud_config_path: ".config/gcloud",
         **database_mounts,
     }
@@ -373,6 +377,11 @@ def __init__(self, env_dir_path: pathlib.Path, port: Optional[int]):
             self.plugins_path = self.get_str_param("plugins_path")
         else:
             self.plugins_path = files.resolve_plugins_path(None, env_dir_path)
+        # Backwards compatibility: don't fail on missing data_path
+        if "data_path" in self.config:
+            self.data_path = self.get_str_param("data_path")
+        else:
+            self.data_path = files.resolve_data_path(None, env_dir_path)
         self.dag_dir_list_interval = self.parse_int_param(
             "dag_dir_list_interval", allowed_range=(0,)
         )
@@ -457,6 +466,7 @@ def __init__(
         location: str,
         dags_path: Optional[str],
         plugins_path: Optional[str] = None,
+        data_path: Optional[str] = None,
         dag_dir_list_interval: int = 10,
         database_engine: str = constants.DatabaseEngine.postgresql,
         memory_limit: Optional[str] = None,
@@ -488,6 +498,7 @@ def __init__(
         self.plugins_path = files.resolve_plugins_path(
             plugins_path, env_dir_path
         )
+        self.data_path = files.resolve_data_path(data_path, env_dir_path)
         self.dag_dir_list_interval = dag_dir_list_interval
         self.database_engine = database_engine
         self.is_database_sqlite3 = (
@@ -563,6 +574,7 @@ def load_from_config(cls, env_dir_path: pathlib.Path, port: Optional[int]):
             location=config.location,
             dags_path=config.dags_path,
             plugins_path=config.plugins_path,
+            data_path=config.data_path,
             dag_dir_list_interval=config.dag_dir_list_interval,
             port=config.port,
             database_engine=config.database_engine,
@@ -582,6 +594,7 @@ def from_source_environment(
         dags_path: Optional[str],
         plugins_path: Optional[str],
         database_engine: str,
+        data_path: Optional[str] = None,
         memory_limit: Optional[str] = None,
         cpu_count: Optional[int] = None,
     ):
@@ -606,6 +619,7 @@ def from_source_environment(
             location=location,
             dags_path=dags_path,
             plugins_path=plugins_path,
+            data_path=data_path,
             dag_dir_list_interval=10,
             port=web_server_port,
             pypi_packages=pypi_packages,
@@ -720,6 +734,7 @@ def write_environment_config_to_config_file(self):
             "composer_project_id": self.project_id,
             "dags_path": self.dags_path,
             "plugins_path": self.plugins_path,
+            "data_path": self.data_path,
             "dag_dir_list_interval": int(self.dag_dir_list_interval),
             "port": int(self.port),
             "database_engine": self.database_engine,
@@ -802,6 +817,7 @@ def create_docker_container(self):
             utils.resolve_kube_config_path(),
             self.requirements_file,
             db_mounts,
+            self.data_path,
         )
         db_vars = db_extras["env_vars"]
         default_vars = self.get_default_environment_variables(db_vars)
@@ -886,6 +902,7 @@ def create_db_docker_container(self):
             utils.resolve_kube_config_path(),
             self.requirements_file,
             db_mounts,
+            self.data_path,
         )
         db_vars = db_extras["env_vars"]
         db_ports = db_extras["ports"]
@@ -967,7 +984,10 @@ def create(self):
         assert_image_exists(self.image_version)
         self.assert_valid_environment_options()
         files.create_environment_directories(
-            self.env_dir_path, self.dags_path, self.plugins_path
+            self.env_dir_path,
+            self.dags_path,
+            self.plugins_path,
+            self.data_path,
         )
         self.create_database_files(skip_if_exist=False)
         self.write_environment_config_to_config_file()
@@ -1104,6 +1124,7 @@ def start(self, assert_not_running=True):
         self.assert_requirements_exist()
         files.assert_dag_path_exists(self.dags_path)
         files.assert_plugins_path_exists(self.plugins_path)
+        files.assert_data_path_exists(self.data_path)
 
         self.create_database_files()
         db_path = (
@@ -1153,6 +1174,7 @@ def print_start_message(self):
                 env_name=self.name,
                 dags_path=self.dags_path,
                 plugins_path=self.plugins_path,
+                data_path=self.data_path,
                 port=self.port,
             )
         )
@@ -1278,6 +1300,7 @@ def prepare_env_description(self, env_status: str) -> str:
                 image_version=self.image_version,
                 dags_path=self.dags_path,
                 plugins_path=self.plugins_path,
+                data_path=self.data_path,
                 gcloud_path=utils.resolve_gcloud_config_path(),
             )
             + (

composer_local_dev/errors.py

@@ -191,6 +191,15 @@ def __init__(self, plugins_path):
         )
 
 
+class DataPathNotExistError(ComposerCliError):
+    """Data path does not exist or is not a directory."""
+
+    def __init__(self, data_path):
+        super().__init__(
+            constants.DATA_PATH_NOT_EXISTS_ERROR.format(data_path=data_path)
+        )
+
+
 def catch_exceptions(func=None):
     """
     Catch exceptions and print user friendly message for common issues.

composer_local_dev/files.py

@@ -126,23 +126,46 @@ def resolve_plugins_path(
     return str(plugins_path.resolve())
 
 
+def resolve_data_path(data_path: Optional[str], env_dir: pathlib.Path) -> str:
+    """
+    Provides and validates path to the data directory.
+    If ``data_path`` is None, the path is constructed from ``env_dir`` path
+    and ``data`` directory.
+    If ``data_path`` is not None, but it does not exist, a warning is raised.
+
+    Returns absolute ``data_path`` path.
+    """
+    if data_path is None:
+        console.get_console().print(constants.DATA_PATH_NOT_PROVIDED_WARN)
+        data_path = env_dir / "data"
+    else:
+        data_path = pathlib.Path(data_path)
+    return str(data_path.resolve())
+
+
 def create_environment_directories(
-    env_dir: pathlib.Path, dags_path: str, plugins_path: str
+    env_dir: pathlib.Path,
+    dags_path: str,
+    plugins_path: str,
+    data_path: Optional[str] = None,
 ):
     """
     Create environment directories (overwriting existing ones).
     Environment directory is a directory which contains configuration files for
     composer local environment and files used by environment such as
     requirements.txt file, dags, data and plugins directories.
     """
-    data_dir = "data"
-    LOG.info(
-        "Creating environment directory %s in " "%s environment directory.",
-        data_dir,
-        env_dir,
+    # Backwards compatibility: when no explicit data_path is given, the data
+    # directory lives inside the environment directory (the original behaviour).
+    data_path = (
+        pathlib.Path(data_path) if data_path is not None else env_dir / "data"
     )
     env_dir.mkdir(exist_ok=True, parents=True)
-    (env_dir / data_dir).mkdir(exist_ok=True)
+    if not data_path.is_dir():
+        console.get_console().print(
+            constants.CREATING_DATA_PATH_WARN.format(data_path=data_path)
+        )
+        data_path.mkdir(parents=True)
 
     dags_path = pathlib.Path(dags_path)
     if not dags_path.is_dir():
@@ -259,3 +282,10 @@ def assert_plugins_path_exists(path: str) -> None:
     if pathlib.Path(path).is_dir():
         return
     raise errors.PluginsPathNotExistError(path)
+
+
+def assert_data_path_exists(path: str) -> None:
+    """Raise an error if data path does not point to existing directory."""
+    if pathlib.Path(path).is_dir():
+        return
+    raise errors.DataPathNotExistError(path)