Skip to content

Commit

Permalink
Add secondary_profiles to profile.py
Browse files Browse the repository at this point in the history
  • Loading branch information
aranke committed Feb 14, 2025
1 parent 7041e58 commit 0691a5f
Show file tree
Hide file tree
Showing 3 changed files with 207 additions and 1 deletion.
40 changes: 39 additions & 1 deletion core/dbt/config/profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ class Profile(HasCredentials):
credentials: Credentials
profile_env_vars: Dict[str, Any]
log_cache_events: bool
secondary_profiles: Dict[str, "Profile"]

def __init__(
self,
Expand All @@ -79,6 +80,7 @@ def __init__(
self.log_cache_events = (
get_flags().LOG_CACHE_EVENTS
) # never available on init, set for adapter instantiation via AdapterRequiredConfig
self.secondary_profiles = {}

def to_profile_info(self, serialize_credentials: bool = False) -> Dict[str, Any]:
"""Unlike to_project_config, this dict is not a mirror of any existing
Expand Down Expand Up @@ -325,13 +327,49 @@ def from_raw_profile_info(
profile_data, profile_name, target_name
)

return cls.from_credentials(
profile = cls.from_credentials(
credentials=credentials,
profile_name=profile_name,
target_name=target_name,
threads=threads,
)

raw_secondary_profiles = profile_data.pop("secondary_profiles", [])

if raw_secondary_profiles:
for p in raw_secondary_profiles:
for secondary_profile_name, secondary_raw_profile in p.items():
if secondary_profile_name in profile.secondary_profiles:
raise DbtProfileError(
f"Secondary profile '{secondary_profile_name}' is already defined"
)

secondary_target_name, secondary_profile_data = cls.render_profile(
secondary_raw_profile, secondary_profile_name, target_override, renderer
)

if secondary_profile_data.get("secondary_profiles"):
raise DbtProfileError(
f"Secondary profile '{secondary_profile_name}' cannot have nested secondary profiles"
)

secondary_threads = secondary_profile_data.pop("threads", DEFAULT_THREADS)
if threads_override is not None:
secondary_threads = threads_override

secondary_credentials: Credentials = cls._credentials_from_profile(
secondary_profile_data, secondary_profile_name, secondary_target_name
)

profile.secondary_profiles[secondary_profile_name] = cls.from_credentials(
credentials=secondary_credentials,
profile_name=secondary_profile_name,
target_name=secondary_target_name,
threads=secondary_threads,
)

return profile

@classmethod
def from_raw_profiles(
cls,
Expand Down
1 change: 1 addition & 0 deletions core/dbt/config/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,7 @@ def from_parts(
project_env_vars=project.project_env_vars,
restrict_access=project.restrict_access,
profile_env_vars=profile.profile_env_vars,
secondary_profiles=profile.secondary_profiles,
profile_name=profile.profile_name,
target_name=profile.target_name,
threads=profile.threads,
Expand Down
167 changes: 167 additions & 0 deletions tests/unit/config/test_profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -292,3 +292,170 @@ def test_profile_with_empty_profile_data(self):
self.default_profile_data, "empty_profile_data", renderer
)
self.assertIn("Profile empty_profile_data in profiles.yml is empty", str(exc.exception))


class TestSecondaryProfiles(BaseConfigTest):
def test_secondary_profiles_basic(self):
profile_data_with_secondary = deepcopy(self.default_profile_data)
profile_data_with_secondary["default"]["outputs"]["postgres"]["secondary_profiles"] = [
{
"secondary_profile_1": {
"target": "secondary_target_1",
"outputs": {
"secondary_target_1": {
"type": "postgres",
"host": "secondary-host",
"port": 1234,
"user": "secondary_user",
"password": "secondary_password",
"schema": "secondary_schema",
"database": "secondary_db",
}
},
}
}
]
renderer = empty_profile_renderer()
profile = dbt.config.Profile.from_raw_profiles(
profile_data_with_secondary, "default", renderer
)

self.assertIn("secondary_profile_1", profile.secondary_profiles)
secondary_profile = profile.secondary_profiles["secondary_profile_1"]
self.assertEqual(secondary_profile.profile_name, "secondary_profile_1")
self.assertEqual(secondary_profile.target_name, "secondary_target_1")
self.assertTrue(isinstance(secondary_profile.credentials, PostgresCredentials))
self.assertEqual(secondary_profile.credentials.host, "secondary-host")
self.assertEqual(secondary_profile.credentials.port, 1234)
self.assertEqual(secondary_profile.credentials.database, "secondary_db")

def test_secondary_profiles_override_threads(self):
profile_data_with_secondary = deepcopy(self.default_profile_data)
profile_data_with_secondary["default"]["outputs"]["postgres"]["secondary_profiles"] = [
{
"secondary_profile_1": {
"target": "secondary_target_1",
"outputs": {
"secondary_target_1": {
"type": "postgres",
"host": "secondary-host",
"port": 1234,
"user": "secondary_user",
"password": "secondary_password",
"schema": "secondary_schema",
"database": "secondary_db",
"threads": 5, # threads defined in secondary profile
}
},
}
}
]
renderer = empty_profile_renderer()
profile = dbt.config.Profile.from_raw_profiles(
profile_data_with_secondary, "default", renderer, threads_override=10
) # override threads to 10

self.assertIn("secondary_profile_1", profile.secondary_profiles)
secondary_profile = profile.secondary_profiles["secondary_profile_1"]
# threads_override should take precedence over threads defined in secondary profile
self.assertEqual(secondary_profile.threads, 10)
self.assertEqual(
profile.threads, 10
) # primary profile should also have overridden threads

def test_secondary_profiles_duplicate_names(self):
profile_data_with_duplicate_secondary = deepcopy(self.default_profile_data)
profile_data_with_duplicate_secondary["default"]["outputs"]["postgres"][
"secondary_profiles"
] = [
{
"secondary_profile_1": {
"target": "secondary_target_1",
"outputs": {
"secondary_target_1": {
"type": "postgres",
"host": "secondary-host",
"port": 1234,
"user": "secondary_user",
"password": "secondary_password",
"schema": "secondary_schema",
"database": "secondary_db",
}
},
}
},
{
"secondary_profile_1": { # Duplicate name
"target": "secondary_target_2",
"outputs": {
"secondary_target_2": {
"type": "postgres",
"host": "another-secondary-host",
"port": 5678,
"user": "another_secondary_user",
"password": "another_secondary_password",
"schema": "another_secondary_schema",
"database": "another_secondary_db",
}
},
}
},
]
renderer = empty_profile_renderer()
with self.assertRaises(dbt.exceptions.DbtProfileError) as exc:
dbt.config.Profile.from_raw_profiles(
profile_data_with_duplicate_secondary, "default", renderer
)
self.assertIn(
"Secondary profile 'secondary_profile_1' is already defined", str(exc.exception)
)

def test_secondary_profiles_nested_secondary(self):
profile_data_with_nested_secondary = deepcopy(self.default_profile_data)
profile_data_with_nested_secondary["default"]["outputs"]["postgres"][
"secondary_profiles"
] = [
{
"secondary_profile_1": {
"target": "secondary_target_1",
"outputs": {
"secondary_target_1": {
"type": "postgres",
"host": "secondary-host",
"port": 1234,
"user": "secondary_user",
"password": "secondary_password",
"schema": "secondary_schema",
"database": "secondary_db",
"secondary_profiles": [ # Nested secondary profiles - should be disallowed
{
"nested_secondary_profile": {
"target": "nested_target",
"outputs": {
"nested_target": {
"type": "postgres",
"host": "nested-host",
"port": 9012,
"user": "nested_user",
"password": "nested_password",
"schema": "nested_schema",
"database": "nested_db",
}
},
}
}
],
}
},
}
}
]
renderer = empty_profile_renderer()
with self.assertRaises(dbt.exceptions.DbtProfileError) as exc:
dbt.config.Profile.from_raw_profiles(
profile_data_with_nested_secondary, "default", renderer
)
self.assertIn(
"Secondary profile 'secondary_profile_1' cannot have nested secondary profiles",
str(exc.exception),
)

0 comments on commit 0691a5f

Please sign in to comment.