dataflow-core 2.0.10__py3-none-any.whl → 2.1.0__py3-none-any.whl

authenticator/dataflowairflowauthenticator.py

@@ -1,83 +1,58 @@
-from .package.models.database import DatabaseManager
-from dataflow.dataflow import Dataflow
-
-from typing import Any, Callable
+from flask import redirect, request
+from flask_appbuilder.security.views import AuthDBView
+from flask_appbuilder.security.views import expose
+from flask_login import login_user
 from airflow.www.security import FabAirflowSecurityManagerOverride
-from airflow.configuration import conf
-
-class DataflowAirflowAuthenticator(FabAirflowSecurityManagerOverride):
-    def __init__(self, wsgi_app: Callable) -> None:
-        self.wsgi_app = wsgi_app
-        self.dataflow = Dataflow()
-        self.airflow_database_url = conf.get("database", "sql_alchemy_conn")
-
-        self.airflow_db_instance = DatabaseManager(self.airflow_database_url)
-        self.airflow_db = next(self.airflow_db_instance.get_session())
-
-    def __call__(self, environ: dict, start_response: Callable) -> Any:
-
-        path = environ.get('PATH_INFO', '')
-        if not path == '/login/':
-            return self.wsgi_app(environ, start_response)
-
+from dataflow.dataflow import Dataflow
+import logging
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+dataflow = Dataflow()
+
+class DataflowAuthDBView(AuthDBView):    
+    @expose('/login/', methods=['GET', 'POST'])
+    def login(self):
+        """
+        Override the default login method to handle custom authentication
+        """
         try:
-            # Extracting browser cookies
-            cookies = environ.get('HTTP_COOKIE', '')
-            user_session_id = None
-            parts = cookies.split('; ')
-            for part in parts:
-                if part.startswith('dataflow_session='):
-                    user_session_id = part
-                    break
-
-            if user_session_id is None:
-                raise Exception("No session id found")
+            session_id = request.cookies.get('dataflow_session')
+            if not session_id:
+                logger.info("No session cookie found, falling back to standard login.")
+                return super().login()
             
-            user_session_id = user_session_id.split('=')[1]
-
-            # Retrieving user details
-            user_data = self.dataflow.auth(user_session_id)
-
-            if user_data is None:
-                raise Exception("No user found for the dataflow_session id")
-            
-            user = self.find_user(user_data["user_name"])
-
-            if not user:
-                user_role = self.find_role(user_data["role"].title())
-                user = self.add_user(
-                    username=user_data["user_name"],
-                    first_name=user_data.get("first_name", ""),
-                    last_name=user_data.get("last_name", ""),
-                    email=user_data.get("email", ""),
-                    role=user_role
+            user_details = dataflow.auth(session_id)
+            logger.info(f"User details retrieved for: {user_details['user_name']}")
+            user = self.appbuilder.sm.find_user(username=user_details['user_name'])
+            if user:
+                logger.info(f"User found: {user}")
+                login_user(user, remember=False)
+            else:
+                user = self.appbuilder.sm.add_user(
+                    username=user_details['user_name'], 
+                    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(user_details.get("base_role", "user").title())
                 )
-
-            environ['REMOTE_USER'] = user.username
-            return self.wsgi_app(environ, start_response)
+                logger.info(f"New user created: {user}")
+                if user:
+                    login_user(user, remember=False)
+            
+            return redirect(self.appbuilder.get_url_for_index)
 
         except Exception as e:
-            return self.wsgi_app(environ, start_response)
-
-    def find_user(self, username=None):
-        """Find user by username or email."""
-        return self.airflow_db.query(self.user_model).filter_by(username=username).one_or_none()
+            logger.error(f"Login failed: {e}")
+            return super().login()
 
-    def find_role(self, role):
-        """Find a role in the database."""
-        return self.airflow_db.query(self.role_model).filter_by(name=role).one_or_none()
-
-    def add_user(self, username, first_name, last_name, email, role, password=""):
-        """Create a user."""
-        user = self.user_model()
-        user.first_name = first_name
-        user.last_name = last_name
-        user.username = username
-        user.email = email
-        user.active = True
-        user.roles = role if isinstance(role, list) else [role]
-        user.password = password
-        self.airflow_db.add(user)
-        self.airflow_db.commit()
-        return user
-        
+class DataflowAirflowAuthenticator(FabAirflowSecurityManagerOverride):
+    authdbview = DataflowAuthDBView
+    
+    def __init__(self, appbuilder):
+        super().__init__(appbuilder)

authenticator/dataflowhubauthenticator.py

@@ -1,6 +1,5 @@
-from .package.configuration import ConfigurationManager
-from .package.models.database import DatabaseManager
-from .package.models import (
+from dataflow.db import get_db
+from dataflow.models import (
     user as m_user,
     session as m_session
     )
@@ -11,15 +10,7 @@ from jupyterhub.auth import Authenticator
 class DataflowHubAuthenticator(Authenticator):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-
-        self.dataflow_config = ConfigurationManager('/dataflow/app/config/dataflow.cfg')
-        self.database_url = self.dataflow_config.get_config_value('database', 'database_url')
-
-        self.db_instance = DatabaseManager(self.database_url)
-        self.db = next(self.db_instance.get_session())
-
-        m_user.Base.metadata.create_all(bind=self.db_instance.get_engine())
-        m_session.Base.metadata.create_all(bind=self.db_instance.get_engine())
+        self.db = next(get_db())
     
     def generate_session_id(self):
         return str(uuid.uuid4())
@@ -39,8 +30,8 @@ class DataflowHubAuthenticator(Authenticator):
 
             # Check if the user already has an existing session
             existing_session = (
-                self.db.query(m_session.Session_table)
-                .filter(m_session.Session_table.user_id == str(user.user_id))
+                self.db.query(m_session.Session)
+                .filter(m_session.Session.user_id == user.user_id)
                 .first()
             )
 
@@ -50,16 +41,16 @@ class DataflowHubAuthenticator(Authenticator):
             else:
                 # Generate a new session_id
                 session_id = self.generate_session_id()
-                query = self.db.query(m_session.Session_table)
-                isSession = query.filter(m_session.Session_table.session_id == session_id).first()
+                query = self.db.query(m_session.Session)
+                isSession = query.filter(m_session.Session.session_id == session_id).first()
 
                 # If session_id(uuid string) already exists in the database, generate a new one
                 while isSession is not None:
                     session_id = self.generate_session_id()
-                    isSession = query.filter(m_session.Session_table.session_id == session_id).first()
+                    isSession = query.filter(m_session.Session.session_id == session_id).first()
 
                 # add session_id to the database
-                db_item = m_session.Session_table(user_id=user.user_id, session_id=session_id)
+                db_item = m_session.Session(user_id=user.user_id, session_id=session_id)
                 self.db.add(db_item)
                 self.db.commit()
                 self.db.refresh(db_item)
@@ -86,4 +77,7 @@ class DataflowHubAuthenticator(Authenticator):
             return user_dict
 
         except Exception as e:
-            return None
+            return None
+        
+        finally:
+            self.db.close()

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/{models/database.py → database_manager.py}

@@ -7,6 +7,7 @@ from sqlalchemy.orm import sessionmaker
 class DatabaseManager:
     def __init__(self, db_url):
         self.db_url = db_url
+        self.engine = self.get_engine()
 
     def get_engine(self):
         try:
@@ -17,8 +18,8 @@ class DatabaseManager:
         
     def get_session(self):
         try:
-            engine = self.get_engine()
-            session = sessionmaker(autocommit=False, autoflush=False, bind=engine or create_engine(self.db_url))
+            engine = self.engine
+            session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
             db = session()
             try:
                 yield db
@@ -29,4 +30,4 @@ class DatabaseManager:
             raise e
     
     def get_base(self):
-        return declarative_base()
+        return declarative_base()

dataflow/dataflow.py

@@ -1,10 +1,8 @@
-import os, requests, shutil, subprocess, datetime
-from .models.database import DatabaseManager
-from .models.environment import JobLogs
-from sqlalchemy.inspection import inspect
+import os, requests
+from .database_manager import DatabaseManager
 from .utils.aws_secrets_manager import SecretsManagerClient
-import json, asyncio, pkg_resources
-from authenticator.package.configuration import ConfigurationManager
+import json
+from .configuration import ConfigurationManager
 
 
 class Dataflow:
@@ -12,7 +10,15 @@ class Dataflow:
         self.secrets_manager = SecretsManagerClient()
 
     def auth(self, session_id: str):
-        """Retrieve user information from the auth API."""
+        """
+        Retrieve and return user information using their session ID.
+        
+        Args:
+            session_id (str): User's session ID from cookies
+            
+        Returns:
+            dict: User information including username, name, email, and role
+        """
         try:
             dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
             auth_api = dataflow_config.get_config_value('auth', 'ui_auth_api')
@@ -30,7 +36,7 @@ class Dataflow:
                 "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"]
+                "role": user_data["base_role"]
             }
             return user_dict
                   
@@ -38,26 +44,124 @@ class Dataflow:
             return e
     
     def variable(self, variable_name: str):
-        """Get variable value from secrets manager."""
+        """
+        Retrieve a Dataflow variable.
+        
+        Args:
+            variable_name (str): Name of the variable to retrieve
+            
+        Returns:
+            str or None: Variable value if found, None otherwise
+        """
         try:
-            host_name = os.environ["HOSTNAME"]
-            user_name = host_name.replace("jupyter-","")
+            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):
+        """
+        Retrieve a Dataflow secret value.
+        
+        Args:
+            secret_name (str): Name of the secret to retrieve
             
-            vault_path = "variables"
-            variable_data =  self.secrets_manager.get_secret_by_key(vault_path, user_name, variable_name)
-            return variable_data['value']
+        Returns:
+            str or None: Secret value if found, None otherwise
+        """
+        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, mode="session"):
+        """
+        Connects with a Dataflow connection.
         
-    def connection(self, conn_id: str):
-        """Get connection details from secrets manager."""
+        Args:
+            conn_id (str): Connection identifier
+            mode (str): Return type - "session" (default) or "engine" or "url"
+            
+        Returns:
+            Session or Engine: SQLAlchemy session or engine based on mode
+        """
         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)
+            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']
@@ -71,7 +175,7 @@ class Dataflow:
 
             connection_string = f"{conn_type}://{user_info}{host}:{port}{db_info}"
 
-            extra = secret.get('extra', '')
+            extra = secret.get('extra', '')            
             if extra:
                 try:
                     extra_params = json.loads(extra)
@@ -82,8 +186,14 @@ class Dataflow:
                     # If 'extra' is not valid JSON, skip adding extra parameters
                     pass
 
+            if mode == "url":
+                return connection_string
+    
             connection_instance = DatabaseManager(connection_string)
-            return next(connection_instance.get_session())
+            if mode == "engine":
+                return connection_instance.get_engine()
+            elif mode == "session":
+                return next(connection_instance.get_session())
         
         except Exception as e:
             return None

dataflow/db.py

@@ -0,0 +1,51 @@
+from .database_manager import DatabaseManager
+from .configuration import ConfigurationManager
+from sqlalchemy.orm import declarative_base
+from sqlalchemy.exc import SQLAlchemyError
+from .utils.logger import CustomLogger
+
+logger = CustomLogger().get_logger(__name__)
+
+dataflow_config = ConfigurationManager('/dataflow/app/config/dataflow.cfg')
+db_url = dataflow_config.get_config_value('database', 'database_url')
+local_db_url = dataflow_config.get_config_value('database', 'single_user_db_url')
+
+db_manager = DatabaseManager(db_url)
+local_db_manager = None
+
+Base = declarative_base()
+Local_Base = declarative_base()
+
+def create_tables(local_db=False):
+    """
+    Create all tables in the database.
+    This is called at the start of the application.
+    """
+    try:
+        if local_db:
+            global local_db_manager
+            if local_db_manager is None:
+                local_db_manager = DatabaseManager(local_db_url) 
+            Local_Base.metadata.create_all(bind=local_db_manager.get_engine())
+        else:
+            Base.metadata.create_all(bind=db_manager.get_engine())
+        logger.info("Database tables created successfully")
+    except SQLAlchemyError as e:
+        error_message = f"Failed to create tables: {str(e)}"
+        logger.error(error_message)
+        raise e
+    
+def get_local_db():
+    """
+    Get a local database session.
+    
+    Yields:
+        Session: Local database session
+    """
+    global local_db_manager
+    if local_db_manager is None:
+        local_db_manager = DatabaseManager(local_db_url)
+    yield from local_db_manager.get_session()
+
+def get_db():
+    yield from db_manager.get_session()

dataflow/environment.py

@@ -1,9 +1,8 @@
 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
-from authenticator.package.configuration import ConfigurationManager
+from .configuration import ConfigurationManager
 
 class EnvironmentManager:
     def __init__(self):

dataflow/models/__init__.py

@@ -0,0 +1,22 @@
+# init for loading models in the application
+
+from .role import Role
+from .user import User
+from .team import Team
+from .environment import Environment
+from .project_details import ProjectDetails
+from .recent_projects import RecentProjects
+from .pinned_projects import PinnedProject
+from .app_types import AppType
+from .blacklist_library import BlacklistedLibrary
+from .environment_status import EnvironmentStatus
+from .session import Session
+from .server_config import ServerConfig
+from .runtime import RuntimeZone
+from .environment_status import EnvironmentStatus
+from .user_team import UserTeam
+from .role_server import RoleServer
+from .variables import Variable
+from .recent_project_studio import RecentProjectStudio
+from .connection import Connection
+from .git_ssh import GitSSH

dataflow/models/app_types.py

@@ -0,0 +1,10 @@
+from sqlalchemy import Column, Integer, String, Boolean
+from dataflow.db import Base
+
+class AppType(Base):
+    __tablename__ = "RUNTIME_APP_TYPE"
+    
+    id = Column(Integer, primary_key=True, autoincrement=True, unique=True)
+    name = Column(String, unique=True, nullable=True)
+    display_name = Column(String, nullable=False)
+    code_based = Column(Boolean, nullable=False)

dataflow/models/blacklist_library.py

@@ -0,0 +1,26 @@
+# models/blacklist_library.py
+from sqlalchemy import Column, Integer, String, UniqueConstraint
+from dataflow.db import Base
+
+class BlacklistedLibrary(Base):
+    """
+    BlacklistedLibrary model represents a table for storing blacklisted libraries with their versions.
+
+    Attributes:
+        id (int): Primary key of the table, auto-incremented.
+        library_name (str): The name of the blacklisted library.
+        version (str): The version of the blacklisted library.
+
+    Unique constraint to ensure the combination of library_name and version is unique.
+    """
+
+    __tablename__ = "BLACKLISTED_LIBRARY"
+    
+    id = Column(Integer, primary_key=True, index=True, doc="Primary key for the library.")
+    library_name = Column(String, index=True, doc="The name of the blacklisted library.")
+    version = Column(String, doc="The version of the blacklisted library.")
+
+    __table_args__ = (
+        UniqueConstraint('library_name', 'version', name='uq_library_version'),
+    )
+    

dataflow/models/connection.py

@@ -0,0 +1,25 @@
+from sqlalchemy import Column, String, Integer, Boolean, DateTime, UniqueConstraint
+from sqlalchemy.sql import func
+from dataflow.db import Base
+
+class Connection(Base):
+    """
+    Database model for storing non-sensitive connection metadata
+    """
+    __tablename__ = "CONNECTION"
+
+    id = Column(Integer, primary_key=True, index=True)
+    conn_id = Column(String, index=True, nullable=False)
+    description = Column(String, nullable=True)
+    conn_type = Column(String, nullable=False)
+    runtime = Column(String, nullable=True)
+    slug = Column(String, nullable=True)
+    status = Column(Boolean, default=False)
+    created_by = Column(String, nullable=True)
+    created_at = Column(DateTime(timezone=True), server_default=func.now())
+    updated_at = Column(DateTime(timezone=True), onupdate=func.now())
+    is_active = Column(Boolean, default=True)
+    
+    __table_args__ = (
+        UniqueConstraint('conn_id', 'runtime', 'slug', 'is_active', 'created_by', name='uq_active_conn_with_runtime_slug'),
+    )