Changes to testing and flags/config

dbt_coves/config/config.py

@@ -158,12 +158,13 @@ class BlueGreenModel(BaseModel):
     production_database: Optional[str] = ""
     staging_database: Optional[str] = ""
     staging_suffix: Optional[str] = ""
-    drop_staging_db: Optional[bool] = False
+    drop_staging_db_if_exists: Optional[bool] = False
     drop_staging_db_after: Optional[int] = None
     drop_staging_db_on_failure: Optional[bool] = False
     dbt_selector: Optional[str] = ""
     defer: Optional[bool] = False
     full_refresh: Optional[bool] = False
+    is_test: Optional[bool] = False
 
 
 class ConfigModel(BaseModel):
@@ -265,12 +266,13 @@ class DbtCovesConfig:
         "blue_green.production_database",
         "blue_green.staging_database",
         "blue_green.staging_suffix",
-        "blue_green.drop_staging_db",
+        "blue_green.drop_staging_db_if_exists",
         "blue_green.drop_staging_db_after",
         "blue_green.drop_staging_db_on_failure",
         "blue_green.dbt_selector",
         "blue_green.defer",
         "blue_green.full_refresh",
+        "blue_green.is_test",
     ]
 
     def __init__(self, flags: DbtCovesFlags) -> None:

dbt_coves/tasks/blue_green/main.py

@@ -47,7 +47,7 @@ def register_parser(cls, sub_parsers, base_subparser):
             help="Green database suffix",
         )
         ext_subparser.add_argument(
-            "--drop-staging-db",
+            "--drop-staging-db-if-exists",
             action="store_true",
             help="Drop staging db after swap",
         )
@@ -72,6 +72,11 @@ def register_parser(cls, sub_parsers, base_subparser):
             help="Perform a full dbt build",
         )
         ext_subparser.add_argument("--defer", action="store_true", help="Run in deferral")
+        ext_subparser.add_argument(
+            "--is-test",
+            action="store_true",
+            help="",
+        )
         return ext_subparser
 
     def get_config_value(self, key):
@@ -102,7 +107,7 @@ def run(self) -> int:
                 f"Production database {self.production_database} cannot be the same as staging database "
                 f"{self.staging_database}"
             )
-        self.drop_staging_db = self.get_config_value("drop_staging_db")
+        self.drop_staging_db_if_exists = self.get_config_value("drop_staging_db_if_exists")
         self.drop_staging_db_after = self.get_config_value("drop_staging_db_after")
         self.con = self.snowflake_connection()
 
@@ -111,6 +116,7 @@ def run(self) -> int:
             self.staging_database,
             self.con,
         )
+
         self._check_and_drop_staging_db()
 
         try:
@@ -129,7 +135,8 @@ def run(self) -> int:
             # Swaps databases: Snowflake sql `alter database {blue} swap with {green}`
             self._swap_databases()
             # drops pre_production (ex production)
-            self.cdb.drop_database()
+            if not self.get_config_value("is_test"):
+                self.cdb.drop_database()
         except Exception as e:
             if self.get_config_value("drop_staging_db_on_failure"):
                 self.cdb.drop_database()
@@ -174,7 +181,7 @@ def _check_and_drop_staging_db(self):
             None
         """
         green_exists = self._check_if_database_exists()
-        if green_exists and self.drop_staging_db:
+        if green_exists and self.drop_staging_db_if_exists:
             if self.drop_staging_db_after:
                 console.print(
                     f"Green database {self.staging_database} exists."

dbt_coves/utils/flags.py

@@ -142,12 +142,13 @@ def __init__(self, cli_parser: ArgumentParser) -> None:
             "production_database": None,
             "staging_database": None,
             "staging_suffix": None,
-            "drop_staging_db": False,
+            "drop_staging_db_if_exists": False,
             "drop_staging_db_after": None,
             "drop_staging_db_on_failure": False,
             "dbt_selector": None,
             "defer": False,
             "full_refresh": False,
+            "is_test": False,
         }
 
     def parse_args(self, cli_args: List[str] = list()) -> None:
@@ -427,8 +428,10 @@ def parse_args(self, cli_args: List[str] = list()) -> None:
                     self.blue_green["staging_database"] = self.args.staging_database
                 if self.args.staging_suffix:
                     self.blue_green["staging_suffix"] = self.args.staging_suffix
-                if self.args.drop_staging_db:
-                    self.blue_green["drop_staging_db"] = self.args.drop_staging_db
+                if self.args.drop_staging_db_if_exists:
+                    self.blue_green[
+                        "drop_staging_db_if_exists"
+                    ] = self.args.drop_staging_db_if_exists
                 if self.args.drop_staging_db_after:
                     self.blue_green["drop_staging_db_after"] = self.args.drop_staging_db_after
                 if self.args.drop_staging_db_on_failure:
@@ -441,3 +444,5 @@ def parse_args(self, cli_args: List[str] = list()) -> None:
                     self.blue_green["defer"] = self.args.defer
                 if self.args.full_refresh:
                     self.blue_green["full_refresh"] = self.args.full_refresh
+                if self.args.is_test:
+                    self.blue_green["is_test"] = self.args.is_test

tests/blue_green/.user.yml

@@ -0,0 +1 @@
+id: f29ef39e-061e-49b7-ab06-93112d5137f0

tests/blue_green/input/dummy_data.sql

@@ -1,7 +1,7 @@
 CREATE DATABASE IF NOT EXISTS DBT_COVES_TEST;
 USE DATABASE DBT_COVES_TEST;
 
-GRANT USAGE ON DATABASE DBT_COVES_TEST TO ROLE TRANSFORMER_DBT;
+GRANT USAGE ON DATABASE DBT_COVES_TEST TO ROLE bot_integration;
 
 CREATE SCHEMA IF NOT EXISTS DBT_COVES_TEST.TESTS_BLUE_GREEN;
 

tests/blue_green/models/blue_green/test_model_2.sql

@@ -1,4 +1,4 @@
-{{ config(materialized='table', schema="BLUE_GREEN") }}
+{{ config(materialized='table') }}
 
 with source_data as (
 

tests/blue_green/settings.yml

@@ -6,3 +6,4 @@ database: "DBT_COVES_TEST"
 blue_green:
   production_database: DBT_COVES_TEST
   staging_suffix: "STG"
+  drop_staging_db_if_exists: True

tests/blue_green_test.py

@@ -39,17 +39,19 @@
 @pytest.fixture(scope="class")
 def snowflake_connection(request):
     # Check env vars
-    assert "USER_SNOWFLAKE" in os.environ
-    assert "PASSWORD_SNOWFLAKE" in os.environ
-    assert "ACCOUNT_SNOWFLAKE" in os.environ
-    assert "WAREHOUSE_SNOWFLAKE" in os.environ
-    assert "ROLE_SNOWFLAKE" in os.environ
-
-    user = os.environ["USER_SNOWFLAKE"]
-    password = os.environ["PASSWORD_SNOWFLAKE"]
-    account = os.environ["ACCOUNT_SNOWFLAKE"]
-    role = os.environ["ROLE_SNOWFLAKE"]
-    warehouse = os.environ["WAREHOUSE_SNOWFLAKE"]
+    assert "DATACOVES__DBT_COVES_TEST__USER" in os.environ
+    assert "DATACOVES__DBT_COVES_TEST__PASSWORD" in os.environ
+    assert "DATACOVES__DBT_COVES_TEST__ACCOUNT" in os.environ
+    assert "DATACOVES__DBT_COVES_TEST__WAREHOUSE" in os.environ
+    assert "DATACOVES__DBT_COVES_TEST__ROLE" in os.environ
+
+    user = os.environ["DATACOVES__DBT_COVES_TEST__USER"]
+    password = os.environ["DATACOVES__DBT_COVES_TEST__PASSWORD"]
+    account = os.environ["DATACOVES__DBT_COVES_TEST__ACCOUNT"]
+    role = os.environ["DATACOVES__DBT_COVES_TEST__ROLE"]
+    warehouse = os.environ["DATACOVES__DBT_COVES_TEST__WAREHOUSE"]
+    database = os.environ["DATACOVES__DBT_COVES_TEST__DATABASE"]
+    schema = "TESTS_BLUE_GREEN"
 
     conn = snowflake.connector.connect(
         user=user,
@@ -61,8 +63,8 @@ def snowflake_connection(request):
 
     request.cls.conn = conn
     request.cls.warehouse = warehouse
-    request.cls.schema = SETTINGS["schema"]
-    request.cls.production_database = SETTINGS["database"]
+    request.cls.schema = schema
+    request.cls.production_database = database
     request.cls.staging_database = (
         f"{request.cls.production_database}_{DBT_COVES_SETTINGS.get('staging_suffix', 'staging')}"
     )
@@ -98,12 +100,15 @@ def test_dbt_coves_bluegreen(self):
             "../../dbt_coves/core/main.py",
             "blue-green",
             "--project-dir",
-            FIXTURE_DIR,
+            str(FIXTURE_DIR),
+            "--profiles-dir",
+            str(FIXTURE_DIR),
             "--production-database",
             self.production_database,
+            "--is-test",
         ]
-        if DBT_COVES_SETTINGS.get("drop_staging_db"):
-            command.append("--drop-staging-db")
+        if DBT_COVES_SETTINGS.get("drop_staging_db_if_exists"):
+            command.append("--drop-staging-db-if-exists")
         if DBT_COVES_SETTINGS.get("dbt_selector"):
             command.extend(["--dbt-selector", DBT_COVES_SETTINGS["dbt_selector"]])
         if DBT_COVES_SETTINGS.get("staging_suffix"):