Merge branch 'main' into clean_get_resource

dev-requirements.txt

@@ -10,6 +10,7 @@ aiosignal==1.2.0
 appdirs==1.4.4
 async-timeout==4.0.2
 attrs==22.1.0
+backoff==2.2.1
 beautifulsoup4==4.11.1
 build==0.8.0
 certifi==2021.10.8

requirements.txt

@@ -10,6 +10,7 @@ aiosignal==1.2.0
 appdirs==1.4.4
 async-timeout==4.0.2
 attrs==22.1.0
+backoff==2.2.1
 beautifulsoup4==4.11.1
 certifi==2021.10.8
 charset-normalizer==2.0.12

setup.cfg

@@ -57,6 +57,7 @@ install_requires =
     Cython>=0.29.26
     PyYAML>=6.0
     appdirs>=1.4.4
+    backoff>=2.2.1 
     beautifulsoup4>=4.10.0
     click>=8.0.3
     jinja2>=3.0.3

src/preset_cli/api/clients/dbt.py

@@ -581,7 +581,7 @@ def __init__(self, auth: Auth):
         self.graphql_client = GraphqlClient(endpoint=GRAPHQL_ENDPOINT)
         self.baseurl = REST_ENDPOINT
 
-        self.session = auth.get_session()
+        self.session = auth.session
         self.session.headers.update(auth.get_headers())
         self.session.headers["User-Agent"] = "Preset CLI"
         self.session.headers["X-Client-Version"] = __version__

src/preset_cli/api/clients/preset.py

@@ -38,7 +38,7 @@ def __init__(self, baseurl: Union[str, URL], auth: Auth):
         self.baseurl = URL(baseurl)
         self.auth = auth
 
-        self.session = auth.get_session()
+        self.session = auth.session
         self.session.headers.update(auth.get_headers())
         self.session.headers["User-Agent"] = "Preset CLI"
         self.session.headers["X-Client-Version"] = __version__

src/preset_cli/api/clients/superset.py

@@ -221,7 +221,7 @@ def __init__(self, baseurl: Union[str, URL], auth: Auth):
         self.baseurl = URL(baseurl)
         self.auth = auth
 
-        self.session = auth.get_session()
+        self.session = auth.session
         self.session.headers.update(auth.get_headers())
         self.session.headers["Referer"] = str(self.baseurl)
         self.session.headers["User-Agent"] = f"Apache Superset Client ({__version__})"

src/preset_cli/auth/jwt.py

@@ -8,7 +8,7 @@
 from preset_cli.auth.token import TokenAuth
 
 
-class JWTAuth(TokenAuth):  # pylint: disable=too-few-public-methods
+class JWTAuth(TokenAuth):  # pylint: disable=too-few-public-methods, abstract-method
     """
     Auth via JWT.
     """

src/preset_cli/auth/main.py

@@ -2,9 +2,9 @@
 Mechanisms for authentication and authorization.
 """
 
-from typing import Dict
+from typing import Any, Dict
 
-import requests
+from requests import Response, Session
 
 
 class Auth:  # pylint: disable=too-few-public-methods
@@ -13,17 +13,34 @@ class Auth:  # pylint: disable=too-few-public-methods
     """
 
     def __init__(self):
-        self.session = requests.Session()
-        self.headers = {}
+        self.session = Session()
+        self.session.hooks["response"].append(self.reauth)
 
-    def get_session(self) -> requests.Session:
+    def get_headers(self) -> Dict[str, str]:  # pylint: disable=no-self-use
         """
-        Return a session.
+        Return headers for auth.
         """
-        return self.session
+        return {}
 
-    def get_headers(self) -> Dict[str, str]:  # pylint: disable=no-self-use
+    def auth(self) -> None:
         """
-        Return headers for auth.
+        Perform authentication, fetching JWT tokens, CSRF tokens, cookies, etc.
         """
-        return self.headers
+        raise NotImplementedError("Must be implemented for reauthorizing")
+
+    # pylint: disable=invalid-name, unused-argument
+    def reauth(self, r: Response, *args: Any, **kwargs: Any) -> Response:
+        """
+        Catch 401 and re-auth.
+        """
+        if r.status_code != 401:
+            return r
+
+        try:
+            self.auth()
+        except NotImplementedError:
+            return r
+
+        self.session.headers.update(self.get_headers())
+        r.request.headers.update(self.get_headers())
+        return self.session.send(r.request, verify=False)

src/preset_cli/auth/password.py

@@ -2,7 +2,7 @@
 Mechanisms for authentication and authorization.
 """
 
-from typing import Optional
+from typing import Dict, Optional
 
 from bs4 import BeautifulSoup
 from yarl import URL
@@ -17,27 +17,30 @@ class UsernamePasswordAuth(Auth):  # pylint: disable=too-few-public-methods
 
     def __init__(self, baseurl: URL, username: str, password: Optional[str] = None):
         super().__init__()
-        self._do_login(baseurl, username, password)
-
-    def _do_login(
-        self,
-        baseurl: URL,
-        username: str,
-        password: Optional[str] = None,
-    ) -> None:
+
+        self.csrf_token: Optional[str] = None
+        self.baseurl = baseurl
+        self.username = username
+        self.password = password
+        self.auth()
+
+    def get_headers(self) -> Dict[str, str]:  # pylint: disable=no-self-use
+        return {"X-CSRFToken": self.csrf_token} if self.csrf_token else {}
+
+    def auth(self) -> None:
         """
         Login to get CSRF token and cookies.
         """
-        response = self.session.get(baseurl / "login/")
+        data = {"username": self.username, "password": self.password}
+
+        response = self.session.get(self.baseurl / "login/")
         soup = BeautifulSoup(response.text, "html.parser")
         input_ = soup.find("input", {"id": "csrf_token"})
         csrf_token = input_["value"] if input_ else None
-
-        data = {"username": username, "password": password}
-
         if csrf_token:
-            self.headers["X-CSRFToken"] = csrf_token
+            self.session.headers["X-CSRFToken"] = csrf_token
             data["csrf_token"] = csrf_token
+            self.csrf_token = csrf_token
 
         # set cookies
-        self.session.post(baseurl / "login/", data=data)
+        self.session.post(self.baseurl / "login/", data=data)

src/preset_cli/auth/preset.py

@@ -0,0 +1,59 @@
+"""
+Preset auth.
+"""
+
+from typing import Dict
+
+import yaml
+from yarl import URL
+
+from preset_cli.auth.lib import get_access_token, get_credentials_path
+from preset_cli.auth.main import Auth
+
+
+class JWTTokenError(Exception):
+    """
+    Exception raised when fetching the JWT fails.
+    """
+
+
+class PresetAuth(Auth):  # pylint: disable=too-few-public-methods
+    """
+    Auth via Preset access token and secret.
+
+    Automatically refreshes the JWT as needed.
+    """
+
+    def __init__(self, baseurl: URL, api_token: str, api_secret: str):
+        super().__init__()
+
+        self.baseurl = baseurl
+        self.api_token = api_token
+        self.api_secret = api_secret
+        self.auth()
+
+    def get_headers(self) -> Dict[str, str]:
+        return {"Authorization": f"Bearer {self.token}"}
+
+    def auth(self) -> None:
+        """
+        Fetch the JWT and store it.
+        """
+        try:
+            self.token = get_access_token(self.baseurl, self.api_token, self.api_secret)
+        except Exception as ex:  # pylint: disable=broad-except
+            raise JWTTokenError("Unable to fetch JWT") from ex
+
+    @classmethod
+    def from_stored_credentials(cls) -> "PresetAuth":
+        """
+        Build auth from stored credentials.
+        """
+        credentials_path = get_credentials_path()
+        if not credentials_path.exists():
+            raise Exception(f"Could not load credentials from {credentials_path}")
+
+        with open(credentials_path, encoding="utf-8") as input_:
+            credentials = yaml.load(input_, Loader=yaml.SafeLoader)
+
+        return PresetAuth(**credentials)

src/preset_cli/auth/token.py

@@ -7,7 +7,7 @@
 from preset_cli.auth.main import Auth
 
 
-class TokenAuth(Auth):  # pylint: disable=too-few-public-methods
+class TokenAuth(Auth):  # pylint: disable=too-few-public-methods, abstract-method
     """
     Auth via a token.
     """

src/preset_cli/cli/main.py

@@ -16,12 +16,8 @@
 from preset_cli.api.clients.preset import PresetClient
 from preset_cli.api.clients.superset import SupersetClient
 from preset_cli.auth.jwt import JWTAuth
-from preset_cli.auth.lib import (
-    get_access_token,
-    get_credentials_path,
-    store_credentials,
-)
-from preset_cli.auth.main import Auth
+from preset_cli.auth.lib import get_credentials_path, store_credentials
+from preset_cli.auth.preset import JWTTokenError, PresetAuth
 from preset_cli.cli.superset.main import superset
 from preset_cli.lib import setup_logging, split_comma
 
@@ -126,7 +122,9 @@ def preset_cli(  # pylint: disable=too-many-branches, too-many-locals, too-many-
         # The user is trying to auth themselves, so skip anything auth-related
         return
 
-    if jwt_token is None:
+    if jwt_token:
+        ctx.obj["AUTH"] = JWTAuth(jwt_token)
+    else:
         if api_token is None or api_secret is None:
             # check for stored credentials
             credentials_path = get_credentials_path()
@@ -161,10 +159,9 @@ def preset_cli(  # pylint: disable=too-many-branches, too-many-locals, too-many-
 
         api_token = cast(str, api_token)
         api_secret = cast(str, api_secret)
-
         try:
-            jwt_token = get_access_token(manager_api_url, api_token, api_secret)
-        except Exception:  # pylint: disable=broad-except
+            ctx.obj["AUTH"] = PresetAuth(manager_api_url, api_token, api_secret)
+        except JWTTokenError:
             click.echo(
                 click.style(
                     "Failed to auth using the provided credentials."
@@ -174,9 +171,6 @@ def preset_cli(  # pylint: disable=too-many-branches, too-many-locals, too-many-
             )
             sys.exit(1)
 
-    # store auth in context so it's used by the Superset SDK
-    ctx.obj["AUTH"] = JWTAuth(jwt_token) if jwt_token else Auth()
-
     if not workspaces and ctx.invoked_subcommand == "superset" and not is_help():
         client = PresetClient(ctx.obj["MANAGER_URL"], ctx.obj["AUTH"])
         click.echo("Choose one or more workspaces (eg: 1-3,5,8-):")

src/preset_cli/cli/superset/sync/dbt/exposures.py

@@ -57,7 +57,7 @@ def get_dashboard_depends_on(
 
     url = client.baseurl / "api/v1/dashboard" / str(dashboard["id"]) / "datasets"
 
-    session = client.auth.get_session()
+    session = client.auth.session
     headers = client.auth.get_headers()
     response = session.get(url, headers=headers)
     response.raise_for_status()
@@ -102,7 +102,7 @@ def sync_exposures(  # pylint: disable=too-many-locals
     for dataset in datasets:
         url = client.baseurl / "api/v1/dataset" / str(dataset["id"]) / "related_objects"
 
-        session = client.auth.get_session()
+        session = client.auth.session
         headers = client.auth.get_headers()
         response = session.get(url, headers=headers)
         response.raise_for_status()

src/preset_cli/cli/superset/sync/native/command.py

@@ -13,7 +13,9 @@
 from typing import Any, Dict, Iterator, Tuple
 from zipfile import ZipFile
 
+import backoff
 import click
+import requests
 import yaml
 from jinja2 import Template
 from sqlalchemy.engine import create_engine
@@ -213,14 +215,17 @@ def import_resources_individually(
     related_configs: Dict[str, Dict[Path, AssetConfig]] = {}
     for resource_name, get_related_uuids in imports:
         for path, config in configs.items():
-            if path.parts[1] == resource_name:
-                asset_configs = {path: config}
-                for uuid in get_related_uuids(config):
-                    asset_configs.update(related_configs[uuid])
-                _logger.info("Importing %s", path.relative_to("bundle"))
-                contents = {str(k): yaml.dump(v) for k, v in asset_configs.items()}
-                import_resources(contents, client, overwrite)
-                related_configs[config["uuid"]] = asset_configs
+            if path.parts[1] != resource_name:
+                continue
+
+            asset_configs = {path: config}
+            for uuid in get_related_uuids(config):
+                asset_configs.update(related_configs[uuid])
+
+            _logger.info("Importing %s", path.relative_to("bundle"))
+            contents = {str(k): yaml.dump(v) for k, v in asset_configs.items()}
+            import_resources(contents, client, overwrite)
+            related_configs[config["uuid"]] = asset_configs
 
 
 def get_charts_uuids(config: AssetConfig) -> Iterator[str]:
@@ -267,6 +272,13 @@ def prompt_for_passwords(path: Path, config: Dict[str, Any]) -> None:
         )
 
 
+@backoff.on_exception(
+    backoff.expo,
+    (requests.exceptions.ConnectionError, requests.exceptions.Timeout),
+    max_time=60,
+    max_tries=5,
+    logger=__name__,
+)
 def import_resources(
     contents: Dict[str, str],
     client: SupersetClient,

tests/auth/jwt_test.py

@@ -44,10 +44,6 @@ def test_jwt_auth_from_stored_credentials(mocker: MockerFixture) -> None:
         api_secret="SECRET",
     )
 
-    # can also pass a URL
-    auth = JWTAuth.from_stored_credentials()
-    assert auth.token == "JWT_TOKEN"
-
     # test for error
     get_credentials_path().exists.return_value = False
     with pytest.raises(Exception) as excinfo:

tests/auth/main_test.py

@@ -3,6 +3,7 @@
 """
 
 from pytest_mock import MockerFixture
+from requests_mock.mocker import Mocker
 
 from preset_cli.auth.main import Auth
 
@@ -11,7 +12,20 @@ def test_auth(mocker: MockerFixture) -> None:
     """
     Tests for the base class ``Auth``.
     """
-    requests = mocker.patch("preset_cli.auth.main.requests")
+    # pylint: disable=invalid-name
+    Session = mocker.patch("preset_cli.auth.main.Session")
 
     auth = Auth()
-    assert auth.get_session() == requests.Session()
+    assert auth.session == Session()
+
+
+def test_reauth(requests_mock: Mocker) -> None:
+    """
+    Test the ``reauth`` hook when authentication fails.
+    """
+    requests_mock.get("http://example.org/", status_code=401)
+
+    # the base class has no reauth
+    auth = Auth()
+    response = auth.session.get("http://example.org/")
+    assert response.status_code == 401