dataflow-core 2.0.9__tar.gz → 2.0.11__tar.gz

{dataflow_core-2.0.9 → dataflow_core-2.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataflow-core
-Version: 2.0.9
+Version: 2.0.11
 Summary: Dataflow core package
 Author: Dataflow
 Author-email: 

{dataflow_core-2.0.9 → dataflow_core-2.0.11}/authenticator/dataflowsupersetauthenticator.py

@@ -1,22 +1,22 @@
-from flask import redirect, request
+from flask import redirect, request, jsonify
 from flask_appbuilder.security.views import AuthDBView
-from superset.security import SupersetSecurityManager
 from flask_appbuilder.security.views import expose
 from flask_login import login_user
-from .package.configuration import ConfigurationManager
-from .package.models.database import DatabaseManager
-from .package.models import (
-    user as m_user,
-    session as m_session
-    )
+from flask_jwt_extended import (
+    create_access_token
+)
+from flask_appbuilder.const import (
+    API_SECURITY_ACCESS_TOKEN_KEY
+)
+from superset.security import SupersetSecurityManager
 from dataflow.dataflow import Dataflow
 
-class CustomAuthDBView(AuthDBView):
+class DataflowAuthDBView(AuthDBView):
     def __init__(self):
         self.dataflow = Dataflow()
 
     @expose('/login/', methods=['GET'])
-    def login(self):
+    def login_redirect(self):
         try:
             session_id = request.cookies.get('dataflow_session')
             
@@ -30,8 +30,7 @@ class CustomAuthDBView(AuthDBView):
                     first_name=user_details.get("first_name", ""),
                     last_name=user_details.get("last_name", ""), 
                     email=user_details.get("email", ""), 
-                    role=self.appbuilder.sm.find_role('Admin'), 
-                    password=""
+                    role=self.appbuilder.sm.find_role('Admin')
                 )
                 if user:
                     login_user(user, remember=False)
@@ -41,8 +40,27 @@ class CustomAuthDBView(AuthDBView):
         except Exception as e:
             return super().login()
 
+    @expose("/login", methods=["POST"])
+    def login_token(self):
+        login_payload = request.get_json()
+        user = self.appbuilder.sm.find_user(username=login_payload["user_name"])
 
-class CustomSecurityManager(SupersetSecurityManager):
-    authdbview = CustomAuthDBView
+        if not user:
+            user = self.appbuilder.sm.add_user(
+                username=login_payload["user_name"], 
+                first_name=login_payload.get("first_name", ""),
+                last_name=login_payload.get("last_name", ""), 
+                email=login_payload.get("email", ""), 
+                role=self.appbuilder.sm.find_role('Admin')
+            )
+        
+        resp = dict()
+        resp[API_SECURITY_ACCESS_TOKEN_KEY] = create_access_token(
+            identity=str(user.id), fresh=True
+        )
+        return jsonify(resp)
+    
+class DataflowSecurityManager(SupersetSecurityManager):
+    authdbview = DataflowAuthDBView
     def __init__(self, appbuilder):
-        super(CustomSecurityManager, self).__init__(appbuilder)
+        super(DataflowSecurityManager, self).__init__(appbuilder)

{dataflow_core-2.0.9 → dataflow_core-2.0.11}/authenticator/package/models/user.py

@@ -1,5 +1,5 @@
 """models.py"""
-from sqlalchemy import Column, Integer, String, LargeBinary, Enum
+from sqlalchemy import Column, Integer, String, LargeBinary, Enum, Boolean
 from sqlalchemy.ext.declarative import declarative_base
 
 #instance for create declarative base
@@ -17,7 +17,6 @@ class User(Base):
     first_name = Column(String)
     last_name = Column(String)
     email = Column(String, unique=True)
-    role = Column(Enum('admin', 'user', name='role_field'), nullable=False)
     image = Column(LargeBinary)
     active = Column(Enum('N', 'Y', name='active_field'), nullable=False, server_default=str("N"))
     password = Column(String, nullable=False)

dataflow_core-2.0.11/dataflow/dataflow.py

@@ -0,0 +1,162 @@
+import os, requests
+from .models.database import DatabaseManager
+from .utils.aws_secrets_manager import SecretsManagerClient
+import json
+from authenticator.package.configuration import ConfigurationManager
+
+
+class Dataflow:
+    def __init__(self):
+        self.secrets_manager = SecretsManagerClient()
+
+    def auth(self, session_id: str):
+        """Retrieve user information from the auth API."""
+        try:
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            auth_api = dataflow_config.get_config_value('auth', 'ui_auth_api')
+            response = requests.get(
+                auth_api,
+                cookies={"dataflow_session": session_id, "jupyterhub-hub-login": ""}
+            )
+            
+            if response.status_code != 200:
+                return response.json()
+            
+            user_data = response.json()
+            user_dict = {
+                "user_name": user_data["user_name"], 
+                "first_name": user_data["first_name"],
+                "last_name": user_data["last_name"] if user_data.get("last_name") else "",
+                "email": user_data["email"],
+                "role": user_data["role"]
+            }
+            return user_dict
+                  
+        except Exception as e:
+            return e
+    
+    def variable(self, variable_name: str):
+        """
+        Get variable value from UI API.
+        """
+        try:
+            host_name = os.environ.get("HOSTNAME", "")
+            user_name = host_name.replace("jupyter-", "") if host_name.startswith("jupyter-") else host_name
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            variable_api = dataflow_config.get_config_value("auth", "db_get_variables")
+            if not variable_api:
+                print("[Dataflow.variable] Variable Unreachable")
+                return None
+
+            if runtime:
+                query_params = {
+                    "variable_key": variable_name,
+                    "runtime": runtime,
+                    "slug": slug
+                }
+                response = requests.get(variable_api, params=query_params)
+                if response.status_code == 200:
+                    response_text = response.text.strip().strip('"')
+                    return response_text
+                
+                query_params["slug"] = "global"
+                response = requests.get(variable_api, params=query_params)
+                if response.status_code == 200:
+                    response_text = response.text.strip().strip('"')
+                    return response_text
+                else: 
+                    return None
+
+            query_params = {
+                "variable_key": variable_name,
+                "runtime": None,
+                "slug": None,
+                "created_by": user_name
+            }
+            response = requests.get(variable_api, params=query_params)
+            if response.status_code == 200:
+                response_text = response.text.strip().strip('"')
+                return response_text
+            else:
+                return None
+        except Exception as e:
+            print(f"[Dataflow.variable] Exception occurred: {e}")
+            return None
+        
+    def secret(self, secret_name: str):
+        """Get secret value from secrets manager for Studio or Runtime."""
+        try:
+            host_name = os.environ.get("HOSTNAME", "")
+            user_name = host_name.replace("jupyter-", "") if host_name.startswith("jupyter-") else host_name
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            secret_api = dataflow_config.get_config_value("auth", "db_get_secrets")
+            if not secret_api:
+                print("[Dataflow.secret] Secret API Unreachable")
+                return None
+
+            query_params = {
+                "secret_key": secret_name,
+                "created_by": user_name
+            }
+
+            if runtime:
+                query_params["runtime"] = runtime
+            if slug:
+                query_params["slug"] = slug
+
+            response = requests.get(secret_api, params=query_params)
+            
+            if response.status_code == 200:
+                response_text = response.text.strip().strip('"')
+                return response_text
+            else:
+                return None
+        except Exception as e:
+            print(f"[Dataflow.secret] Exception occurred: {e}")
+            return None
+
+    def connection(self, conn_id: str):
+        """Get connection details from secrets manager."""
+        try:
+            host_name = os.environ["HOSTNAME"]
+            user_name=host_name.replace("jupyter-","")
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+            
+            vault_path = "connections"
+            secret = self.secrets_manager.get_secret_by_key(vault_path, user_name, conn_id, runtime, slug)
+
+            conn_type = secret['conn_type'].lower()
+            username = secret['login']
+            password = secret.get('password', '')
+            host = secret['host']
+            port = secret['port']
+            database = secret.get('schemas', '')
+
+            user_info = f"{username}:{password}@" if password else f"{username}@"
+            db_info = f"/{database}" if database else ""
+
+            connection_string = f"{conn_type}://{user_info}{host}:{port}{db_info}"
+
+            extra = secret.get('extra', '')            
+            if extra:
+                try:
+                    extra_params = json.loads(extra)
+                    if extra_params:
+                        extra_query = "&".join(f"{key}={value}" for key, value in extra_params.items())
+                        connection_string += f"?{extra_query}"
+                except json.JSONDecodeError:
+                    # If 'extra' is not valid JSON, skip adding extra parameters
+                    pass
+
+            connection_instance = DatabaseManager(connection_string)
+            return next(connection_instance.get_session())
+        
+        except Exception as e:
+            return None

{dataflow_core-2.0.9 → dataflow_core-2.0.11}/dataflow/environment.py

@@ -1,5 +1,4 @@
 import os, shutil, subprocess, datetime
-from .models.database import DatabaseManager
 from .models.environment import JobLogs, Environment
 import json, asyncio, pkg_resources
 from sqlalchemy.orm import Session

{dataflow_core-2.0.9 → dataflow_core-2.0.11}/dataflow/utils/aws_secrets_manager.py

@@ -1,6 +1,7 @@
 import boto3
 from botocore.exceptions import ClientError, EndpointConnectionError, NoCredentialsError
 import json
+from .json_handler import JsonHandler
 
 class SecretsManagerClient:
     """
@@ -11,6 +12,7 @@ class SecretsManagerClient:
         json_handler: An instance of JsonHandler for handling JSON operations.
     """
     def __init__(self):
+        self.json_handler = JsonHandler()
         try:
             self.client = boto3.client('secretsmanager')
         except EndpointConnectionError as e:
@@ -21,14 +23,16 @@ class SecretsManagerClient:
             raise Exception(f"Failed to initialize SecretsManagerClient: No AWS credentials found. {e}")
 
 
-    def get_secret_by_key(self, vault_path, user_name, secret_key: str):
+    def get_secret_by_key(self, vault_path, user_name, conn_id: str, runtime, slug):
         """
-        Get information about a specific secret.
+        Get information about a specific secret using a dynamically constructed vault path.
 
         Args:
-            vault_path (str): The vault path.
+            vault_path (str): The base vault path.
             user_name (str): The user name.
-            secret_key (str): The key of the secret to retrieve.
+            conn_id (str): The key of the secret to retrieve.
+            runtime (str, optional): The runtime environment. Defaults to None.
+            slug (str, optional): The slug identifier. Defaults to None.
 
         Returns:
             str: Information about the secret in JSON format.
@@ -37,25 +41,17 @@ class SecretsManagerClient:
             Exception: If the operation fails.
         """
         try:
-            if not user_name:
-                raise Exception("user_name is required when secret_key is provided")
-
-            secret_name = f"{user_name}/{vault_path}/{secret_key}"
+            if runtime and slug:
+                secret_name = f"{runtime}/{slug}/{vault_path}/{conn_id}"
+            else:
+                secret_name = f"{user_name}/{vault_path}/{conn_id}"
+            
             response = self.client.get_secret_value(SecretId=secret_name)
-            secret_metadata = self.client.describe_secret(SecretId=secret_name)
             secret_data = json.loads(response.get('SecretString'))
 
-            if secret_data.get('is_active') == 'Y':
-                secret_info={
-                    "Name": secret_key,
-                    "Description": secret_metadata.get('Description')
-                }
-                secret_info.update(secret_data)
-                return secret_info
-            else:
-                raise Exception(f"Secret named '{secret_key}' is not active")
+            return secret_data
         except ClientError as e:
             if e.response['Error']['Code'] == 'ResourceNotFoundException':
-                raise Exception(f"Secret named '{secret_key}' not found")
+                raise Exception(f"Secret named '{conn_id}' not found with path '{secret_name}'")
             else:
-                raise Exception(f"Failed to get secret '{secret_key}': {e}")
+                raise Exception(f"Failed to get secret '{conn_id}': {e}")

dataflow_core-2.0.11/dataflow/utils/json_handler.py

@@ -0,0 +1,33 @@
+import json 
+
+class JsonHandler:
+    """
+    Helper class for handling JSON serialization and deserialization.
+    """
+
+    def __init__(self):
+        pass
+
+    def dict_to_json(self, data_dict):
+        """
+        Serialize a dictionary to JSON string.
+
+        Args:
+            data_dict (dict): The dictionary to serialize.
+
+        Returns:
+            str: The JSON string representation of the dictionary.
+        """
+        return json.dumps(data_dict)
+
+    def json_to_dict(self, json_string):
+        """
+        Deserialize a JSON string to dictionary.
+
+        Args:
+            json_string (str): The JSON string to deserialize.
+
+        Returns:
+            dict: The dictionary representation of the JSON string.
+        """
+        return json.loads(json_string)

{dataflow_core-2.0.9 → dataflow_core-2.0.11}/dataflow_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataflow-core
-Version: 2.0.9
+Version: 2.0.11
 Summary: Dataflow core package
 Author: Dataflow
 Author-email: 

{dataflow_core-2.0.9 → dataflow_core-2.0.11}/dataflow_core.egg-info/SOURCES.txt

@@ -21,6 +21,7 @@ dataflow/scripts/clone_environment.sh
 dataflow/scripts/create_environment.sh
 dataflow/utils/__init__.py
 dataflow/utils/aws_secrets_manager.py
+dataflow/utils/json_handler.py
 dataflow_core.egg-info/PKG-INFO
 dataflow_core.egg-info/SOURCES.txt
 dataflow_core.egg-info/dependency_links.txt

{dataflow_core-2.0.9 → dataflow_core-2.0.11}/setup.py

@@ -14,7 +14,7 @@ class PostInstall(install):
 
 setup(
     name="dataflow-core",
-    version="2.0.9",
+    version="2.0.11",
     packages=find_packages(include=["dataflow", "dataflow.*", "authenticator", "authenticator.*"]),
     include_package_data=True,
     package_data={

dataflow_core-2.0.9/dataflow/dataflow.py

@@ -1,89 +0,0 @@
-import os, requests, shutil, subprocess, datetime
-from .models.database import DatabaseManager
-from .models.environment import JobLogs
-from sqlalchemy.inspection import inspect
-from .utils.aws_secrets_manager import SecretsManagerClient
-import json, asyncio, pkg_resources
-from authenticator.package.configuration import ConfigurationManager
-
-
-class Dataflow:
-    def __init__(self):
-        self.secrets_manager = SecretsManagerClient()
-
-    def auth(self, session_id: str):
-        """Retrieve user information from the auth API."""
-        try:
-            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
-            auth_api = dataflow_config.get_config_value('auth', 'ui_auth_api')
-            response = requests.get(
-                auth_api,
-                cookies={"dataflow_session": session_id, "jupyterhub-hub-login": ""}
-            )
-            
-            if response.status_code != 200:
-                return response.json()
-            
-            user_data = response.json()
-            user_dict = {
-                "user_name": user_data["user_name"], 
-                "first_name": user_data["first_name"],
-                "last_name": user_data["last_name"] if user_data.get("last_name") else "",
-                "email": user_data["email"],
-                "role": user_data["role"]
-            }
-            return user_dict
-                  
-        except Exception as e:
-            return e
-    
-    def variable(self, variable_name: str):
-        """Get variable value from secrets manager."""
-        try:
-            host_name = os.environ["HOSTNAME"]
-            user_name = host_name.replace("jupyter-","")
-            
-            vault_path = "variables"
-            variable_data =  self.secrets_manager.get_secret_by_key(vault_path, user_name, variable_name)
-            return variable_data['value']
-            
-        except Exception as e:
-            return None
-        
-    def connection(self, conn_id: str):
-        """Get connection details from secrets manager."""
-        try:
-            host_name = os.environ["HOSTNAME"]
-            user_name=host_name.replace("jupyter-","")
-            
-            vault_path = "connections"
-            secret = self.secrets_manager.get_secret_by_key(vault_path, user_name, conn_id)
-
-            conn_type = secret['conn_type'].lower()
-            username = secret['login']
-            password = secret.get('password', '')
-            host = secret['host']
-            port = secret['port']
-            database = secret.get('schemas', '')
-
-            user_info = f"{username}:{password}@" if password else f"{username}@"
-            db_info = f"/{database}" if database else ""
-
-            connection_string = f"{conn_type}://{user_info}{host}:{port}{db_info}"
-
-            extra = secret.get('extra', '')
-            if extra:
-                try:
-                    extra_params = json.loads(extra)
-                    if extra_params:
-                        extra_query = "&".join(f"{key}={value}" for key, value in extra_params.items())
-                        connection_string += f"?{extra_query}"
-                except json.JSONDecodeError:
-                    # If 'extra' is not valid JSON, skip adding extra parameters
-                    pass
-
-            connection_instance = DatabaseManager(connection_string)
-            return next(connection_instance.get_session())
-        
-        except Exception as e:
-            return None