Skip to content

Commit

Permalink
Make PostgresUserPasswordProfileMapping schema argument optional (#683)
Browse files Browse the repository at this point in the history 
Made PostgresUserPasswordProfileMapping schema profile argument optional

Closes: #675

Co-authored-by: Tatiana Al-Chueyr <[email protected]>
  • Loading branch information
@FouziaTariq @tatiana
FouziaTariq and tatiana authored Feb 29, 2024
1 parent ffa724a commit 1bf78b4
Show file tree
Hide file tree
Showing 4 changed files with 116 additions and 20 deletions.
14 changes: 9 additions & 5 deletions cosmos/profiles/postgres/user_pass.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ class PostgresUserPasswordProfileMapping(BaseProfileMapping):
"user",
"password",
"dbname",
"schema",
]
secret_fields = [
"password",
Expand All @@ -47,14 +46,19 @@ def profile(self) -> dict[str, Any | None]:
"password": self.get_env_var_format("password"),
}

if "schema" in self.profile_args:
profile["schema"] = self.profile_args["schema"]

return self.filter_null(profile)

@property
def mock_profile(self) -> dict[str, Any | None]:
"Gets mock profile. Defaults port to 5432."
parent_mock = super().mock_profile

return {
profile_dict = {
"port": 5432,
**parent_mock,
**super().mock_profile,
}
user_defined_schema = self.profile_args.get("schema")
if user_defined_schema:
profile_dict["schema"] = user_defined_schema
return profile_dict
2 changes: 1 addition & 1 deletion tests/dbt/test_graph.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -457,7 +457,7 @@ def test_load_via_dbt_ls_with_sources(load_method):
),
)
getattr(dbt_graph, load_method)()
assert len(dbt_graph.nodes) == 4
assert len(dbt_graph.nodes) >= 4
assert "source.simple.main.movies_ratings" in dbt_graph.nodes
assert "exposure.simple.weekly_metrics" in dbt_graph.nodes

Expand Down
4 changes: 2 additions & 2 deletions tests/profiles/postgres/test_pg_user_pass.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,9 +83,9 @@ def test_connection_claiming() -> None:
assert not profile_mapping.can_claim_connection()

# also test when there's no schema
conn = Connection(**potential_values) # type: ignore
conn = Connection(**{k: v for k, v in potential_values.items() if k != "schema"})
with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn):
profile_mapping = PostgresUserPasswordProfileMapping(conn, {})
profile_mapping = PostgresUserPasswordProfileMapping(conn, {"schema": None})
assert not profile_mapping.can_claim_connection()

# if we have them all, it should claim
Expand Down
116 changes: 104 additions & 12 deletions tests/sample/manifest_source.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,10 @@
"exposure.simple.weekly_metrics"
],
"source.simple.main.movies_ratings": [
"model.simple.movies_ratings_simplified"
]
"model.simple.movies_ratings_simplified",
"test.simple.source_not_null_imdb_movies_ratings_X.e684bf90f4"
],
"test.simple.source_not_null_imdb_movies_ratings_X.e684bf90f4": []
},
"disabled": {},
"docs": {
Expand All @@ -28,7 +30,7 @@
"config": {
"enabled": true
},
"created_at": 1696859933.549042,
"created_at": 1697205180.995924,
"depends_on": {
"macros": [],
"nodes": [
Expand Down Expand Up @@ -1016,7 +1018,7 @@
"node_color": null,
"show": true
},
"macro_sql": "{% macro dates_in_range(start_date_str, end_date_str=none, in_fmt=\"%Y%m%d\", out_fmt=\"%Y%m%d\") %}\n {% set end_date_str = start_date_str if end_date_str is none else end_date_str %}\n\n {% set start_date = convert_datetime(start_date_str, in_fmt) %}\n {% set end_date = convert_datetime(end_date_str, in_fmt) %}\n\n {% set day_count = (end_date - start_date).days %}\n {% if day_count < 0 %}\n {% set msg -%}\n Partition start date is after the end date ({{ start_date }}, {{ end_date }})\n {%- endset %}\n\n {{ exceptions.raise_compiler_error(msg, model) }}\n {% endif %}\n\n {% set date_list = [] %}\n {% for i in range(0, day_count + 1) %}\n {% set the_date = (modules.datetime.timedelta(days=i) + start_date) %}\n {% if not out_fmt %}\n {% set _ = date_list.append(the_date) %}\n {% else %}\n {% set _ = date_list.append(the_date.strftime(out_fmt)) %}\n {% endif %}\n {% endfor %}\n\n {{ return(date_list) }}\n{% endmacro %}",
"macro_sql": "{% macro dates_in_range(start_date_str, end_date_str=none, in_fmt=\"%Y%m%d\", out_fmt=\"%Y%m%d\") %}\n {% set end_date_str = start_date_str if end_date_str is none else end_date_str %}\n\n {% set start_date = convert_datetime(start_date_str, in_fmt) %}\n {% set end_date = convert_datetime(end_date_str, in_fmt) %}\n\n {% set day_count = (end_date - start_date).days %}\n {% if day_count < 0 %}\n {% set msg -%}\n Partiton start date is after the end date ({{ start_date }}, {{ end_date }})\n {%- endset %}\n\n {{ exceptions.raise_compiler_error(msg, model) }}\n {% endif %}\n\n {% set date_list = [] %}\n {% for i in range(0, day_count + 1) %}\n {% set the_date = (modules.datetime.timedelta(days=i) + start_date) %}\n {% if not out_fmt %}\n {% set _ = date_list.append(the_date) %}\n {% else %}\n {% set _ = date_list.append(the_date.strftime(out_fmt)) %}\n {% endif %}\n {% endfor %}\n\n {{ return(date_list) }}\n{% endmacro %}",
"meta": {},
"name": "dates_in_range",
"original_file_path": "macros/etc/datetime.sql",
Expand Down Expand Up @@ -6945,8 +6947,8 @@
"dbt_schema_version": "https://schemas.getdbt.com/dbt/manifest/v8.json",
"dbt_version": "1.4.0",
"env": {},
"generated_at": "2023-10-09T13:58:53.320926Z",
"invocation_id": "1790dc18-1177-4ca0-b993-a8eeb59f0c4c",
"generated_at": "2023-10-13T13:58:46.591195Z",
"invocation_id": "4f4b6d15-b4ab-4683-bbe5-efd94824b1d9",
"project_id": "8dbdda48fb8748d6746f1965824e966a",
"send_anonymous_usage_stats": true,
"user_id": "4bdc3972-5c9f-4f16-90bd-3769a225fbe6"
Expand Down Expand Up @@ -6986,7 +6988,7 @@
"tags": [],
"unique_key": null
},
"created_at": 1696859933.5317938,
"created_at": 1697205180.9909241,
"database": "database",
"deferred": false,
"depends_on": {
Expand Down Expand Up @@ -7062,7 +7064,7 @@
"tags": [],
"unique_key": null
},
"created_at": 1696859933.537527,
"created_at": 1697205180.984051,
"database": "database",
"deferred": false,
"depends_on": {
Expand Down Expand Up @@ -7103,6 +7105,84 @@
"unrendered_config": {
"materialized": "table"
}
},
"test.simple.source_not_null_imdb_movies_ratings_X.e684bf90f4": {
"alias": "source_not_null_imdb_movies_ratings_X",
"build_path": null,
"checksum": {
"checksum": "",
"name": "none"
},
"column_name": "X",
"columns": {},
"compiled_path": null,
"config": {
"alias": null,
"database": null,
"enabled": true,
"error_if": "!= 0",
"fail_calc": "count(*)",
"limit": null,
"materialized": "test",
"meta": {},
"schema": "dbt_test__audit",
"severity": "ERROR",
"store_failures": null,
"tags": [],
"warn_if": "!= 0",
"where": null
},
"created_at": 1697205181.009168,
"database": "database",
"deferred": false,
"depends_on": {
"macros": [
"macro.dbt.test_not_null"
],
"nodes": [
"source.simple.main.movies_ratings"
]
},
"description": "",
"docs": {
"node_color": null,
"show": true
},
"file_key_name": "sources.imdb",
"fqn": [
"simple",
"source_not_null_imdb_movies_ratings_X"
],
"language": "sql",
"meta": {},
"metrics": [],
"name": "source_not_null_imdb_movies_ratings_X",
"original_file_path": "models/source.yml",
"package_name": "simple",
"patch_path": null,
"path": "source_not_null_imdb_movies_ratings_X.sql",
"raw_code": "{{ test_not_null(**_dbt_generic_test_kwargs) }}",
"refs": [],
"relation_name": null,
"resource_type": "test",
"schema": "main_dbt_test__audit",
"sources": [
[
"imdb",
"movies_ratings"
]
],
"tags": [],
"test_metadata": {
"kwargs": {
"column_name": "X",
"model": "{{ get_where_subquery(source('imdb', 'movies_ratings')) }}"
},
"name": "not_null",
"namespace": null
},
"unique_id": "test.simple.source_not_null_imdb_movies_ratings_X.e684bf90f4",
"unrendered_config": {}
}
},
"parent_map": {
Expand All @@ -7115,18 +7195,30 @@
"model.simple.top_animations": [
"model.simple.movies_ratings_simplified"
],
"source.simple.main.movies_ratings": []
"source.simple.main.movies_ratings": [],
"test.simple.source_not_null_imdb_movies_ratings_X.e684bf90f4": [
"source.simple.main.movies_ratings"
]
},
"selectors": {},
"sources": {
"source.simple.main.movies_ratings": {
"columns": {},
"columns": {
"X": {
"data_type": null,
"description": "",
"meta": {},
"name": "X",
"quote": null,
"tags": []
}
},
"config": {
"enabled": true
},
"created_at": 1696859933.549542,
"created_at": 1697205181.0098429,
"database": "database",
"description": "Ratings by movie",
"description": "Ratings by movie\n",
"external": null,
"fqn": [
"simple",
Expand Down

0 comments on commit 1bf78b4

Please sign in to comment.