Docker Calculation fix + Reset DB command added (#3459)

docker/run_local_docker.sh

@@ -16,7 +16,7 @@ cd ../ && mvn -DskipTests clean package
 echo "Prepare Docker volume for the operators"
 cd docker/local-metadata
 echo "Starting Local Docker Containers"
-docker-compose down && docker-compose up --build -d
+docker compose down && docker compose up --build -d
 until curl -s -f -o /dev/null "http://localhost:8585/api/v1/tables/name/bigquery_gcp.shopify.fact_sale"; do
     printf '.'
     sleep 2

ingestion/src/metadata/cli/docker.py

@@ -9,135 +9,156 @@ from datetime import timedelta
 import click
 import requests as requests
 
+from metadata.generated.schema.entity.data.table import Table
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
+
 logger = logging.getLogger(__name__)
-calc_gb = 1024 * 1024 * 1000
+calc_gb = 1024 * 1024 * 1024
 min_memory_limit = 6 * calc_gb
 
 
-def run_docker(start, stop, pause, resume, clean, file_path, env_file_path):
+def start_docker(docker, start_time, file_path):
+    logger.info("Running docker compose for OpenMetadata..")
+    click.secho("It may take some time on the first run", fg="bright_yellow")
+    if file_path:
+        docker.compose.up(detach=True, build=True)
+    else:
+        docker.compose.up(detach=True)
+
+    logger.info(
+        "Ran docker compose for OpenMetadata successfully.\nWaiting for ingestion to complete.."
+    )
+    metadata_config = MetadataServerConfig.parse_obj(
+        {
+            "api_endpoint": "http://localhost:8585/api",
+            "auth_provider_type": "no-auth",
+        }
+    )
+    logging.getLogger("metadata.ingestion.ometa.ometa_api").disabled = True
+    ometa_client = OpenMetadata(metadata_config)
+    while True:
+        try:
+            resp = ometa_client.get_by_name(
+                entity=Table, fqdn="bigquery_gcp.shopify.dim_customer"
+            )
+            if not resp:
+                raise Exception("Error")
+            break
+        except Exception:
+            sys.stdout.write(".")
+            sys.stdout.flush()
+            time.sleep(5)
+    logging.getLogger("metadata.ingestion.ometa.ometa_api").disabled = False
+    elapsed = time.time() - start_time
+    logger.info(
+        f"Time took to get OpenMetadata running: {str(timedelta(seconds=elapsed))}"
+    )
+    click.secho(
+        "\n✅ OpenMetadata is up and running",
+        fg="bright_green",
+    )
+    click.secho(
+        """\nOpen http://localhost:8585 in your browser to access OpenMetadata..
+        \nTo checkout Ingestion via Airflow, go to http://localhost:8080 \n(username: admin, password: admin)
+        """,
+        fg="bright_blue",
+    )
+    click.secho(
+        """We are available on Slack , https://slack.open-metadata.org/ . Reach out to us if you have any questions.
+        \nIf you like what we are doing, please consider giving us a star on github at https://github.com/open-metadata/OpenMetadata. 
+It helps OpenMetadata reach wider audience and helps our community.\n""",
+        fg="bright_magenta",
+    )
+
+
+def env_file_check(env_file_path):
+    if env_file_path is not None:
+        if env_file_path == "":
+            raise ValueError("Please provide path to env file")
+        else:
+            logger.info(f"Using env file from {env_file_path}")
+            return pathlib.Path(env_file_path)
+
+
+def file_path_check(file_path):
+    if file_path is None:
+        docker_compose_file_path = (
+            pathlib.Path(tempfile.gettempdir()) / "docker-compose.yml"
+        )
+        if not docker_compose_file_path.exists():
+            logger.info(
+                "Downloading latest docker compose file from openmetadata repository..."
+            )
+            r = requests.get(
+                "https://raw.githubusercontent.com/open-metadata/OpenMetadata/main/docker/metadata/docker-compose.yml"
+            )
+            with open(docker_compose_file_path, "wb") as docker_compose_file_handle:
+                docker_compose_file_handle.write(r.content)
+                docker_compose_file_handle.close()
+    else:
+        if file_path == "":
+            raise ValueError("Please Provide Path to local docker-compose.yml file")
+        else:
+            logger.info(f"Using docker compose file from {file_path}")
+            docker_compose_file_path = pathlib.Path(file_path)
+    return docker_compose_file_path
+
+
+def run_docker(start, stop, pause, resume, clean, file_path, env_file_path, reset_db):
     try:
         from python_on_whales import DockerClient
 
-        from metadata.ingestion.ometa.ometa_api import OpenMetadata
-        from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
-
         docker = DockerClient(compose_project_name="openmetadata", compose_files=[])
 
-        logger.info("Checking if docker compose is installed....")
+        logger.info("Checking if docker compose is installed..")
         if not docker.compose.is_installed():
             raise Exception("Docker Compose CLI is not installed on the system.")
 
         docker_info = docker.info()
 
-        logger.info("Checking if docker service is running....")
+        logger.info("Checking if docker service is running..")
         if not docker_info.id:
             raise Exception("Docker Service is not up and running.")
 
-        logger.info("Checking openmetadata memory constraints....")
+        logger.info("Checking openmetadata memory constraints..")
         if docker_info.mem_total < min_memory_limit:
             raise MemoryError
 
         # Check for -f <Path>
         start_time = time.time()
-        if file_path is None:
+        docker_compose_file_path = file_path_check(file_path)
-            docker_compose_file_path = (
+        env_file = env_file_check(env_file_path)
-                pathlib.Path(tempfile.gettempdir()) / "docker-compose.yml"
-            )
-            if not docker_compose_file_path.exists():
-                logger.info(
-                    "Downloading latest docker compose file from openmetadata repository..."
-                )
-                r = requests.get(
-                    "https://raw.githubusercontent.com/open-metadata/OpenMetadata/main/docker/metadata/docker-compose.yml"
-                )
-                with open(docker_compose_file_path, "wb") as docker_compose_file_handle:
-                    docker_compose_file_handle.write(r.content)
-        else:
-            if file_path == "":
-                raise ValueError("Please Provide Path to local docker-compose.yml file")
-            else:
-                logger.info(f"Using docker compose file from {file_path}")
-                docker_compose_file_path = pathlib.Path(file_path)
-
-        env_file = None
-        if env_file_path is not None:
-            if env_file_path == "":
-                raise ValueError("Please provide path to env file")
-            else:
-                logger.info(f"Using env file from {env_file_path}")
-                env_file = pathlib.Path(env_file_path)
-
         # Set up Docker Client Config with docker compose file path
         docker = DockerClient(
             compose_project_name="openmetadata",
             compose_files=[docker_compose_file_path],
             compose_env_file=env_file,
         )
-
         if start:
-            logger.info("Running docker compose for Open Metadata....")
+            start_docker(docker, start_time, file_path)
-            if file_path:
-                docker.compose.up(detach=True, build=True)
-            else:
-                docker.compose.up(detach=True)
-
-            logger.info(
-                "Docker Compose for Open Metadata successful. Waiting for ingestion to complete...."
-            )
-            metadata_config = MetadataServerConfig.parse_obj(
-                {
-                    "api_endpoint": "http://localhost:8585/api",
-                    "auth_provider_type": "no-auth",
-                }
-            )
-
-            ometa_client = OpenMetadata(metadata_config).client
-            while True:
-                try:
-                    ometa_client.get(f"/tables/name/bigquery_gcp.shopify.dim_customer")
-                    break
-                except Exception as err:
-                    sys.stdout.flush()
-                    time.sleep(5)
-            elapsed = time.time() - start_time
-            logger.info(
-                f"Time took to get OpenMetadata running: {str(timedelta(seconds=elapsed))}"
-            )
-            click.secho(
-                "\n✅ OpenMetadata is up and running",
-                fg="bright_green",
-            )
-            click.secho(
-                """\nOpen http://localhost:8585 in your browser to access OpenMetadata..
-                \nTo checkout Ingestion via Airflow, go to http://localhost:8080 \n(username: admin, password: admin)
-                """,
-                fg="bright_blue",
-            )
-            click.secho(
-                """We are available on Slack , https://slack.open-metadata.org/ . Reach out to us if you have any questions.
-                \nIf you like what we are doing, please consider giving us a star on github at https://github.com/open-metadata/OpenMetadata. 
-        It helps OpenMetadata reach wider audience and helps our community.\n""",
-                fg="bright_magenta",
-            )
         if pause:
-            logger.info("Pausing docker compose for Open Metadata....")
+            logger.info("Pausing docker compose for OpenMetadata..")
             docker.compose.pause()
-            logger.info("Pausing docker compose for Open Metadata Successful.")
+            logger.info("Pausing docker compose for OpenMetadata successful.")
         if resume:
-            logger.info("Resuming docker compose for Open Metadata....")
+            logger.info("Resuming docker compose for OpenMetadata..")
             docker.compose.unpause()
-            logger.info("Resuming docker compose for Open Metadata Successful.")
+            logger.info("Resuming docker compose for OpenMetadata Successful.")
         if stop:
-            logger.info("Stopping docker compose for Open Metadata....")
+            logger.info("Stopping docker compose for OpenMetadata..")
             docker.compose.stop()
-            logger.info("docker compose for Open Metadata stopped successfully.")
+            logger.info("Docker compose for OpenMetadata stopped successfully.")
+        if reset_db:
+
+            reset_db_om(docker)
         if clean:
             logger.info(
-                "Stopping docker compose for Open Metadata and removing images, networks, volumes...."
+                "Stopping docker compose for OpenMetadata and removing images, networks, volumes.."
             )
             docker.compose.down(remove_orphans=True, remove_images="all", volumes=True)
             logger.info(
-                "Stopped docker compose for Open Metadata and removing images, networks, volumes."
+                "Stopped docker compose for OpenMetadata and removing images, networks, volumes."
             )
             if file_path is None:
                 docker_compose_file_path.unlink()
@@ -145,10 +166,29 @@ def run_docker(start, stop, pause, resume, clean, file_path, env_file_path):
     except MemoryError:
         click.secho(
             f"Please Allocate More memory to Docker.\nRecommended: 6GB\nCurrent: "
-            f"{round(float(dict(docker_info).get('mem_total')) / calc_gb, 2)}",
+            f"{round(float(dict(docker_info).get('mem_total')) / calc_gb)}",
             fg="red",
         )
     except Exception as err:
         logger.debug(traceback.format_exc())
         logger.debug(traceback.print_exc())
         click.secho(str(err), fg="red")
+
+
+def reset_db_om(docker):
+    if docker.container.inspect("openmetadata_server").state.running:
+        click.secho(
+            f"Resetting OpenMetadata.\nThis will clear out all the data",
+            fg="red",
+        )
+        docker.container.execute(
+            container="openmetadata_server",
+            tty=True,
+            command=[
+                "/bin/bash",
+                "-c",
+                "./openmetadata-*/bootstrap/bootstrap_storage.sh drop-create-all",
+            ],
+        )
+    else:
+        click.secho("OpenMetadata Instance is not up and running", fg="yellow")

ingestion/src/metadata/cmd.py

@@ -197,13 +197,16 @@ def report(config: str) -> None:
     type=click.Path(exists=True, dir_okay=False),
     required=False,
 )
-def docker(start, stop, pause, resume, clean, file_path, env_file_path) -> None:
+@click.option("--reset-db", help="Reset OpenMetadata Data", is_flag=True)
+def docker(
+    start, stop, pause, resume, clean, file_path, env_file_path, reset_db
+) -> None:
     """
     Checks Docker Memory Allocation
     Run Latest Release Docker - metadata docker --start
     Run Local Docker - metadata docker --start -f path/to/docker-compose.yml
     """
-    run_docker(start, stop, pause, resume, clean, file_path, env_file_path)
+    run_docker(start, stop, pause, resume, clean, file_path, env_file_path, reset_db)
 
 
 @metadata.command()