feat(helm): allow the postgresql scheme to be configurable (#8126)

dagster-io · Jun 1, 2022 · b4b4e33 · b4b4e33
1 parent bc038fe
commit b4b4e33
Show file tree

Hide file tree

Showing 8 changed files with 53 additions and 30 deletions.
diff --git a/helm/dagster/schema/schema/charts/dagster/subschema/postgresql.py b/helm/dagster/schema/schema/charts/dagster/subschema/postgresql.py
@@ -15,4 +15,5 @@ class PostgreSQL(BaseModel):
     postgresqlPassword: str
     postgresqlDatabase: str
     postgresqlParams: dict
+    postgresqlScheme: str
     service: Service
diff --git a/helm/dagster/schema/schema_tests/test_instance.py b/helm/dagster/schema/schema_tests/test_instance.py
@@ -58,8 +58,9 @@ def helm_template() -> HelmTemplate:
     )
 
 
+@pytest.mark.parametrize("postgresql_scheme", ["postgresql", "postgresql+psycopg2"])
 @pytest.mark.parametrize("storage", ["schedule_storage", "run_storage", "event_log_storage"])
-def test_storage_postgres_db_config(template: HelmTemplate, storage: str):
+def test_storage_postgres_db_config(template: HelmTemplate, postgresql_scheme: str, storage: str):
     postgresql_username = "username"
     postgresql_host = "1.1.1.1"
     postgresql_database = "database"
@@ -75,6 +76,7 @@ def test_storage_postgres_db_config(template: HelmTemplate, storage: str):
             postgresqlHost=postgresql_host,
             postgresqlDatabase=postgresql_database,
             postgresqlParams=postgresql_params,
+            postgresqlScheme=postgresql_scheme,
             service=Service(port=postgresql_port),
         )
     )
@@ -95,6 +97,7 @@ def test_storage_postgres_db_config(template: HelmTemplate, storage: str):
     assert postgres_db["db_name"] == postgresql_database
     assert postgres_db["port"] == postgresql_port
     assert postgres_db["params"] == postgresql_params
+    assert postgres_db["scheme"] == postgresql_scheme
 
 
 def test_k8s_run_launcher_config(template: HelmTemplate):

diff --git a/helm/dagster/templates/configmap-instance.yaml b/helm/dagster/templates/configmap-instance.yaml
@@ -22,14 +22,7 @@ data:
       module: dagster_postgres.schedule_storage
       class: PostgresScheduleStorage
       config:
-        postgres_db:
-          username: {{ .Values.postgresql.postgresqlUsername }}
-          password:
-            env: DAGSTER_PG_PASSWORD
-          hostname: {{ include "dagster.postgresql.host" . }}
-          db_name:  {{ .Values.postgresql.postgresqlDatabase	}}
-          port: {{ .Values.postgresql.service.port }}
-          params: {{ .Values.postgresql.postgresqlParams | toYaml | nindent 12 }}
+        {{- include "dagsterYaml.postgresql.config" . | indent 8 }}
 
     run_launcher:
       {{- $runLauncherType := .Values.runLauncher.type }}
@@ -46,27 +39,13 @@ data:
       module: dagster_postgres.run_storage
       class: PostgresRunStorage
       config:
-        postgres_db:
-          username: {{ .Values.postgresql.postgresqlUsername }}
-          password:
-            env: DAGSTER_PG_PASSWORD
-          hostname: {{ include "dagster.postgresql.host" . }}
-          db_name:  {{ .Values.postgresql.postgresqlDatabase	}}
-          port: {{ .Values.postgresql.service.port }}
-          params: {{ .Values.postgresql.postgresqlParams | toYaml | nindent 12 }}
+        {{- include "dagsterYaml.postgresql.config" . | indent 8 }}
 
     event_log_storage:
       module: dagster_postgres.event_log
       class: PostgresEventLogStorage
       config:
-        postgres_db:
-          username: {{ .Values.postgresql.postgresqlUsername }}
-          password:
-            env: DAGSTER_PG_PASSWORD
-          hostname: {{ include "dagster.postgresql.host" . }}
-          db_name:  {{ .Values.postgresql.postgresqlDatabase	}}
-          port: {{ .Values.postgresql.service.port }}
-          params: {{ .Values.postgresql.postgresqlParams | toYaml | nindent 12 }}
+        {{- include "dagsterYaml.postgresql.config" . | indent 8 }}
 
     {{- if and (.Values.dagsterDaemon.enabled) (.Values.dagsterDaemon.runCoordinator.enabled) }}
     {{- $runCoordinatorType := .Values.dagsterDaemon.runCoordinator.type }}

diff --git a/helm/dagster/templates/helpers/instance/_postgresql.tpl b/helm/dagster/templates/helpers/instance/_postgresql.tpl
@@ -0,0 +1,11 @@
+{{- define "dagsterYaml.postgresql.config" }}
+postgres_db:
+  username: {{ .Values.postgresql.postgresqlUsername }}
+  password:
+    env: DAGSTER_PG_PASSWORD
+  hostname: {{ include "dagster.postgresql.host" . }}
+  db_name: {{ .Values.postgresql.postgresqlDatabase	}}
+  port: {{ .Values.postgresql.service.port }}
+  params: {{- .Values.postgresql.postgresqlParams | toYaml | nindent 4 }}
+  scheme: {{ .Values.postgresql.postgresqlScheme }}
+{{- end }}
diff --git a/helm/dagster/values.schema.json b/helm/dagster/values.schema.json
diff --git a/helm/dagster/values.yaml b/helm/dagster/values.yaml
@@ -688,6 +688,11 @@ postgresql:
   # see: https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-PARAMKEYWORDS
   postgresqlParams: {}
 
+  # When set, overrides the default `postgresql` scheme for the connection string.
+  # (e.g., set to `postgresql+pg8000` to use the `pg8000` library).
+  # See: https://docs.sqlalchemy.org/en/13/dialects/postgresql.html#dialect-postgresql
+  postgresqlScheme: "postgresql"
+
   service:
     port: 5432
 

diff --git a/python_modules/libraries/dagster-postgres/dagster_postgres/utils.py b/python_modules/libraries/dagster-postgres/dagster_postgres/utils.py
@@ -34,6 +34,7 @@ def pg_config():
                 "db_name": StringSource,
                 "port": Field(IntSource, is_required=False, default_value=5432),
                 "params": Field(Permissive(), is_required=False, default_value={}),
+                "scheme": Field(StringSource, is_required=False, default_value="postgresql"),
             },
             is_required=False,
         ),
@@ -58,8 +59,10 @@ def pg_url_from_config(config_value):
         return get_conn_string(**config_value["postgres_db"])
 
 
-def get_conn_string(username, password, hostname, db_name, port="5432", params=None):
-    uri = f"postgresql://{quote(username)}:{quote(password)}@{hostname}:{port}/{db_name}"
+def get_conn_string(
+    username, password, hostname, db_name, port="5432", params=None, scheme="postgresql"
+):
+    uri = f"{scheme}://{quote(username)}:{quote(password)}@{hostname}:{port}/{db_name}"
 
     if params:
         query_string = f"{urlencode(params, quote_via=quote)}"

diff --git a/python_modules/libraries/dagster-postgres/dagster_postgres_tests/test_instance.py b/python_modules/libraries/dagster-postgres/dagster_postgres_tests/test_instance.py
@@ -231,17 +231,33 @@ def test_conn_str():
     db_name = "dagster"
     hostname = "database-city.com"
 
+    url_wo_scheme = r"has%40init:%3Afull%3A%20of%20junk%21%40%3F@database-city.com:5432/dagster"
+
     conn_str = get_conn_string(
         username=username,
         password=password,
         db_name=db_name,
         hostname=hostname,
     )
-    assert (
-        conn_str
-        == r"postgresql://has%40init:%3Afull%3A%20of%20junk%21%40%3F@database-city.com:5432/dagster"
+    assert conn_str == f"postgresql://{url_wo_scheme}"
+    parsed = urlparse(conn_str)
+    assert unquote(parsed.username) == username
+    assert unquote(parsed.password) == password
+    assert parsed.hostname == hostname
+    assert parsed.scheme == "postgresql"
+
+    custom_scheme = "postgresql+dialect"
+    conn_str = get_conn_string(
+        username=username,
+        password=password,
+        db_name=db_name,
+        hostname=hostname,
+        scheme=custom_scheme,
     )
+
+    assert conn_str == f"postgresql+dialect://{url_wo_scheme}"
     parsed = urlparse(conn_str)
     assert unquote(parsed.username) == username
     assert unquote(parsed.password) == password
     assert parsed.hostname == hostname
+    assert parsed.scheme == custom_scheme