dataflow-core 2.0.10__tar.gz → 2.1.0__tar.gz

{dataflow_core-2.0.10 → dataflow_core-2.1.0}/PKG-INFO

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

dataflow_core-2.1.0/authenticator/dataflowairflowauthenticator.py

@@ -0,0 +1,58 @@
+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 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:
+            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_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())
+                )
+                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:
+            logger.error(f"Login failed: {e}")
+            return super().login()
+
+class DataflowAirflowAuthenticator(FabAirflowSecurityManagerOverride):
+    authdbview = DataflowAuthDBView
+    
+    def __init__(self, appbuilder):
+        super().__init__(appbuilder)

{dataflow_core-2.0.10 → dataflow_core-2.1.0}/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()

{dataflow_core-2.0.10 → dataflow_core-2.1.0}/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.10/dataflow/models/database.py → dataflow_core-2.1.0/dataflow/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_core-2.1.0/dataflow/dataflow.py

@@ -0,0 +1,199 @@
+import os, requests
+from .database_manager import DatabaseManager
+from .utils.aws_secrets_manager import SecretsManagerClient
+import json
+from .configuration import ConfigurationManager
+
+
+class Dataflow:
+    def __init__(self):
+        self.secrets_manager = SecretsManagerClient()
+
+    def auth(self, session_id: str):
+        """
+        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')
+            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["base_role"]
+            }
+            return user_dict
+                  
+        except Exception as e:
+            return e
+    
+    def variable(self, variable_name: str):
+        """
+        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.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
+            
+        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.
+        
+        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, 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
+
+            if mode == "url":
+                return connection_string
+    
+            connection_instance = DatabaseManager(connection_string)
+            if mode == "engine":
+                return connection_instance.get_engine()
+            elif mode == "session":
+                return next(connection_instance.get_session())
+        
+        except Exception as e:
+            return None

dataflow_core-2.1.0/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_core-2.0.10 → dataflow_core-2.1.0}/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_core-2.1.0/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_core-2.1.0/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_core-2.1.0/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_core-2.1.0/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'),
+    )

{dataflow_core-2.0.10 → dataflow_core-2.1.0}/dataflow/models/environment.py

@@ -1,9 +1,8 @@
 from sqlalchemy import Column, Integer, String, Boolean, Text, ForeignKey, DateTime
-from sqlalchemy.orm import relationship, declarative_base
+from sqlalchemy.orm import relationship
 from sqlalchemy.sql import func
 from datetime import datetime
-
-Base = declarative_base()
+from dataflow.db import Base
 
 class EnvironmentAttributes(Base):
     """
@@ -53,7 +52,7 @@ class ArchivedEnvironment(EnvironmentAttributes):
 
 
 class JobLogs(Base):
-    __tablename__ = "JOB_LOGS"
+    __tablename__ = "JOB_LOG"
 
     id = Column(Integer, primary_key=True, index=True)
     created_at = Column(DateTime, default=datetime.now)
@@ -61,4 +60,5 @@ class JobLogs(Base):
     log_file_name = Column(String, unique=True, nullable=False)
     log_file_location = Column(String, nullable=False)
     status = Column(String)
-    created_by = Column(String)
+    created_by = Column(String)
+

dataflow_core-2.1.0/dataflow/models/environment_status.py

@@ -0,0 +1,17 @@
+"""models.py"""
+from sqlalchemy import Column, Integer, String, ForeignKey, DateTime
+from sqlalchemy.sql import func
+from dataflow.db import Base
+
+class EnvironmentStatus(Base):
+    """
+    Table ENVIRONMENT_STATUS
+    """
+
+    __tablename__='ENVIRONMENT_STATUS'
+
+    id = Column(Integer, ForeignKey('ENVIRONMENT.id'), primary_key=True, nullable=False)
+    status = Column(String, nullable=False)
+    comment = Column(String)
+    status_changed_date = Column(DateTime, server_default=func.now(), nullable=False)
+    

dataflow_core-2.1.0/dataflow/models/git_ssh.py

@@ -0,0 +1,18 @@
+"""models.py"""
+from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, UniqueConstraint
+from sqlalchemy.sql import func
+from dataflow.db import Base
+
+class GitSSH(Base):
+    __tablename__ = 'GIT_SSH'
+
+    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
+    user_name = Column(String, ForeignKey('USER.user_name', ondelete="CASCADE"), nullable=False)
+    description = Column(String)
+    key_name = Column(String, nullable=False)
+    created_date = Column(DateTime, server_default=func.now(), nullable=False)
+    last_used_date = Column(DateTime)
+
+    __table_args__ = (
+        UniqueConstraint(user_name, key_name, name='user_name_key_name_unique'),
+    )

dataflow_core-2.1.0/dataflow/models/pinned_projects.py

@@ -0,0 +1,15 @@
+from sqlalchemy import Column, Integer, ForeignKey, DateTime, UniqueConstraint
+from dataflow.db import Base
+from datetime import datetime
+
+class PinnedProject(Base):
+    __tablename__ = "PINNED_PROJECT"
+
+    id = Column(Integer, primary_key=True)
+    user_id = Column(Integer, ForeignKey('USER.user_id'), index=True)
+    project_id = Column(Integer, ForeignKey('PROJECT_DETAIL.project_id'), index=True)
+    pinned_at = Column(DateTime, default=datetime.utcnow)
+
+    __table_args__ = (
+        UniqueConstraint("user_id", "project_id", name="uix_user_project"),
+    )