Merge pull request #1 from RelationalAI/andrew/clone_functionality

Clone DB from src DB
RelationalAI · Aug 15, 2023 · 865114f · 865114f
2 parents 091f402 + dcccf06
commit 865114f
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 7 deletions.
diff --git a/cli/args.py b/cli/args.py
@@ -24,6 +24,12 @@ def parse() -> Namespace:
         required=True,
         type=str
     )
+    parser.add_argument(
+        "--source-database",
+        help="RAI database for clone",
+        required=False,
+        type=str
+    )
     parser.add_argument(
         "--engine",
         help="RAI engine",

diff --git a/cli/runner.py b/cli/runner.py
@@ -39,7 +39,7 @@ def start():
         # Skip infrastructure setup during recovery
         if not args.recover and not args.recover_step:
             # Create db and disable IVM in case of enabled flag
-            resource_manager.create_database(args.drop_db, args.disable_ivm)
+            resource_manager.create_database(args.drop_db, args.disable_ivm, args.source_database)
         # Init workflow executor
         parameters = {
             workflow.constants.REL_CONFIG_DIR: args.rel_config_dir,

diff --git a/workflow/manager.py b/workflow/manager.py
@@ -25,17 +25,18 @@ def cleanup_resources(self) -> None:
         rai.delete_database(self.logger, self.rai_config)
         rai.delete_engine(self.logger, self.rai_config)
 
-    def create_database(self, deleted_db: bool = False, disable_ivm: bool = False) -> None:
+    def create_database(self, delete_db: bool = False, disable_ivm: bool = False, source_db=None) -> None:
         """
         Create RAI database. If `delete_db` is True then manager should delete previous db with the same name if it
          does exist before creation. `disable_ivm` indicate if manager need to disable IVM for a new db.
-        :param deleted_db:      delete database flag
+        :param delete_db:       delete database flag
         :param disable_ivm:     disable ivm flag
+        :param source_db:       source database name
         :return:
         """
-        if deleted_db:
+        if delete_db:
             self.deleted_database()
-        rai.create_database(self.logger, self.rai_config)
+        rai.create_database(self.logger, self.rai_config, source_db)
         if disable_ivm:
             self.logger.info(f"Disabling IVM for `{self.rai_config.database}`")
             rai.execute_query(self.logger, self.rai_config, q.DISABLE_IVM, readonly=False)

diff --git a/workflow/rai.py b/workflow/rai.py
@@ -81,16 +81,19 @@ def engine_exist(logger: logging.Logger, rai_config: RaiConfig) -> bool:
     return bool(api.get_engine(rai_config.ctx, rai_config.engine))
 
 
-def create_database(logger: logging.Logger, rai_config: RaiConfig) -> None:
+def create_database(logger: logging.Logger, rai_config: RaiConfig, source_db=None) -> None:
     """
     Create RAI DB specified in RAI config. The HTTP error 409(Conflict) swallowed since DB already exists.
     :param logger:      logger
     :param rai_config:  RAI config
+    :param source_db:   source DB to clone from
     :return:
     """
     logger.info(f"Creating database `{rai_config.database}`")
+    if source_db:
+        logger.info(f"Use `{source_db}` database for clone")
     try:
-        api.create_database(rai_config.ctx, rai_config.database)
+        api.create_database(rai_config.ctx, rai_config.database, source_db)
     except HTTPError as e:
         if e.status == 409:
             logger.info(f"Database '{rai_config.database}' already exists")