dataflow-core 2.1.7__tar.gz → 2.1.8__tar.gz

{dataflow_core-2.1.7 → dataflow_core-2.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataflow-core
-Version: 2.1.7
+Version: 2.1.8
 Summary: Dataflow core package
 Author: Dataflow
 Author-email: 
@@ -9,6 +9,8 @@ Requires-Dist: boto3
 Requires-Dist: psycopg2-binary
 Requires-Dist: pymysql
 Requires-Dist: requests
+Requires-Dist: azure-identity
+Requires-Dist: azure-keyvault-secrets
 Dynamic: author
 Dynamic: requires-dist
 Dynamic: summary

{dataflow_core-2.1.7 → dataflow_core-2.1.8}/authenticator/dataflowhubauthenticator.py

@@ -109,7 +109,7 @@ class DataflowBaseAuthenticator(Authenticator):
                 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
@@ -247,7 +247,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
@@ -269,48 +273,47 @@ 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)
             self.log.info(f"Azure AD OAuth completed for user: {username}, session_id={session_id}")
             return {
-                "name": username,
+                "name": db_user.first_name,
                 "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.8/dataflow/dataflow.py

@@ -0,0 +1,323 @@
+import os, requests
+from .database_manager import DatabaseManager
+import json
+from .configuration import ConfigurationManager
+
+
+class Dataflow:
+    """
+    Dataflow class to interact with Dataflow services.
+    """
+    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")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+
+            variable_api = None
+            if runtime and slug:
+                variable_api = dataflow_config.get_config_value("auth", "variable_ui_api")
+            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
+
+            if runtime:
+                query_params = {
+                    "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 = {
+                "key": variable_name,
+            }
+            response = requests.get(variable_api, params=query_params)
+            
+            # Handle different HTTP status codes gracefully
+            if response.status_code == 404:
+                return None  # Variable not found
+            elif response.status_code >= 500:
+                response.raise_for_status()  # Let server errors propagate
+            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 response.text.strip().strip('"')
+        
+        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")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            if runtime:
+                secret_api = dataflow_config.get_config_value("auth", "secret_ui_api")
+            else:
+                secret_api = dataflow_config.get_config_value("auth", "secret_manager_api")
+            if not secret_api:
+                print("[Dataflow.secret] Secret API Unreachable")
+                return None
+
+            query_params = {
+                "key": secret_name
+            }
+
+            if runtime:
+                query_params["runtime"] = runtime
+            if slug:
+                query_params["slug"] = slug
+
+            response = requests.get(secret_api, params=query_params)
+            
+            # Handle different HTTP status codes gracefully
+            if response.status_code == 404:
+                return None  # Secret not found
+            elif response.status_code >= 500:
+                response.raise_for_status()  # Let server errors propagate
+            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 response.text.strip().strip('"')
+        
+        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")
+            
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            if runtime:
+                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.")
+            
+            query_params = {
+                "conn_id": conn_id
+            }
+
+            if runtime:
+                query_params["runtime"] = runtime
+            if slug:
+                query_params["slug"] = slug
+
+            response = requests.get(connection_api, params=query_params)
+            
+            # Handle different HTTP status codes gracefully
+            if response.status_code == 404:
+                raise RuntimeError(f"[Dataflow.connection] Connection '{conn_id}' not found!")
+            elif response.status_code >= 500:
+                response.raise_for_status()  # Let server errors propagate
+            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":
+                with open('/home/jovyan/log.txt', 'w') as log_file:
+                    log_file.write(f"Connection details for {conn_id}: {connection_details}\n")
+                print(f"connection_details: {connection_details}")
+                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")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            if runtime and slug:
+                variableorsecret_api = dataflow_config.get_config_value("auth", "variableorsecret_ui_api")
+                query_params = {
+                    "key": key,
+                    "runtime": runtime,
+                    "slug": slug
+                }
+            elif host_name:
+                variableorsecret_api = dataflow_config.get_config_value("auth", "variableorsecret_manager_api")
+                query_params = {
+                    "key": key
+                }
+            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)
+            
+            # Handle different HTTP status codes gracefully
+            if response.status_code == 404:
+                return None  # Variable/secret not found
+            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
+
+            response_text = response.text.strip().strip('"')
+            return response_text
+            
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.variable_or_secret] Failed to fetch '{key}'") from e

dataflow_core-2.1.8/dataflow/schemas/connection.py

@@ -0,0 +1,84 @@
+"""schemas/connection.py"""
+
+from pydantic import BaseModel, field_validator
+from typing import Optional
+from datetime import datetime
+from enum import Enum
+
+
+class ConnectionType(str, Enum):
+    """Enum for supported connection types."""
+    POSTGRESQL = "PostgreSQL"
+    MYSQL = "MySQL"
+
+
+class ConnectionBase(BaseModel):
+    """Base connection model with common fields."""
+    conn_id: str
+    conn_type: ConnectionType
+    description: Optional[str] = None
+    host: str
+    schemas: Optional[str] = None
+    password: str
+    login: str
+    port: int
+    extra: Optional[str] = None
+
+    @field_validator("conn_id")
+    def validate_conn_id(cls, v) -> str:
+        import re
+        if not isinstance(v, str):
+            raise ValueError("Connection ID must be a string.")
+        if len(v) > 20:
+            raise ValueError("Connection ID must be at most 20 characters long.")
+        if not re.fullmatch(r"[A-Za-z0-9-]+", v):
+            raise ValueError(
+                "Connection ID can only contain letters, numbers, and hyphens (-)!"
+            )
+        return v
+
+    @field_validator("conn_type")
+    def validate_conn_type(cls, v) -> ConnectionType:
+        if isinstance(v, str):
+            try:
+                return ConnectionType(v)
+            except ValueError:
+                raise ValueError(f'conn_type must be one of {[e.value for e in ConnectionType]}')
+        return v
+
+
+class ConnectionSave(ConnectionBase):
+    """Model for creating a new connection."""
+    pass
+
+
+class ConnectionUpdate(BaseModel):
+    """Model for updating an existing connection."""
+    conn_type: Optional[ConnectionType] = None
+    description: Optional[str] = None
+    host: Optional[str] = None
+    schemas: Optional[str] = None
+    login: Optional[str] = None
+    password: Optional[str] = None
+    port: Optional[int] = None
+    extra: Optional[str] = None
+
+    @field_validator("conn_type")
+    def validate_conn_type(cls, v) -> Optional[ConnectionType]:
+        if v is None:
+            return v
+        if isinstance(v, str):
+            # Convert string to enum if needed
+            try:
+                return ConnectionType(v)
+            except ValueError:
+                raise ValueError(f'conn_type must be one of {[e.value for e in ConnectionType]}')
+        return v
+
+
+class ConnectionRead(ConnectionBase):
+    """Model for reading/displaying connection data."""
+    pass
+
+    class Config:
+        from_attributes = True

dataflow_core-2.1.8/dataflow/schemas/git_ssh.py

@@ -0,0 +1,50 @@
+"""schemas/git_ssh.py"""
+
+from pydantic import BaseModel, field_validator
+from typing import Optional, Literal
+from datetime import datetime
+
+
+class SSHBase(BaseModel):
+    """Base SSH key model with common fields."""
+    key_name: str
+    description: Optional[str] = None
+
+    @field_validator("key_name")
+    def validate_key_name(cls, v) -> str:
+        import re
+        if not isinstance(v, str):
+            raise ValueError("SSH key name must be a string.")
+        if len(v) > 20:
+            raise ValueError("SSH key name must be at most 20 characters long.")
+        if not re.fullmatch(r"[A-Za-z0-9-]+", v):
+            raise ValueError(
+                "SSH key name can only contain letters, numbers, and hyphens (-)!"
+            )
+        return v
+
+
+class SSHSave(SSHBase):
+    """Model for creating a new SSH key."""
+    public_key: str
+    private_key: str
+
+
+class SSHUpdate(BaseModel):
+    """Model for updating an existing SSH key."""
+    description: Optional[str] = None
+    public_key: Optional[str] = None
+    private_key: Optional[str] = None
+
+
+class SSHRead(SSHBase):
+    """Model for reading/displaying SSH key data."""
+    public_key: str
+    private_key: str
+    created_date: Optional[datetime] = None
+
+    class Config:
+        from_attributes = True
+
+    class Config:
+        from_attributes = True

dataflow_core-2.1.8/dataflow/schemas/secret.py

@@ -0,0 +1,44 @@
+"""schemas/secret.py"""
+
+from pydantic import BaseModel, field_validator
+from typing import Optional
+from datetime import datetime
+
+
+class SecretBase(BaseModel):
+    """Base secret model with common fields."""
+    key: str
+    value: str
+    description: Optional[str] = None
+
+    @field_validator("key")
+    def validate_key(cls, v) -> str:
+        import re
+        if not isinstance(v, str):
+            raise ValueError("Secret key must be a string.")
+        if len(v) > 20:
+            raise ValueError("Secret key must be at most 20 characters long.")
+        if not re.fullmatch(r"[A-Za-z0-9-]+", v):
+            raise ValueError(
+                "Secret key can only contain letters, numbers, and hyphens (-)!"
+            )
+        return v
+
+
+class SecretSave(SecretBase):
+    """Model for creating a new secret."""
+    pass
+
+
+class SecretUpdate(BaseModel):
+    """Model for updating an existing secret."""
+    value: Optional[str] = None
+    description: Optional[str] = None
+
+
+class SecretRead(SecretBase):
+    """Model for reading/displaying secret data."""
+    created_date: Optional[datetime] = None
+
+    class Config:
+        from_attributes = True

dataflow_core-2.1.8/dataflow/secrets_manager/__init__.py

@@ -0,0 +1,13 @@
+# secrets_manager/__init__.py
+
+from .factory import get_secret_manager
+from .service import SecretsService
+
+# 1. Call the factory to get the configured low-level secret manager
+#    (e.g., an instance of AWSSecretsManager or AzureKeyVault).
+#    This happens only once when the package is first imported.
+secret_manager_instance = get_secret_manager()
+
+# 2. Create the single, high-level service instance that the rest of
+#    your application will use. It wraps the low-level instance.
+secrets_service = SecretsService(secret_manager=secret_manager_instance)

dataflow_core-2.1.8/dataflow/secrets_manager/factory.py

@@ -0,0 +1,59 @@
+# secrets_manager/factory.py
+import os
+from .interface import SecretManager
+from .providers.aws_manager import AWSSecretsManager
+from .providers.azure_manager import AzureKeyVault
+from ..configuration import ConfigurationManager
+
+# A custom exception for clear error messages
+class SecretProviderError(Exception):
+    pass
+
+def get_secret_manager() -> SecretManager:
+    """
+    Factory function to get the configured secret manager instance.
+
+    Reads the cloud provider configuration from dataflow_auth.cfg
+    to determine which cloud provider's secret manager to instantiate.
+    """
+    try:
+        # dataflow_config = None
+        # if os.getenv('HOSTNAME'):
+        #     dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+        # else:
+        dataflow_config = ConfigurationManager('/dataflow/app/config/dataflow.cfg')
+    except Exception as e:
+        raise SecretProviderError(
+            f"Failed to read cloud provider configuration: {str(e)}. "
+            "Please check that the configuration file exists and contains the 'cloud' section."
+        )
+
+    provider = dataflow_config.get_config_value('cloudProvider', 'cloud')
+    if not provider:
+        raise SecretProviderError(
+            "The cloud provider is not configured in config file. "
+            "Please set the 'cloud' value in the 'cloud' section to 'aws' or 'azure'."
+        )
+
+    provider = provider.lower()
+    print(f"Initializing secret manager for provider: {provider}")
+
+    if provider == "aws":
+        return AWSSecretsManager()
+
+    elif provider == "azure":
+        vault_url = dataflow_config.get_config_value('cloudProvider', 'key_vault')
+        if not vault_url:
+            raise SecretProviderError(
+                "AZURE_VAULT_URL must be set when using the Azure provider."
+            )
+        return AzureKeyVault(vault_url=vault_url)
+
+    # You can easily add more providers here in the future
+    # elif provider == "gcp":
+    #     return GCPSecretManager()
+
+    else:
+        raise SecretProviderError(
+            f"Unsupported secret provider: '{provider}'. Supported providers are: aws, azure."
+        )