dataflow-core 2.1.8__tar.gz → 2.1.18rc3__tar.gz

{dataflow_core-2.1.8 → dataflow_core-2.1.18rc3}/PKG-INFO

@@ -1,16 +1,19 @@
 Metadata-Version: 2.4
 Name: dataflow-core
-Version: 2.1.8
+Version: 2.1.18rc3
 Summary: Dataflow core package
 Author: Dataflow
 Author-email: 
 Requires-Dist: sqlalchemy
+Requires-Dist: alembic
 Requires-Dist: boto3
 Requires-Dist: psycopg2-binary
 Requires-Dist: pymysql
 Requires-Dist: requests
 Requires-Dist: azure-identity
 Requires-Dist: azure-keyvault-secrets
+Requires-Dist: google-auth
+Requires-Dist: google-cloud-secret-manager
 Dynamic: author
 Dynamic: requires-dist
 Dynamic: summary

{dataflow_core-2.1.8 → dataflow_core-2.1.18rc3}/authenticator/dataflowairflowauthenticator.py

@@ -18,9 +18,14 @@ dataflow = Dataflow()
 class DataflowAuthDBView(AuthDBView):    
     @expose('/login/', methods=['GET', 'POST'])
     def login(self):
+
+        """This method checks for a 'dataflow_session' cookie, retrieves user details from Dataflow,
+        and logs in or creates the user in Airflow accordingly. 
+        If the cookie is not present, it falls back to the standard login process.
+        
+        Overrides the default login method to integrate with Dataflow authentication.
         """
-        Override the default login method to handle custom authentication
-        """
+
         try:
             session_id = request.cookies.get('dataflow_session')
             if not session_id:
@@ -52,6 +57,9 @@ class DataflowAuthDBView(AuthDBView):
             return super().login()
 
 class DataflowAirflowAuthenticator(FabAirflowSecurityManagerOverride):
+    
+    """Custom Security Manager to integrate Airflow authentication with Dataflow."""
+    
     authdbview = DataflowAuthDBView
     
     def __init__(self, appbuilder):

{dataflow_core-2.1.8 → dataflow_core-2.1.18rc3}/authenticator/dataflowhubauthenticator.py

@@ -1,6 +1,4 @@
-import os
-import uuid
-import re
+import os, uuid, re, hashlib, secrets
 from datetime import datetime, timedelta
 from zoneinfo import ZoneInfo
 from traitlets import Bool, Unicode
@@ -8,9 +6,17 @@ from jupyterhub.auth import Authenticator
 from oauthenticator.google import GoogleOAuthenticator
 from oauthenticator.azuread import AzureAdOAuthenticator
 from dataflow.db import get_db
-from dataflow.models import user as m_user, session as m_session, role as m_role
+from dataflow.models import user as m_user, session as m_session
+from sqlalchemy import or_
 
 class DataflowBaseAuthenticator(Authenticator):
+    
+    """Base Authenticator to handle Dataflow authentication and session management.
+    Provides methods to authenticate users via Dataflow credentials, manage sessions.
+    
+    Overrides JupyterHub's Authenticator class.
+    """
+    
     enable_dataflow_auth = Bool(True, config=True, help="Enable username/password authentication")
 
     def __init__(self, **kwargs):
@@ -25,10 +31,20 @@ class DataflowBaseAuthenticator(Authenticator):
             raise
 
     def generate_session_id(self):
+
+        """Generate and return a unique session ID using UUID4."""
+
         return str(uuid.uuid4())
 
     def set_session_cookie(self, handler, session_id):
-        expires = datetime.now(ZoneInfo("UTC")) + timedelta(days=365)
+
+        """Set the dataflow_session cookie in the user's browser.
+        
+        Args:
+            handler: The request handler to set the cookie on.
+            session_id: The session ID to set in the cookie."""
+
+        expires = datetime.now(ZoneInfo("UTC")) + timedelta(days=60)
         host = handler.request.host
         domain = '.'.join(host.split('.')[-2:]) if len(host.split('.')) >= 2 else host
         handler.set_cookie(
@@ -44,19 +60,22 @@ class DataflowBaseAuthenticator(Authenticator):
         self.log.info(f"Set session cookie: dataflow_session={session_id} for host={host}")
 
     def get_or_create_session(self, user_id):
-        existing_session = (
-            self.db.query(m_session.Session)
-            .filter(m_session.Session.user_id == str(user_id))
-            .first()
-        )
-        if existing_session:
-            self.log.info(f"Reusing existing session: {existing_session.session_id}")
-            return existing_session.session_id
+
+        """Retrieve existing session ID for user or create a new one.
+        
+        Args:
+            user_id: The ID of the user to get or create a session for.
+        
+        Returns:
+            session_id (str): The existing or newly created session ID.
+        """
+
         session_id = self.generate_session_id()
         while self.db.query(m_session.Session).filter(
             m_session.Session.session_id == session_id
         ).first():
             session_id = self.generate_session_id()
+
         db_item = m_session.Session(user_id=user_id, session_id=session_id)
         self.db.add(db_item)
         self.db.commit()
@@ -65,6 +84,16 @@ class DataflowBaseAuthenticator(Authenticator):
         return session_id
     
     def check_blocked_users(self, username, authenticated):
+        
+        """Check if the authenticated user is blocked based on allowed_users list.
+        
+        Args:
+            username (str): The username of the authenticated user.
+            authenticated (dict|None): The authentication data returned from authenticate method.
+        
+        Returns:
+            username (str|None): The username if not blocked, else None."""
+        
         self.log.info(f"Checking blocked users for {username}: authenticated={authenticated}, allowed_users={self.allowed_users}")
 
         if not authenticated:
@@ -77,42 +106,54 @@ class DataflowBaseAuthenticator(Authenticator):
 
         return super().check_blocked_users(username, authenticated)
 
-    def get_applicant_role_id(self):
-        """Get the role ID for 'Applicant' role"""
-        try:
-            applicant_role = (
-                self.db.query(m_role.Role)
-                .filter(m_role.Role.name == "Applicant")
-                .first()
-            )
-            if applicant_role:
-                return applicant_role.id
-            else:
-                self.log.warning("Applicant role not found in database")
-                return None
-        except Exception as e:
-            self.log.error(f"Error getting Applicant role: {str(e)}")
-            return None
-
     def extract_username_from_email(self, email):
-        """Extract username from email by removing domain"""
+
+        """Extract username from email by removing domain
+        
+        Args:
+            email (str): User's email address
+        
+        Returns:
+            username (str): Extracted username after removing domain
+        """
+        
         if '@' in email:
             return email.split('@')[0]
         return email
 
+    def generate_secure_password(self):
+        
+        """Generate secure random password hash
+        
+        Returns:
+            password_hash (str): Securely hashed password
+        """
+        
+        salt = secrets.token_hex(16)
+        random_uuid = str(uuid.uuid4())
+        hash_obj = hashlib.sha256((random_uuid + salt).encode())
+        return hash_obj.hexdigest()
+
     def create_new_user(self, email, first_name=None, last_name=None):
-        """Create a new user with Applicant role"""
+        
+        """Create a new user with Applicant role
+        
+        Args:
+            email (str): User's email address
+            first_name (str): User's first name
+            last_name (str): User's last name
+        
+        Returns:
+            new_user (m_user.User|None): Created user object or None if creation failed
+        """
+        
         try:
-            role_id = self.get_applicant_role_id()
-            if not role_id:
-                self.log.error("Cannot create user: Applicant role not found")
-                return None
-
             username = self.extract_username_from_email(email)
             username = re.sub(r'[^a-z0-9]', '', username.lower())
             if not username:
                 self.log.error("Cannot create user: Username is empty")
                 return None
+            
             existing_user = (
                 self.db.query(m_user.User)
                 .filter(m_user.User.user_name == username)
@@ -122,7 +163,7 @@ class DataflowBaseAuthenticator(Authenticator):
                 counter = 1
                 original_username = username
                 while existing_user:
-                    username = f"{original_username}_{counter}"
+                    username = f"{original_username}{counter}"
                     existing_user = (
                         self.db.query(m_user.User)
                         .filter(m_user.User.user_name == username)
@@ -130,13 +171,13 @@ class DataflowBaseAuthenticator(Authenticator):
                     )
                     counter += 1
 
+            secure_password = self.generate_secure_password()
             new_user = m_user.User(
                 user_name=username,
                 first_name=first_name or username,
                 last_name=last_name or "",
                 email=email,
-                role_id=role_id,
-                password='user@123',
+                password=secure_password,
             )
             
             self.db.add(new_user)
@@ -152,24 +193,42 @@ class DataflowBaseAuthenticator(Authenticator):
             return None
 
     async def authenticate_dataflow(self, handler, data):
+
+        """Authenticate user using Dataflow username/password.
+
+        Args:
+            handler: The request handler.
+            data: The authentication data containing username and password.
+        
+        Returns:
+            dict|None: Authentication result with username and session_id if successful, else None.
+        """
+
         if not (self.enable_dataflow_auth and isinstance(data, dict) and data.get("username") and data.get("password")):
             return None
-        username = data["username"]
+        user_name_or_email = data["username"]
         password = data["password"]
-        self.log.info(f"Attempting Dataflow authentication for user: {username}")
+        self.log.info(f"Attempting Dataflow authentication for user: {user_name_or_email}")
         try:
             user = (
                 self.db.query(m_user.User)
-                .filter(m_user.User.user_name == username)
+                .filter(
+                    or_(
+                        m_user.User.email == user_name_or_email,
+                        m_user.User.user_name == user_name_or_email
+                    )
+                )
                 .first()
             )
+
             if not user or user.password != password:
-                self.log.warning(f"Dataflow authentication failed for user: {username}")
+                self.log.warning(f"Dataflow authentication failed for user: {user_name_or_email}")
                 return None
+            
             session_id = self.get_or_create_session(user.user_id)
             self.set_session_cookie(handler, session_id)
-            self.log.info(f"Dataflow authentication successful for user: {username}")
-            return {"name": username, "session_id": session_id, "auth_state": {}}
+            self.log.info(f"Dataflow authentication successful for user: {user.user_name}")
+            return {"name": user.user_name, "session_id": session_id, "auth_state": {}}
         except Exception as e:
             self.log.error(f"Dataflow authentication error: {str(e)}")
             return None
@@ -177,6 +236,18 @@ class DataflowBaseAuthenticator(Authenticator):
             self.db.close()
 
 class DataflowGoogleAuthenticator(DataflowBaseAuthenticator, GoogleOAuthenticator):
+
+    """Authenticator to handle Google OAuth authentication with Dataflow integration.
+    
+    Overrides
+      - DataflowBaseAuthenticator
+      - GoogleOAuthenticator
+    
+    Requires Google OAuth credentials.
+      - google_client_id
+      - google_client_secret
+    """
+
     dataflow_oauth_type = Unicode(
         default_value="google",
         config=True,
@@ -195,6 +266,17 @@ class DataflowGoogleAuthenticator(DataflowBaseAuthenticator, GoogleOAuthenticato
                       f"enable_dataflow_auth={self.enable_dataflow_auth}")
 
     async def authenticate(self, handler, data):
+
+        """Authenticate user using Google OAuth with Dataflow integration.
+
+        Args:
+            handler: The request handler.
+            data: The authentication data.
+        
+        Returns:
+            dict|None: Authentication result with username and session_id if successful, else None.
+        """
+
         self.log.info(f"Authenticate called with data: {data}, request_uri: {handler.request.uri}")
         result = await self.authenticate_dataflow(handler, data)
         if result:
@@ -243,6 +325,19 @@ class DataflowGoogleAuthenticator(DataflowBaseAuthenticator, GoogleOAuthenticato
             self.db.close()
 
 class DataflowAzureAuthenticator(DataflowBaseAuthenticator, AzureAdOAuthenticator):
+
+    """Authenticator to handle Azure AD OAuth authentication with Dataflow integration.
+
+    Overrides
+      - DataflowBaseAuthenticator
+      - AzureAdOAuthenticator
+
+    Requires Azure AD OAuth credentials.
+      - azure_client_id
+      - azure_client_secret
+      - azure_tenant_id
+    """
+
     azure_client_id = Unicode(config=True, help="Azure AD OAuth client ID")
     azure_client_secret = Unicode(config=True, help="Azure AD OAuth client secret")
     azure_tenant_id = Unicode(config=True, help="Azure AD tenant ID")
@@ -264,6 +359,17 @@ class DataflowAzureAuthenticator(DataflowBaseAuthenticator, AzureAdOAuthenticato
                       f"enable_dataflow_auth={self.enable_dataflow_auth}")
 
     async def authenticate(self, handler, data):
+
+        """Authenticate user using Azure AD OAuth with Dataflow integration.
+
+        Args:
+            handler: The request handler.
+            data: The authentication data.
+        
+        Returns:
+            dict|None: Authentication result with username and session_id if successful, else None.
+        """
+        
         result = await self.authenticate_dataflow(handler, data)
         if result:
             return result
@@ -302,7 +408,7 @@ class DataflowAzureAuthenticator(DataflowBaseAuthenticator, AzureAdOAuthenticato
             self.set_session_cookie(handler, session_id)
             self.log.info(f"Azure AD OAuth completed for user: {username}, session_id={session_id}")
             return {
-                "name": db_user.first_name,
+                "name": username,
                 "session_id": session_id,
                 "auth_state": user.get("auth_state", {})
             }

{dataflow_core-2.1.8 → dataflow_core-2.1.18rc3}/authenticator/dataflowsupersetauthenticator.py

@@ -12,6 +12,7 @@ from superset.security import SupersetSecurityManager
 from dataflow.dataflow import Dataflow
 
 class DataflowAuthDBView(AuthDBView):
+
     def __init__(self):
         self.dataflow = Dataflow()
 
@@ -33,14 +34,15 @@ class DataflowAuthDBView(AuthDBView):
     @expose('/login/', methods=['GET', "POST"])
     def login(self):
         """
-        Handles both GET and POST login requests.
+        This method handles authentication for superset in Dataflow.
 
-        - GET: Used for browser-based login. Authenticates using session cookie and redirects to home.
-        - POST: Used for API-based login. Returns JWT access token for programmatic access.
+        Methods:
+        - GET: 
+            Used for browser-based login. Authenticates using session cookie and redirects to home.
 
-        Returns:
-            - Redirect to home page (GET)
-            - JSON response with access token (POST)
+        - POST: 
+            Used for API-based login. Returns JWT access token for programmatic access.
+            Returns JSON response with access token
         """
         if request.method == "GET":
             session_id = request.cookies.get('dataflow_session')
@@ -71,6 +73,9 @@ class DataflowAuthDBView(AuthDBView):
             return jsonify(resp)
     
 class DataflowSecurityManager(SupersetSecurityManager):
+
+    """Custom Security Manager integrating Dataflow authentication with superset."""
+    
     authdbview = DataflowAuthDBView
     def __init__(self, appbuilder):
         super(DataflowSecurityManager, self).__init__(appbuilder)

{dataflow_core-2.1.8 → dataflow_core-2.1.18rc3}/dataflow/configuration.py

@@ -20,6 +20,13 @@ class ConfigurationManager:
     def get_config_value(self, section, option):
         """
         Get configuration value
+
+        Args:
+            section (str): The section in the config file.
+            option (str): The option within the section.
+        
+        Returns:
+            str | None: The configuration value or None if not found.
         """
         try:
             return self.config.get(section, option)

{dataflow_core-2.1.8 → dataflow_core-2.1.18rc3}/dataflow/database_manager.py

@@ -5,11 +5,22 @@ from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
 class DatabaseManager:
+
+    """Manages database connections and sessions."""
+
     def __init__(self, db_url):
         self.db_url = db_url
         self.engine = self.get_engine()
 
     def get_engine(self):
+        
+        """
+        Create a new SQLAlchemy engine instance.
+
+        Returns:
+            Engine: The SQLAlchemy engine instance.
+        """
+        
         try:
             engine = create_engine(self.db_url)
             return engine
@@ -17,6 +28,12 @@ class DatabaseManager:
             raise e
         
     def get_session(self):
+        """
+        Create a new SQLAlchemy session.
+
+        Returns:
+            Session: The SQLAlchemy session instance.
+        """
         try:
             engine = self.engine
             session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
@@ -30,4 +47,10 @@ class DatabaseManager:
             raise e
     
     def get_base(self):
+        """
+        Get the declarative base class for the ORM models.
+
+        Returns:
+            Base: The declarative base class.
+        """
         return declarative_base()