Skip to content

Commit

Permalink
Clean up imports
Browse files Browse the repository at this point in the history
  • Loading branch information
kaapstorm committed Jan 28, 2024
1 parent 43ee6c9 commit 7239efa
Show file tree
Hide file tree
Showing 13 changed files with 64 additions and 55 deletions.
1 change: 1 addition & 0 deletions hq_superset/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 4 additions & 0 deletions hq_superset/hq_domain.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
def before_request_hook():
return ensure_domain_selected()


def after_request_hook(response):
# On logout clear domain cookie
logout_views = [
Expand All @@ -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",
Expand All @@ -28,10 +30,12 @@ def after_request_hook(response):
"static",
]


def 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
# Ensure necessary roles, permissions and DB schemas are created for the domain
Expand Down
1 change: 0 additions & 1 deletion hq_superset/oauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from superset.security import SupersetSecurityManager

from .utils import (
DOMAIN_PREFIX,
SESSION_OAUTH_RESPONSE_KEY,
SESSION_USER_DOMAINS_KEY,
)
Expand Down
3 changes: 2 additions & 1 deletion hq_superset/tasks.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import logging
import os

from superset.extensions import celery_app

import logging
from .utils import AsyncImportHelper

logger = logging.getLogger(__name__)
Expand Down
7 changes: 4 additions & 3 deletions hq_superset/tests/base_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
22 changes: 15 additions & 7 deletions hq_superset/tests/test_hq_domain.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:[
Expand Down
10 changes: 6 additions & 4 deletions hq_superset/tests/test_oauth.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down
1 change: 1 addition & 0 deletions hq_superset/tests/test_utils.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import doctest

from hq_superset.utils import get_column_dtypes

from .utils import TEST_DATASOURCE


Expand Down
9 changes: 5 additions & 4 deletions hq_superset/tests/test_views.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,18 @@
import json
import jwt
import os
import pickle

from io import StringIO
from unittest.mock import patch
from flask import session, redirect

import jwt
from flask import redirect, session
from sqlalchemy.sql import text

from hq_superset.utils import (
SESSION_USER_DOMAINS_KEY,
get_schema_name_for_domain,
)

from .base_test import HQDBTestCase
from .utils import TEST_DATASOURCE

Expand Down Expand Up @@ -233,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'
Expand Down
7 changes: 5 additions & 2 deletions hq_superset/tests/utils.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -21,8 +24,8 @@ def inner(*args, **kwargs):

@unit_testing_only
def setup_hq_db():
from superset.commands.database.create import CreateDatabaseCommand
import superset
from superset.commands.database.create import CreateDatabaseCommand
try:
get_hq_database()
except NoResultFound:
Expand Down
19 changes: 8 additions & 11 deletions hq_superset/utils.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand Down
32 changes: 13 additions & 19 deletions hq_superset/views.py
Original file line number Diff line number Diff line change
@@ -1,47 +1,41 @@
import logging
import ast
import logging
import os

import pandas as pd
import superset
import sqlalchemy

from datetime import datetime
from io import BytesIO
from sqlalchemy.dialects import postgresql
from flask import Response, abort, g, redirect, request, flash, url_for

import superset
from flask import Response, abort, flash, g, redirect, request, url_for
from flask_appbuilder import expose
from flask_appbuilder.security.decorators import has_access, permission_name
from sqlalchemy.dialects import postgresql
from superset import db
from superset.connectors.sqla.models import SqlaTable
from superset.commands.dataset.delete import DeleteDatasetCommand
from superset.commands.dataset.exceptions import (
DatasetDeleteFailedError,
DatasetForbiddenError,
DatasetNotFoundError,
)
from superset.connectors.sqla.models import SqlaTable
from superset.models.core import Database
from superset.sql_parse import Table
from superset.views.base import BaseSupersetView
from superset.extensions import cache_manager

from .hq_domain import user_domains
from .oauth import get_valid_cchq_oauth_token
from .tasks import refresh_hq_datasource_task
from .utils import (
get_column_dtypes,
get_datasource_details_url,
get_datasource_export_url,
get_datasource_list_url,
get_schema_name_for_domain,
get_hq_database,
parse_date,
ASYNC_DATASOURCE_IMPORT_LIMIT_IN_BYTES,
AsyncImportHelper,
DomainSyncUtil,
get_datasource_file,
download_datasource,
get_column_dtypes,
get_datasource_defn,
ASYNC_DATASOURCE_IMPORT_LIMIT_IN_BYTES,
get_datasource_file,
get_datasource_list_url,
get_hq_database,
get_schema_name_for_domain,
parse_date,
)

logger = logging.getLogger(__name__)
Expand Down
3 changes: 0 additions & 3 deletions setup.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,3 @@
"""
Setup.py for superset-patchup
"""
from setuptools import setup, find_packages

setup(
Expand Down

0 comments on commit 7239efa

Please sign in to comment.