Upgrade to apache-superset==3.1.0

.github/workflows/tests.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-20.04
     strategy:
       matrix:
-        python-version: [3.8, 3.9]
+        python-version: [3.9, '3.10', 3.11]
     env:
       PYTHONPATH: ${{ github.workspace }}
     services:
@@ -49,7 +49,6 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install .
           pip install -r requirements_test.txt
       - name: Create shared_dir
         run: |

hq_superset/__init__.py

@@ -14,6 +14,7 @@ def flask_app_mutator(app):
     # Import the views (which assumes the app is initialized) here
     # return
     from superset.extensions import appbuilder
+
     from . import hq_domain, views
     appbuilder.add_view(views.HQDatasourceView, 'Update HQ Datasource', menu_cond=lambda *_: False)
     appbuilder.add_view(views.SelectDomainView, 'Select a Domain', menu_cond=lambda *_: False)

hq_superset/hq_domain.py

@@ -6,6 +6,7 @@
 def before_request_hook():
     return ensure_domain_selected()
 
+
 def after_request_hook(response):
     # On logout clear domain cookie
     logout_views = [
@@ -16,6 +17,7 @@ def after_request_hook(response):
         response.set_cookie('hq_domain', '', expires=0)
     return response
 
+
 DOMAIN_EXCLUDED_VIEWS = [
     "AuthOAuthView.login",
     "AuthOAuthView.logout",
@@ -28,9 +30,11 @@ def after_request_hook(response):
     "static",
 ]
 
+
 def is_user_admin():
-    from superset.views.base import is_user_admin
-    return is_user_admin()
+    from superset import security_manager
+    return security_manager.is_admin()
+
 
 def ensure_domain_selected():
     # Check if a hq_domain cookie is set

hq_superset/oauth.py

@@ -7,7 +7,6 @@
 from superset.security import SupersetSecurityManager
 
 from .utils import (
-    DOMAIN_PREFIX,
     SESSION_OAUTH_RESPONSE_KEY,
     SESSION_USER_DOMAINS_KEY,
 )

hq_superset/tasks.py

@@ -1,7 +1,8 @@
+import logging
 import os
+
 from superset.extensions import celery_app
 
-import logging
 from .utils import AsyncImportHelper
 
 logger = logging.getLogger(__name__)

hq_superset/tests/base_test.py

@@ -6,11 +6,12 @@
 import os
 import shutil
 
-from sqlalchemy.sql import text
-from flask_appbuilder import SQLA
 from flask_testing import TestCase
+from sqlalchemy.sql import text
 from superset.app import create_app
-from hq_superset.utils import get_hq_database, DOMAIN_PREFIX
+
+from hq_superset.utils import DOMAIN_PREFIX, get_hq_database
+
 from .utils import setup_hq_db
 
 superset_test_home = os.path.join(os.path.dirname(__file__), ".test_superset")
@@ -39,14 +40,14 @@ def setUp(self):
 
     def tearDown(self):
         # Drop HQ DB Schemas
-        engine = self.hq_db.get_sqla_engine()
-        with engine.connect() as connection:
-            results = connection.execute(text("SELECT schema_name FROM information_schema.schemata"))
-            domain_schemas = []
-            for schema, in results.fetchall():
-                if schema.startswith(DOMAIN_PREFIX):
-                    domain_schemas.append(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE; COMMIT;')
-            if domain_schemas:
-                sql = "; ".join(domain_schemas) + ";"
-                connection.execute(text(sql))
+        with self.hq_db.get_sqla_engine_with_context() as engine:
+            with engine.connect() as connection:
+                results = connection.execute(text("SELECT schema_name FROM information_schema.schemata"))
+                domain_schemas = []
+                for schema, in results.fetchall():
+                    if schema.startswith(DOMAIN_PREFIX):
+                        domain_schemas.append(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE; COMMIT;')
+                if domain_schemas:
+                    sql = "; ".join(domain_schemas) + ";"
+                    connection.execute(text(sql))
         super(HQDBTestCase, self).tearDown()

hq_superset/tests/test_hq_domain.py

@@ -2,14 +2,22 @@
 
 from flask import g
 
-from hq_superset.hq_domain import (user_domains,
-    is_valid_user_domain, ensure_domain_selected,
-    DOMAIN_EXCLUDED_VIEWS, before_request_hook, after_request_hook)
-from hq_superset.utils import (SESSION_USER_DOMAINS_KEY, DomainSyncUtil,
-        get_hq_database, get_schema_name_for_domain)
+from hq_superset.hq_domain import (
+    DOMAIN_EXCLUDED_VIEWS,
+    after_request_hook,
+    before_request_hook,
+    ensure_domain_selected,
+    is_valid_user_domain,
+    user_domains,
+)
+from hq_superset.utils import (
+    SESSION_USER_DOMAINS_KEY,
+    DomainSyncUtil,
+    get_schema_name_for_domain,
+)
+
+from .base_test import HQDBTestCase, SupersetTestCase
 from .utils import setup_hq_db
-from .base_test import SupersetTestCase, HQDBTestCase
-
 
 MOCK_DOMAIN_SESSION = {
     SESSION_USER_DOMAINS_KEY:[
@@ -112,11 +120,11 @@ def setUp(self):
 
     def test_schema_gets_created(self):
         schema_name = get_schema_name_for_domain(self.domain)
-        engine = self.hq_db.get_sqla_engine()
-        self.assertFalse(
-            engine.dialect.has_schema(engine, schema_name),
-        )
-        DomainSyncUtil._ensure_schema_created(self.domain)
-        self.assertTrue(
-            engine.dialect.has_schema(engine, schema_name),
-        )
+        with self.hq_db.get_sqla_engine_with_context() as engine:
+            self.assertFalse(
+                engine.dialect.has_schema(engine, schema_name),
+            )
+            DomainSyncUtil._ensure_schema_created(self.domain)
+            self.assertTrue(
+                engine.dialect.has_schema(engine, schema_name),
+            )

hq_superset/tests/test_oauth.py

@@ -1,12 +1,14 @@
 import datetime
-import jwt
+from unittest.mock import patch
 
-from unittest.mock import patch, MagicMock
 from flask import session
 
 from hq_superset.oauth import OAuthSessionExpired, get_valid_cchq_oauth_token
-from hq_superset.utils import (SESSION_USER_DOMAINS_KEY, 
-    SESSION_OAUTH_RESPONSE_KEY)
+from hq_superset.utils import (
+    SESSION_OAUTH_RESPONSE_KEY,
+    SESSION_USER_DOMAINS_KEY,
+)
+
 from .base_test import SupersetTestCase
 
 

hq_superset/tests/test_utils.py

@@ -1,6 +1,7 @@
 import doctest
 
 from hq_superset.utils import get_column_dtypes
+
 from .utils import TEST_DATASOURCE
 
 

hq_superset/tests/test_views.py

@@ -1,19 +1,19 @@
-import datetime
 import json
-import jwt
 import os
 import pickle
-
 from io import StringIO
-from unittest.mock import patch, MagicMock
-from flask import session, redirect
+from unittest.mock import patch
+
+import jwt
+from flask import redirect, session
 from sqlalchemy.sql import text
-from superset.connectors.sqla.models import SqlaTable
 
-from hq_superset.oauth import OAuthSessionExpired, get_valid_cchq_oauth_token
-from hq_superset.utils import (SESSION_USER_DOMAINS_KEY, 
-    SESSION_OAUTH_RESPONSE_KEY, get_schema_name_for_domain)
-from .base_test import SupersetTestCase, HQDBTestCase
+from hq_superset.utils import (
+    SESSION_USER_DOMAINS_KEY,
+    get_schema_name_for_domain,
+)
+
+from .base_test import HQDBTestCase
 from .utils import TEST_DATASOURCE
 
 
@@ -117,7 +117,9 @@ def setUp(self):
         self.app.appbuilder.sm.oauth_remotes = {"commcare": self.oauth_mock}
 
         gamma_role = self.app.appbuilder.sm.find_role('Gamma')
-        self.user = self.app.appbuilder.sm.add_user(**self.oauth_mock.user_json, role=[gamma_role])
+        self.user = self.app.appbuilder.sm.find_user(self.oauth_mock.user_json['username'])
+        if not self.user:
+            self.user = self.app.appbuilder.sm.add_user(**self.oauth_mock.user_json, role=[gamma_role])
 
     def login(self, client):
         # bypass oauth-workflow by skipping login and oauth flow
@@ -147,11 +149,11 @@ def _assert_hq_domain_cookie(self, client, response, domain):
             self.assertTrue('hq_domain' not in response.request.cookies)
 
     def _assert_pg_schema_exists(self, domain, exists):
-        engine = self.hq_db.get_sqla_engine()
-        self.assertEqual(
-            engine.dialect.has_schema(engine, get_schema_name_for_domain(domain)),
-            exists
-        )
+        with self.hq_db.get_sqla_engine_with_context() as engine:
+            self.assertEqual(
+                engine.dialect.has_schema(engine, get_schema_name_for_domain(domain)),
+                exists
+            )
 
     def test_redirects_to_domain_select_after_login(self):
         with self.app.test_client() as client:
@@ -232,7 +234,7 @@ def test_datasource_upload(self, *args):
     def test_trigger_datasource_refresh(self, *args):
         from hq_superset.views import (
             ASYNC_DATASOURCE_IMPORT_LIMIT_IN_BYTES,
-            trigger_datasource_refresh
+            trigger_datasource_refresh,
         )
         domain = 'test1'
         ds_name = 'ds_name'
@@ -304,15 +306,15 @@ def _test_upload(test_data, expected_output):
                 self.assertEqual(datasets['result'][0]['schema'], get_schema_name_for_domain('test1'))
                 self.assertEqual(datasets['result'][0]['table_name'], ucr_id)
                 self.assertEqual(datasets['result'][0]['description'], ds_name)
-                engine = self.hq_db.get_sqla_engine()
-                with engine.connect() as connection:
-                    result = connection.execute(text(
-                        'SELECT doc_id FROM hqdomain_test1.test1_ucr1'
-                    )).fetchall()
-                    self.assertEqual(
-                        result,
-                        expected_output
-                    )
+                with self.hq_db.get_sqla_engine_with_context() as engine:
+                    with engine.connect() as connection:
+                        result = connection.execute(text(
+                            'SELECT doc_id FROM hqdomain_test1.test1_ucr1'
+                        )).fetchall()
+                        self.assertEqual(
+                            result,
+                            expected_output
+                        )
                 # Check that updated dataset is reflected in the list view
                 client.get('/hq_datasource/list/', follow_redirects=True)
                 self.assert_context('ucr_id_to_pks', {'test1_ucr1': 1})

hq_superset/tests/utils.py

@@ -1,12 +1,15 @@
 from functools import wraps
+
 from sqlalchemy.orm.exc import NoResultFound
-from hq_superset.utils import get_hq_database, HQ_DB_CONNECTION_NAME
+
+from hq_superset.utils import HQ_DB_CONNECTION_NAME, get_hq_database
 
 # @pytest.fixture(scope="session", autouse=True)
 # def manage_ucr_db(request):
 #     # setup_ucr_db()
 #     request.addfinalizer(clear_ucr_db)
 
+
 def unit_testing_only(fn):
     import superset
 
@@ -21,13 +24,12 @@ def inner(*args, **kwargs):
 
 @unit_testing_only
 def setup_hq_db():
-    from superset.databases.commands.create import CreateDatabaseCommand
     import superset
+    from superset.commands.database.create import CreateDatabaseCommand
     try:
         get_hq_database()
     except NoResultFound:
         CreateDatabaseCommand(
-            None, 
             {
                 'sqlalchemy_uri': superset.app.config.get('HQ_DATA_DB'),
                 'engine': 'PostgreSQL', 

hq_superset/utils.py

@@ -1,22 +1,20 @@
 import os
-import pandas
-import sqlalchemy
-
-
 from contextlib import contextmanager
 from datetime import date, datetime
-from superset.extensions import cache_manager
-from flask_login import current_user
 from zipfile import ZipFile
 
-
+import pandas
+import sqlalchemy
+from flask_login import current_user
+from superset.extensions import cache_manager
 
 DOMAIN_PREFIX = "hqdomain_"
 SESSION_USER_DOMAINS_KEY = "user_hq_domains"
 SESSION_OAUTH_RESPONSE_KEY = "oauth_response"
 HQ_DB_CONNECTION_NAME = "HQ Data"
-# ~5MB
-ASYNC_DATASOURCE_IMPORT_LIMIT_IN_BYTES = 5000000
+
+ASYNC_DATASOURCE_IMPORT_LIMIT_IN_BYTES = 5_000_000  # ~5MB
+
 
 def get_datasource_export_url(domain, datasource_id):
     return f"a/{domain}/configurable_reports/data_sources/export/{datasource_id}/?format=csv"
@@ -153,9 +151,9 @@ def _ensure_schema_perm_created(self, domain):
     def _ensure_schema_created(domain):
         schema_name = get_schema_name_for_domain(domain)
         database = get_hq_database()
-        engine = database.get_sqla_engine()
-        if not engine.dialect.has_schema(engine, schema_name):
-            engine.execute(sqlalchemy.schema.CreateSchema(schema_name))
+        with database.get_sqla_engine_with_context() as engine:
+            if not engine.dialect.has_schema(engine, schema_name):
+                engine.execute(sqlalchemy.schema.CreateSchema(schema_name))
 
     def re_eval_roles(self, existing_roles, new_domain_role):
         # Filter out other domain roles
@@ -182,8 +180,6 @@ def sync_domain_role(self, domain):
         self.sm.get_session.commit()
 
 
-
-
 @contextmanager
 def get_datasource_file(path):
     with ZipFile(path) as zipfile:
@@ -205,6 +201,7 @@ def download_datasource(provider, oauth_token, domain, datasource_id):
 
     return path, len(response.content)
 
+
 def get_datasource_defn(provider, oauth_token, domain, datasource_id):
     url = get_datasource_details_url(domain, datasource_id)
     response = provider.get(url, token=oauth_token)