clean up create_pg_connection (#8165)

dagster-io · Jun 2, 2022 · 974af9f · 974af9f
1 parent 5c35d39
commit 974af9f
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 15 deletions.
diff --git a/python_modules/libraries/dagster-postgres/dagster_postgres/event_log/event_log.py b/python_modules/libraries/dagster-postgres/dagster_postgres/event_log/event_log.py
@@ -214,7 +214,7 @@ def store_asset_event(self, event):
             conn.execute(query)
 
     def _connect(self):
-        return create_pg_connection(self._engine, pg_alembic_config(__file__), "event log")
+        return create_pg_connection(self._engine)
 
     def run_connection(self, run_id=None):
         return self._connect()

diff --git a/python_modules/libraries/dagster-postgres/dagster_postgres/run_storage/run_storage.py b/python_modules/libraries/dagster-postgres/dagster_postgres/run_storage/run_storage.py
@@ -119,11 +119,7 @@ def create_clean_storage(postgres_url, should_autocreate_tables=True):
         return PostgresRunStorage(postgres_url, should_autocreate_tables)
 
     def connect(self):
-        return create_pg_connection(
-            self._engine,
-            pg_alembic_config(__file__),
-            "run",
-        )
+        return create_pg_connection(self._engine)
 
     def upgrade(self):
         with self.connect() as conn:

diff --git a/..._modules/libraries/dagster-postgres/dagster_postgres/schedule_storage/schedule_storage.py b/..._modules/libraries/dagster-postgres/dagster_postgres/schedule_storage/schedule_storage.py
@@ -111,7 +111,7 @@ def create_clean_storage(postgres_url, should_autocreate_tables=True):
         return PostgresScheduleStorage(postgres_url, should_autocreate_tables)
 
     def connect(self, run_id=None):  # pylint: disable=arguments-differ, unused-argument
-        return create_pg_connection(self._engine, pg_alembic_config(__file__), "schedule")
+        return create_pg_connection(self._engine)
 
     def upgrade(self):
         alembic_config = pg_alembic_config(__file__)

diff --git a/python_modules/libraries/dagster-postgres/dagster_postgres/utils.py b/python_modules/libraries/dagster-postgres/dagster_postgres/utils.py
@@ -156,15 +156,8 @@ def pg_alembic_config(dunder_file, script_location=None):
 
 
 @contextmanager
-def create_pg_connection(engine, _alembic_config, storage_type_desc=None):
+def create_pg_connection(engine):
     check.inst_param(engine, "engine", sqlalchemy.engine.Engine)
-    check.opt_str_param(storage_type_desc, "storage_type_desc", "")
-
-    if storage_type_desc:
-        storage_type_desc += " "
-    else:
-        storage_type_desc = ""
-
     conn = None
     try:
         # Retry connection to gracefully handle transient connection issues