dataflow-core 2.1.6__tar.gz → 2.1.18rc2__tar.gz

{dataflow_core-2.1.6 → dataflow_core-2.1.18rc2}/PKG-INFO

@@ -1,14 +1,19 @@
 Metadata-Version: 2.4
 Name: dataflow-core
-Version: 2.1.6
+Version: 2.1.18rc2
 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.6 → dataflow_core-2.1.18rc2}/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,7 +6,8 @@ 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):
     enable_dataflow_auth = Bool(True, config=True, help="Enable username/password authentication")
@@ -28,7 +27,7 @@ class DataflowBaseAuthenticator(Authenticator):
         return str(uuid.uuid4())
 
     def set_session_cookie(self, handler, session_id):
-        expires = datetime.now(ZoneInfo("UTC")) + timedelta(days=365)
+        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 +43,12 @@ 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
         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()
@@ -77,42 +69,28 @@ 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"""
         if '@' in email:
             return email.split('@')[0]
         return email
 
+    def generate_secure_password(self):
+        """Generate a secure random password hash"""
+        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"""
         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-Za-z0-9]', '', username)
+            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 +100,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,14 +108,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,
-                active='Y',
-                password='user@123',
+                password=secure_password,
             )
             
             self.db.add(new_user)
@@ -155,22 +132,29 @@ class DataflowBaseAuthenticator(Authenticator):
     async def authenticate_dataflow(self, handler, data):
         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
@@ -248,7 +232,11 @@ class DataflowAzureAuthenticator(DataflowBaseAuthenticator, AzureAdOAuthenticato
     azure_client_secret = Unicode(config=True, help="Azure AD OAuth client secret")
     azure_tenant_id = Unicode(config=True, help="Azure AD tenant ID")
     azure_scope = Unicode("openid profile email", config=True, help="Azure AD OAuth scopes")
-
+    dataflow_oauth_type = Unicode(
+        default_value="google",
+        config=True,
+        help="The OAuth provider type for DataflowHub (e.g., github, google)"
+    )
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.client_id = self.azure_client_id
@@ -270,32 +258,30 @@ class DataflowAzureAuthenticator(DataflowBaseAuthenticator, AzureAdOAuthenticato
             if not user:
                 self.log.warning("Azure AD OAuth authentication failed: No user data returned")
                 return None
-            
-            email = user.get("email") or user.get("preferred_username")
+
+            auth_state = user.get("auth_state", {})
+            user_info = auth_state.get("user", {}) if auth_state else {}
+            email = user_info.get("upn")
             if not email:
-                self.log.warning("Azure AD OAuth authentication failed: No email in user data")
+                self.log.warning("Azure AD OAuth authentication failed: No upn in user data")
                 return None
-            
+
             db_user = (
                 self.db.query(m_user.User)
                 .filter(m_user.User.email == email)
                 .first()
             )
-            
+
             if not db_user:
                 self.log.info(f"User with email {email} not found in Dataflow database, creating new user")
-                # Extract additional info from user data if available
-                auth_state = user.get("auth_state", {})
-                user_info = auth_state.get("user", {}) if auth_state else {}
                 
-                first_name = user_info.get("given_name") or user.get("given_name")
-                last_name = user_info.get("family_name") or user.get("family_name")
+                first_name = user_info.get("name") or user.get("name")
                 
-                db_user = self.create_new_user(email, first_name, last_name)
+                db_user = self.create_new_user(email, first_name, last_name=None)
                 if not db_user:
                     self.log.error(f"Failed to create new user for email: {email}")
                     return None
-            
+
             username = db_user.user_name
             session_id = self.get_or_create_session(db_user.user_id)
             self.set_session_cookie(handler, session_id)
@@ -305,13 +291,14 @@ class DataflowAzureAuthenticator(DataflowBaseAuthenticator, AzureAdOAuthenticato
                 "session_id": session_id,
                 "auth_state": user.get("auth_state", {})
             }
+
         except Exception as e:
             self.log.error(f"Azure AD OAuth authentication error: {str(e)}", exc_info=True)
             return None
         finally:
             self.db.close()
 
-auth_type = os.environ.get("DATAFLOW_AUTH_TYPE", "google")
+auth_type = os.environ.get("DATAFLOW_OAUTH_TYPE", "google")
 
 if auth_type == "google":
     BaseAuthenticator = DataflowGoogleAuthenticator

dataflow_core-2.1.18rc2/dataflow/dataflow.py

@@ -0,0 +1,423 @@
+import os, requests
+from .database_manager import DatabaseManager
+import json
+import base64
+from .configuration import ConfigurationManager
+
+
+class Dataflow:
+    """
+    Dataflow class to interact with Dataflow services.
+    """
+
+    @staticmethod
+    def _json_parse(value):  
+        try:
+            result = json.loads(value)
+            if isinstance(result, str):
+                try:
+                    return json.loads(result)
+                except json.JSONDecodeError:
+                    return result
+            return result
+        except (json.JSONDecodeError, TypeError):
+            return value
+    
+    def _parse_response_data(self, response):
+        """Parse response data based on datatype field or fallback to JSON parsing."""
+        data = response.json()
+        if not isinstance(data, dict):
+            raise ValueError("Internal Dataflow Error!")
+        value = data.get('value', '')
+        if data.get('datatype') == 'json':
+            return self._json_parse(value)
+        else:
+            return value
+
+    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", "")
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+            org_id = os.environ.get("ORGANIZATION")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            query_params = {
+                "key": variable_name,
+            }
+
+            variable_api = None
+            if runtime and slug:
+                variable_api = dataflow_config.get_config_value("auth", "variable_ui_api")
+                query_params["runtime"] = runtime
+                query_params["slug"] = slug
+                query_params["org_id"] = org_id
+            elif host_name:
+                variable_api = dataflow_config.get_config_value("auth", "variable_manager_api")
+            else:
+                raise Exception("Cannot run dataflow methods here!") 
+            
+            if not variable_api:
+                print("[Dataflow.variable] Variable Unreachable")
+                return None
+        
+            response = requests.get(variable_api, params=query_params)
+            
+            if response.status_code == 404:
+                return None
+            elif response.status_code >= 500:
+                response.raise_for_status()
+            elif response.status_code >= 400:
+                print(f"[Dataflow.variable] Client error {response.status_code} for variable '{variable_name}'")
+                return None
+            elif response.status_code != 200:
+                print(f"[Dataflow.variable] Unexpected status {response.status_code} for variable '{variable_name}'")
+                return None
+           
+            return self._parse_response_data(response)
+
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.variable] Failed to fetch variable '{variable_name}'") from e
+        
+        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", "")
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+            org_id = os.environ.get("ORGANIZATION")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            query_params = {
+                "key": secret_name
+            }
+
+            if runtime:
+                secret_api = dataflow_config.get_config_value("auth", "secret_ui_api")
+                query_params["runtime"] = runtime
+                query_params["slug"] = slug
+                query_params["org_id"] = org_id
+            else:
+                secret_api = dataflow_config.get_config_value("auth", "secret_manager_api")
+            if not secret_api:
+                print("[Dataflow.secret] Secret API Unreachable")
+                return None
+
+            response = requests.get(secret_api, params=query_params)
+
+            if response.status_code == 404:
+                return None
+            elif response.status_code >= 500:
+                response.raise_for_status()
+            elif response.status_code >= 400:
+                print(f"[Dataflow.secret] Client error {response.status_code} for secret '{secret_name}'")
+                return None
+            elif response.status_code != 200:
+                print(f"[Dataflow.secret] Unexpected status {response.status_code} for secret '{secret_name}'")
+                return None
+
+            return self._parse_response_data(response)
+            
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.secret] Failed to fetch secret '{secret_name}'") from e
+        except Exception as e:
+            print(f"[Dataflow.secret] Exception occurred: {e}")
+            return None
+        
+    def secret_file(self, secret_name: str):
+        """
+        Retrieve a Dataflow secret file.
+
+        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", "")
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+            org_id = os.environ.get("ORGANIZATION")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            query_params = {
+                "key": secret_name
+            }
+
+            if runtime:
+                secret_api = dataflow_config.get_config_value("auth", "secret_ui_api")
+                query_params["runtime"] = runtime
+                query_params["slug"] = slug
+                query_params["org_id"] = org_id
+            else:
+                secret_api = dataflow_config.get_config_value("auth", "secret_manager_api")
+            if not secret_api:
+                print("[Dataflow.secret] Secret API Unreachable")
+                return None
+
+            response = requests.get(secret_api, params=query_params)
+
+            if response.status_code == 404:
+                return None
+            elif response.status_code >= 500:
+                response.raise_for_status()
+            elif response.status_code >= 400:
+                print(f"[Dataflow.secret] Client error {response.status_code} for secret '{secret_name}'")
+                return None
+            elif response.status_code != 200:
+                print(f"[Dataflow.secret] Unexpected status {response.status_code} for secret '{secret_name}'")
+                return None
+            
+            response_data = response.json()
+            if response.status_code == 200 and response_data.get('filename'):
+                # For runtime mode, create file and return filepath
+                if runtime:
+                    import tempfile
+                    from pathlib import Path
+
+                    # Create /tmp/secrets directory if it doesn't exist
+                    secrets_dir = Path("/tmp/secrets")
+                    secrets_dir.mkdir(parents=True, exist_ok=True)
+                    
+                    # Get filename and content
+                    filename = response_data.get('filename')
+                    file_content = response_data.get('value')
+                    
+                    if not filename or not file_content:
+                        print(f"[Dataflow.secret] Missing filename or content for secret '{secret_name}'")
+                        return None
+                    
+                    file_path = os.path.join(secrets_dir, filename)
+        
+                    # Detect if content is Base64 encoded binary or text
+                    try:
+                        # Try to decode as Base64
+                        decoded_content = base64.b64decode(file_content, validate=True)
+                        # Check if it contains non-printable characters (likely binary)
+                        is_binary = not all(32 <= byte <= 126 or byte in (9, 10, 13) for byte in decoded_content[:100])
+                        
+                        if is_binary:
+                            # Write as binary
+                            with open(file_path, 'wb') as f:
+                                f.write(decoded_content)
+                        else:
+                            # Decode and write as text
+                            with open(file_path, 'w', encoding='utf-8') as f:
+                                f.write(decoded_content.decode('utf-8'))
+                    except Exception:
+                        # Not Base64 or decode failed, treat as text
+                        with open(file_path, 'w', encoding='utf-8') as f:
+                            f.write(file_content)
+                    return str(file_path)
+                else:
+                    # For non-runtime mode, return the value as-is
+                    return response_data.get('value')
+            else:
+                print(f"[Dataflow.secret] No file found for secret '{secret_name}'! If it is a non-file secret, please use the 'secret' method.")
+                return None
+        
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.secret] Failed to fetch secret '{secret_name}'") from e
+        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"]
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+            org_id = os.environ.get("ORGANIZATION")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            query_params = {
+                "conn_id": conn_id
+            }
+
+            if runtime:
+                query_params["runtime"] = runtime
+                query_params["org_id"] = org_id
+                query_params["slug"] = slug
+                connection_api = dataflow_config.get_config_value("auth", "connection_ui_api")
+            elif host_name:
+                connection_api = dataflow_config.get_config_value("auth", "connection_manager_api")
+            else:
+                raise Exception("Cannot run dataflow methods here! HOSTNAME or RUNTIME env variable not set.")
+
+            response = requests.get(connection_api, params=query_params)
+            
+            if response.status_code == 404:
+                raise RuntimeError(f"[Dataflow.connection] Connection '{conn_id}' not found!")
+            elif response.status_code >= 500:
+                response.raise_for_status()
+            elif response.status_code >= 400:
+                raise RuntimeError(f"[Dataflow.connection] Client error {response.status_code} for connection '{conn_id}'")
+            elif response.status_code != 200:
+                raise RuntimeError(f"[Dataflow.connection] Unexpected status {response.status_code} for connection '{conn_id}'")
+                
+            connection_details = response.json()
+
+            if not connection_details:
+                raise RuntimeError(f"[Dataflow.connection] Connection '{conn_id}' not found!")
+            
+            if mode == "dict":
+                return dict(connection_details)
+
+            conn_type = connection_details['conn_type'].lower()
+            username = connection_details['login']
+            password = connection_details.get('password', '')
+            host = connection_details['host']
+            port = connection_details['port']
+            database = connection_details.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 = connection_details.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())
+            else:
+                raise ValueError(f"Unsupported mode: {mode}. Use 'session', 'engine', 'url'.")
+        
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.connection] Failed to fetch connection '{conn_id}'") from e
+
+        except Exception as e:
+            raise RuntimeError(f"[Dataflow.connection] Error connecting to '{conn_id}': {str(e)}") from e
+        
+    def variable_or_secret(self, key: str):
+        """
+        Retrieve a variable or secret by key.
+        
+        Args:
+            key (str): Key of the variable or secret
+            
+        Returns:
+            str or None: Value if found, None otherwise
+        """
+        try:
+            host_name = os.environ.get("HOSTNAME", "")
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+            org_id = os.environ.get("ORGANIZATION")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            query_params = {
+                    "key": key
+                }
+            
+            if runtime:
+                variableorsecret_api = dataflow_config.get_config_value("auth", "variableorsecret_ui_api")
+                query_params["runtime"] = runtime
+                query_params["slug"] = slug
+                query_params["org_id"] = org_id
+            elif host_name:
+                variableorsecret_api = dataflow_config.get_config_value("auth", "variableorsecret_manager_api")
+            else:
+                raise Exception("Cannot run dataflow methods here!") 
+
+            if not variableorsecret_api:
+                print("[Dataflow.variable_or_secret] Variable/Secret Unreachable")
+                return None
+            
+            response = requests.get(variableorsecret_api, params=query_params)
+            
+            if response.status_code == 404:
+                return None
+            elif response.status_code >= 500:
+                response.raise_for_status()  # Let server errors propagate
+            elif response.status_code >= 400:
+                print(f"[Dataflow.variable_or_secret] Client error {response.status_code} for key '{key}'")
+                return None
+            elif response.status_code != 200:
+                print(f"[Dataflow.variable_or_secret] Unexpected status {response.status_code} for key '{key}'")
+                return None
+
+            return self._parse_response_data(response)
+            
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.variable_or_secret] Failed to fetch '{key}'") from e