feat: create function for get_sqla_engine with context #21790

hughhhh · 2022-10-12T20:28:51Z

SUMMARY

First step in allowing enabling ssh tunneling by allowing the get_sqla_engine function to be a context manager. This will allow us to add logic before and after yielding the engine for spinning up and tearing down the tunnel on each connection.

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TESTING INSTRUCTIONS

ADDITIONAL INFORMATION

Has associated issue:
Required feature flags:
Changes UI
Includes DB Migration (follow approval process in SIP-59)
- Migration is atomic, supports rollback & is backwards-compatible
- Confirm DB migration upgrade and downgrade tested
- Runtime estimates and downtime expectations provided
Introduces new feature or API
Removes existing feature or API

codecov · 2022-10-12T21:05:31Z

Codecov Report

Merging #21790 (63d4a9c) into master (8f4415b) will increase coverage by 0.87%.
The diff coverage is 71.42%.

@@            Coverage Diff             @@
##           master   #21790      +/-   ##
==========================================
+ Coverage   66.18%   67.06%   +0.87%     
==========================================
  Files        1805     1805              
  Lines       69066    69416     +350     
  Branches     7369     7369              
==========================================
+ Hits        45712    46554     +842     
+ Misses      21448    20955     -493     
- Partials     1906     1907       +1

Flag	Coverage Δ
hive	`53.24% <71.42%> (+0.32%)`	⬆️
javascript	`53.33% <ø> (-0.01%)`	⬇️
mysql	`78.50% <71.42%> (?)`
postgres	`78.57% <71.42%> (?)`
presto	`53.14% <71.42%> (+0.32%)`	⬆️
python	`81.59% <71.42%> (+1.67%)`	⬆️
sqlite	`77.04% <71.42%> (+0.14%)`	⬆️
unit	`51.36% <42.85%> (+0.29%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
superset/models/core.py	`89.79% <71.42%> (+1.08%)`	⬆️
...ols/MetricControl/AdhocMetricEditPopover/index.jsx	`74.46% <0.00%> (-1.07%)`	⬇️
superset/examples/tabbed_dashboard.py	`0.00% <0.00%> (ø)`
superset/initialization/__init__.py	`91.50% <0.00%> (+0.04%)`	⬆️
superset/db_engine_specs/base.py	`89.69% <0.00%> (+0.15%)`	⬆️
superset/datasets/schemas.py	`97.75% <0.00%> (+0.36%)`	⬆️
superset/views/core.py	`76.06% <0.00%> (+0.45%)`	⬆️
superset/common/query_object.py	`94.38% <0.00%> (+0.51%)`	⬆️
superset/connectors/sqla/models.py	`91.01% <0.00%> (+0.51%)`	⬆️
superset/commands/importers/v1/utils.py	`93.50% <0.00%> (+1.29%)`	⬆️
... and 32 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

…tunnel-refactor-get-sqla-engine

Antonio-RiveroMartnez · 2022-10-20T18:27:42Z

superset/models/core.py

@@ -362,6 +362,18 @@ def get_effective_user(self, object_url: URL) -> Optional[str]:
            else None
        )

+    @contextmanager
+    def get_sqla_engine_with_context(


What is the difference between having this new get_sqla_engine_with_context with the decorator VS using it in our existing get_sqla_engine ? I Mean, couldn't we just use the existing one and make use or not of the new functionality the decorator brings when needed? or would that mean changing a ton of places where the get_sqla_engine is being used right now?

For now there's no difference, but Hugh is planning to add support for SSH tunneling, which would require a setup phase before the engine is created, and a teardown after. In order for the SSH tunnel to work everywhere we will need to replace all existing calls with the new context manager.

Oh ok ok, makes sense then 😎 thanks!

Antonio-RiveroMartnez · 2022-10-20T18:32:47Z

tests/integration_tests/model_tests.py

-        db = make_url(model.get_sqla_engine().url).database
-        self.assertEqual("prod", db)
+        with model.get_sqla_engine_with_context() as engine:
+            db = make_url(engine.url).database


I see we have this make_url_safe here: superset/databases/utils.py, would that serve the same purpose? If so, can we change to that one or must we keep using this?

make_url_safe prevents password from showing up in the logs when make_url fails for some reason. It should always be used in production, but in the tests it's fine to use make_url instead.

Good to know 👍 Thanks!

betodealmeida

What's the plan for modifying the files in superset/?

betodealmeida · 2022-10-20T23:48:57Z

tests/integration_tests/access_tests.py

+            perm_data = ROLE_TABLES_PERM_DATA.copy()
+            perm_data["database"][0]["schema"][0]["name"] = schema

-        response = self.client.post(
-            "/superset/override_role_permissions/",
-            data=json.dumps(perm_data),
-            content_type="application/json",
-        )
-        self.assertEqual(201, response.status_code)
+            response = self.client.post(
+                "/superset/override_role_permissions/",
+                data=json.dumps(perm_data),
+                content_type="application/json",
+            )
+            self.assertEqual(201, response.status_code)

-        updated_override_me = security_manager.find_role("override_me")
-        self.assertEqual(1, len(updated_override_me.permissions))
-        birth_names = self.get_table(name="birth_names")
-        self.assertEqual(
-            birth_names.perm, updated_override_me.permissions[0].view_menu.name
-        )
-        self.assertEqual(
-            "datasource_access", updated_override_me.permissions[0].permission.name
-        )
+            updated_override_me = security_manager.find_role("override_me")
+            self.assertEqual(1, len(updated_override_me.permissions))
+            birth_names = self.get_table(name="birth_names")
+            self.assertEqual(
+                birth_names.perm, updated_override_me.permissions[0].view_menu.name
+            )
+            self.assertEqual(
+                "datasource_access", updated_override_me.permissions[0].permission.name
+            )


I think this doesn't have to be inside the context manager.

betodealmeida · 2022-10-20T23:50:44Z

tests/integration_tests/access_tests.py

+            override_me = security_manager.find_role("override_me")
+            override_me.permissions.append(
+                security_manager.find_permission_view_menu(
+                    view_menu_name=self.get_table(name="energy_usage").perm,
+                    permission_name="datasource_access",
+                )
            )
-        )
-        db.session.flush()
+            db.session.flush()

-        perm_data = ROLE_TABLES_PERM_DATA.copy()
-        perm_data["database"][0]["schema"][0]["name"] = schema
+            perm_data = ROLE_TABLES_PERM_DATA.copy()
+            perm_data["database"][0]["schema"][0]["name"] = schema

-        response = self.client.post(
-            "/superset/override_role_permissions/",
-            data=json.dumps(perm_data),
-            content_type="application/json",
-        )
-        self.assertEqual(201, response.status_code)
-        updated_override_me = security_manager.find_role("override_me")
-        self.assertEqual(1, len(updated_override_me.permissions))
-        birth_names = self.get_table(name="birth_names")
-        self.assertEqual(
-            birth_names.perm, updated_override_me.permissions[0].view_menu.name
-        )
-        self.assertEqual(
-            "datasource_access", updated_override_me.permissions[0].permission.name
-        )
+            response = self.client.post(
+                "/superset/override_role_permissions/",
+                data=json.dumps(perm_data),
+                content_type="application/json",
+            )
+            self.assertEqual(201, response.status_code)
+            updated_override_me = security_manager.find_role("override_me")
+            self.assertEqual(1, len(updated_override_me.permissions))
+            birth_names = self.get_table(name="birth_names")
+            self.assertEqual(
+                birth_names.perm, updated_override_me.permissions[0].view_menu.name
+            )
+            self.assertEqual(
+                "datasource_access", updated_override_me.permissions[0].permission.name
+            )


betodealmeida · 2022-10-20T23:53:05Z

superset/models/core.py

@@ -362,6 +362,18 @@ def get_effective_user(self, object_url: URL) -> Optional[str]:
            else None
        )

+    @contextmanager
+    def get_sqla_engine_with_context(


Eventually I think we want to rename this to get_sqla_engine, since we want this to be the one and only way to create an engine.

betodealmeida · 2022-10-20T23:57:26Z

tests/integration_tests/reports/commands_tests.py

+    with database.get_sqla_engine_with_context() as engine:
+        engine.execute("CREATE TABLE test_table AS SELECT 1 as first, 2 as second")
+        engine.execute("INSERT INTO test_table (first, second) VALUES (1, 2)")
+        engine.execute("INSERT INTO test_table (first, second) VALUES (3, 4)")

    yield db.session
    database.get_sqla_engine().execute("DROP TABLE test_table")


hughhhh · 2022-10-21T17:26:13Z

What's the plan for modifying the files in superset/?

Create get_sqla_engine_with_context() function and switch out places in our testing suite that straight forward (this PR)
Update get_sqla_engine() to _get_sqla_engine() making it a private function and switch out all remaining places that get_sqla_engine() to use get_sqla_engine_with_context()
- open to changing the name to get_sqla_engine for the contextmanager and then change the name of the original function to something else

I initially wanted to take incremental approach with 1 first to verify the get_sqla_engine_with_context won't break things then have a follow up PR that will close up the remaining work let me know what you think

@betodealmeida

Antonio-RiveroMartnez

LGTM!

betodealmeida

Looks great!

create function for get_sqla_engine with context

f6490fa

pull-request-size bot added the size/M label Oct 12, 2022

hughhhh marked this pull request as ready for review October 14, 2022 17:14

hughhhh added 3 commits October 15, 2022 11:41

Merge branch 'master' of https://github.com/apache/superset into ssh-…

fc0b0d5

…tunnel-refactor-get-sqla-engine

refactor

897c921

Merge branch 'master' of https://github.com/apache/superset into ssh-…

b824ab3

…tunnel-refactor-get-sqla-engine

pull-request-size bot added size/S and removed size/M labels Oct 17, 2022

update test 1

2497b72

pull-request-size bot added size/L and removed size/S labels Oct 17, 2022

hughhhh added 4 commits October 17, 2022 15:54

go with original pattern

296d8f2

oops

83130ca

fix reference

4b4613f

Merge branch 'master' of https://github.com/apache/superset into ssh-…

ca24e25

…tunnel-refactor-get-sqla-engine

hughhhh force-pushed the ssh-tunnel-refactor-get-sqla-engine branch 2 times, most recently from d53079d to 03029dd Compare October 19, 2022 23:56

update other test

c006017

hughhhh force-pushed the ssh-tunnel-refactor-get-sqla-engine branch from 03029dd to c006017 Compare October 20, 2022 00:03

hughhhh requested review from betodealmeida, eschutho, AAfghahi and pkdotson October 20, 2022 14:20

Antonio-RiveroMartnez reviewed Oct 20, 2022

View reviewed changes

betodealmeida reviewed Oct 20, 2022

View reviewed changes

hughhhh added 2 commits October 21, 2022 10:57

refactor

cbb400c

address comments

63d4a9c

Antonio-RiveroMartnez approved these changes Oct 24, 2022

View reviewed changes

betodealmeida approved these changes Oct 25, 2022

View reviewed changes

hughhhh merged commit 7600da8 into master Oct 25, 2022

mistercrunch added the 🚢 2.1.3 label Feb 18, 2024

mistercrunch added 🏷️ bot A label used by `supersetbot` to keep track of which PR where auto-tagged with release labels 🚢 2.1.0 and removed 🚢 2.1.3 labels Mar 13, 2024

mistercrunch deleted the ssh-tunnel-refactor-get-sqla-engine branch March 26, 2024 16:20

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: create function for get_sqla_engine with context #21790

feat: create function for get_sqla_engine with context #21790

hughhhh commented Oct 12, 2022 •

edited

Loading

codecov bot commented Oct 12, 2022 •

edited

Loading

Antonio-RiveroMartnez Oct 20, 2022

betodealmeida Oct 20, 2022

Antonio-RiveroMartnez Oct 21, 2022

Antonio-RiveroMartnez Oct 20, 2022 •

edited

Loading

betodealmeida Oct 20, 2022

Antonio-RiveroMartnez Oct 21, 2022

betodealmeida left a comment

betodealmeida Oct 20, 2022

betodealmeida Oct 20, 2022

betodealmeida Oct 20, 2022

betodealmeida Oct 20, 2022

hughhhh commented Oct 21, 2022 •

edited

Loading

Antonio-RiveroMartnez left a comment

betodealmeida left a comment

feat: create function for get_sqla_engine with context #21790

feat: create function for get_sqla_engine with context #21790

Conversation

hughhhh commented Oct 12, 2022 • edited Loading

SUMMARY

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TESTING INSTRUCTIONS

ADDITIONAL INFORMATION

codecov bot commented Oct 12, 2022 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Antonio-RiveroMartnez Oct 20, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

betodealmeida left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hughhhh commented Oct 21, 2022 • edited Loading

Antonio-RiveroMartnez left a comment

Choose a reason for hiding this comment

betodealmeida left a comment

Choose a reason for hiding this comment

hughhhh commented Oct 12, 2022 •

edited

Loading

codecov bot commented Oct 12, 2022 •

edited

Loading

Antonio-RiveroMartnez Oct 20, 2022 •

edited

Loading

hughhhh commented Oct 21, 2022 •

edited

Loading