diff --git a/dbt_coves/config/config.py b/dbt_coves/config/config.py index ef744910..a9e350d8 100644 --- a/dbt_coves/config/config.py +++ b/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: diff --git a/dbt_coves/tasks/blue_green/main.py b/dbt_coves/tasks/blue_green/main.py index c3fc2de5..2c3bcee5 100644 --- a/dbt_coves/tasks/blue_green/main.py +++ b/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." diff --git a/dbt_coves/utils/flags.py b/dbt_coves/utils/flags.py index cd31f7ba..fe258c23 100644 --- a/dbt_coves/utils/flags.py +++ b/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 diff --git a/tests/blue_green/.user.yml b/tests/blue_green/.user.yml new file mode 100644 index 00000000..0da24ae8 --- /dev/null +++ b/tests/blue_green/.user.yml @@ -0,0 +1 @@ +id: f29ef39e-061e-49b7-ab06-93112d5137f0 diff --git a/tests/blue_green/input/dummy_data.sql b/tests/blue_green/input/dummy_data.sql index 7ef00d28..774aa97f 100644 --- a/tests/blue_green/input/dummy_data.sql +++ b/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; diff --git a/tests/blue_green/models/blue_green/test_model_2.sql b/tests/blue_green/models/blue_green/test_model_2.sql index 18b60383..381d90f2 100644 --- a/tests/blue_green/models/blue_green/test_model_2.sql +++ b/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 ( diff --git a/tests/blue_green/settings.yml b/tests/blue_green/settings.yml index 4e7f00b6..bf49f05c 100644 --- a/tests/blue_green/settings.yml +++ b/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 diff --git a/tests/blue_green_test.py b/tests/blue_green_test.py index 458ff41f..435cd1a5 100644 --- a/tests/blue_green_test.py +++ b/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"):