preset-io · joaopamaral · Sep 11, 2024 · Sep 12, 2024 · Sep 12, 2024 · Sep 12, 2024
diff --git a/.gitignore b/.gitignore
@@ -56,3 +56,5 @@ MANIFEST
 .pyre/
 .pyre_configuration
 .watchmanconfig
+
+backend-sdk
diff --git a/Makefile b/Makefile
@@ -1,12 +1,11 @@
 pyenv: .python-version
 
 .python-version: setup.cfg
- if [ -z "`pyenv virtualenvs | grep backend-sdk`" ]; then\
- pyenv virtualenv backend-sdk;\
- fi
- if [ ! -f .python-version ]; then\
- pyenv local backend-sdk;\
+ if [ ! -d "backend-sdk" ]; then \
+ pyenv local 3.11; \
+ python -m venv backend-sdk; \
  fi
+ backend-sdk/bin/activate; \
  pip install -e '.[testing]'
  touch .python-version
 

diff --git a/setup.cfg b/setup.cfg
@@ -71,7 +71,7 @@ install_requires =
  requests>=2.26.0
  rich>=12.3.0
  sqlalchemy>=1.4,<2
- sqlglot>=19,<23
+ sqlglot>=23.0.2,<24
  tabulate>=0.8.9
  typing-extensions>=4.0.1
  yarl>=1.7.2

diff --git a/src/preset_cli/api/clients/superset.py b/src/preset_cli/api/clients/superset.py
@@ -767,6 +767,8 @@ def import_zip(
  )
  validate_response(response)
 
+ _logger.debug(response.text)
+
  payload = response.json()
 
  return payload["message"] == "OK"

diff --git a/src/preset_cli/auth/main.py b/src/preset_cli/auth/main.py
@@ -1,13 +1,14 @@
 """
 Mechanisms for authentication and authorization.
 """
-
+import logging
 from typing import Any, Dict
 
 from requests import Response, Session
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
 
+_logger = logging.getLogger(__name__)
 
 class Auth: # pylint: disable=too-few-public-methods
  """
@@ -46,6 +47,8 @@ def reauth(self, r: Response, *args: Any, **kwargs: Any) -> Response:
  if r.status_code != 401:
  return r
 
+ _logger.info('Token expired. Re-authenticating...')
+
  try:
  self.auth()
  except NotImplementedError:

diff --git a/src/preset_cli/auth/superset.py b/src/preset_cli/auth/superset.py
@@ -4,7 +4,6 @@
 
 from typing import Dict, Optional
 
-from bs4 import BeautifulSoup
 from yarl import URL
 
 from preset_cli.auth.main import Auth
@@ -28,24 +27,31 @@ def __init__(self, baseurl: URL, username: str, password: Optional[str] = None):
  def get_headers(self) -> Dict[str, str]:
  return {"X-CSRFToken": self.csrf_token} if self.csrf_token else {}
 
+ def get_access_token(self):
+ body = {"username": self.username, "password": self.password, "provider": "ldap"}
+ if "Referer" in self.session.headers:
+ del self.session.headers["Referer"]
+ response = self.session.post(self.baseurl / "api/v1/security/login", json=body)
+ response.raise_for_status()
+ return response.json()["access_token"]
+
+ def get_csrf_token(self):
+ response = self.session.get(self.baseurl / "api/v1/security/csrf_token/")
+ response.raise_for_status()
+ return response.json()["result"]
+
  def auth(self) -> None:
  """
  Login to get CSRF token and cookies.
  """
- data = {"username": self.username, "password": self.password}
+ self.session.headers["Authorization"] = f"Bearer {self.get_access_token()}"
+ csrf_token = self.get_csrf_token()
 
- 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
  if csrf_token:
  self.session.headers["X-CSRFToken"] = csrf_token
- data["csrf_token"] = csrf_token
+ self.session.headers["Referer"] = str(self.baseurl / "api/v1/security/csrf_token/")
  self.csrf_token = csrf_token
 
- # set cookies
- self.session.post(self.baseurl / "login/", data=data)
-
 
 class SupersetJWTAuth(TokenAuth): # pylint: disable=abstract-method
  """

diff --git a/tests/auth/superset_test.py b/tests/auth/superset_test.py
@@ -14,11 +14,13 @@ def test_username_password_auth(requests_mock: Mocker) -> None:
  Tests for the username/password authentication mechanism.
  """
  csrf_token = "CSFR_TOKEN"
+ access_token = "ACCESS_TOKEN"
  requests_mock.get(
- "https://superset.example.org/login/",
- text=f'<html><body><input id="csrf_token" value="{csrf_token}"></body></html>',
+ "https://superset.example.org/api/v1/security/csrf_token/",
+ json={'result': csrf_token},
  )
- requests_mock.post("https://superset.example.org/login/")
+ requests_mock.post("https://superset.example.org/api/v1/security/login",
+ json={'access_token': access_token})
 
  auth = UsernamePasswordAuth(
  URL("https://superset.example.org/"),
@@ -29,21 +31,18 @@ def test_username_password_auth(requests_mock: Mocker) -> None:
  "X-CSRFToken": csrf_token,
  }
 
- assert (
- requests_mock.last_request.text
- == "username=admin&password=password123&csrf_token=CSFR_TOKEN"
- )
-
 
 def test_username_password_auth_no_csrf(requests_mock: Mocker) -> None:
  """
  Tests for the username/password authentication mechanism.
  """
+ access_token = "ACCESS_TOKEN"
  requests_mock.get(
- "https://superset.example.org/login/",
- text="<html><body>WTF_CSRF_ENABLED = False</body></html>",
+ "https://superset.example.org/api/v1/security/csrf_token/",
+ json={'result': None},
  )
- requests_mock.post("https://superset.example.org/login/")
+ requests_mock.post("https://superset.example.org/api/v1/security/login",
+ json={'access_token': access_token})
 
  auth = UsernamePasswordAuth(
  URL("https://superset.example.org/"),
@@ -53,9 +52,6 @@ def test_username_password_auth_no_csrf(requests_mock: Mocker) -> None:
  # pylint: disable=use-implicit-booleaness-not-comparison
  assert auth.get_headers() == {}
 
- assert requests_mock.last_request.text == "username=admin&password=password123"
-
-
 def test_jwt_auth_superset(mocker: MockerFixture) -> None:
  """
  Test the ``JWTAuth`` authentication mechanism for Superset tenant.