Flowfile 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

flowfile_core/flowfile/database_connection_manager/db_connections.py

@@ -1,6 +1,8 @@
 from flowfile_core.schemas.input_schema import FullDatabaseConnection, FullDatabaseConnectionInterface
+from flowfile_core.schemas.cloud_storage_schemas import FullCloudStorageConnection, FullCloudStorageConnectionInterface
 from sqlalchemy.orm import Session
-from flowfile_core.database.models import DatabaseConnection as DBConnectionModel, Secret
+from flowfile_core.database.models import (DatabaseConnection as DBConnectionModel, Secret,
+                                           CloudStorageConnection as DBCloudStorageConnection)
 from flowfile_core.secret_manager.secret_manager import store_secret, SecretInput, decrypt_secret
 from flowfile_core.database.connection import get_db_context
 
@@ -53,6 +55,18 @@ def get_database_connection(db: Session, connection_name: str, user_id: int) ->
     return db_connection
 
 
+def get_cloud_connection(db: Session, connection_name: str, user_id: int) -> DBCloudStorageConnection | None:
+    """
+    Get a cloud storage connection by its name and user ID.
+    """
+    db_connection = db.query(DBCloudStorageConnection).filter(
+        DBCloudStorageConnection.connection_name == connection_name,
+        DBCloudStorageConnection.user_id == user_id
+    ).first()
+
+    return db_connection
+
+
 def get_database_connection_schema(db: Session, connection_name: str, user_id: int) -> FullDatabaseConnection | None:
     """
     Get a database connection schema by its name and user ID.
@@ -84,6 +98,20 @@ def get_local_database_connection(connection_name: str, user_id: int) -> FullDat
         return get_database_connection_schema(db, connection_name, user_id)
 
 
+def get_local_cloud_connection(connection_name: str, user_id: int) -> FullCloudStorageConnection | None:
+    """
+    Get a cloud storage connection schema by its name and user ID.
+    Args:
+        connection_name (str): The name of the cloud storage connection.
+        user_id (int): The ID of the user who owns the connection.
+
+    Returns:
+        FullCloudStorageConnection | None: The cloud storage connection schema if found, otherwise None.
+    """
+    with get_db_context() as db:
+        return get_cloud_connection_schema(db, connection_name, user_id)
+
+
 def delete_database_connection(db: Session, connection_name: str, user_id: int) -> None:
     """
     Delete a database connection by its name and user ID.
@@ -102,7 +130,8 @@ def delete_database_connection(db: Session, connection_name: str, user_id: int)
         db.commit()
 
 
-def database_connection_interface_from_db_connection(db_connection: DBConnectionModel) -> FullDatabaseConnectionInterface:
+def database_connection_interface_from_db_connection(
+        db_connection: DBConnectionModel) -> FullDatabaseConnectionInterface:
     """
     Convert a database connection from the database model to the interface model.
     """
@@ -137,3 +166,188 @@ def get_all_database_connections_interface(db: Session, user_id: int) -> list[Fu
             raise TypeError(f"Expected a DBConnectionModel instance, got {type(db_connection)}")
 
     return result
+
+
+def store_cloud_connection(db: Session, connection: FullCloudStorageConnection, user_id: int) -> DBCloudStorageConnection:
+    """
+    Placeholder function to store a cloud database connection.
+    This function should be implemented based on specific cloud provider requirements.
+    """
+    existing_database_connection = get_cloud_connection(db, connection.connection_name, user_id)
+    if existing_database_connection:
+        raise ValueError(
+            f"Database connection with name '{connection.connection_name}' already exists for user {user_id}."
+            f" Please use a unique connection name or delete the existing connection first."
+        )
+    if connection.aws_secret_access_key is not None:
+        aws_secret_access_key_ref_id = store_secret(db,
+                                                    SecretInput(name=connection.connection_name + "_aws_secret_access_key",
+                                                                value=connection.aws_secret_access_key), user_id).id
+    else:
+        aws_secret_access_key_ref_id = None
+    if connection.azure_client_secret is not None:
+        azure_client_secret_ref_id = store_secret(db,
+                                                  SecretInput(name=connection.connection_name + "azure_client_secret",
+                                                              value=connection.azure_client_secret), user_id).id
+    else:
+        azure_client_secret_ref_id = None
+    if connection.azure_account_key is not None:
+        azure_account_key_ref_id = store_secret(db, SecretInput(name=connection.connection_name + "azure_account_key",
+                                                                value=connection.azure_account_key), user_id).id
+    else:
+        azure_account_key_ref_id = None
+
+    db_cloud_connection = DBCloudStorageConnection(
+        connection_name=connection.connection_name,
+        storage_type=connection.storage_type,
+        auth_method=connection.auth_method,
+        user_id=user_id,
+
+            # AWS S3 fields
+        aws_region=connection.aws_region,
+        aws_access_key_id=connection.aws_access_key_id,
+        aws_role_arn=connection.aws_role_arn,
+        aws_secret_access_key_id=aws_secret_access_key_ref_id,
+        aws_allow_unsafe_html=connection.aws_allow_unsafe_html,
+
+        # Azure ADLS fields
+        azure_account_name=connection.azure_account_name,
+        azure_tenant_id=connection.azure_tenant_id,
+        azure_client_id=connection.azure_client_id,
+        azure_account_key_id=azure_account_key_ref_id,
+        azure_client_secret_id=azure_client_secret_ref_id,
+
+        # Common fields
+        endpoint_url=connection.endpoint_url,
+        verify_ssl=connection.verify_ssl
+    )
+    db.add(db_cloud_connection)
+    db.commit()
+    db.refresh(db_cloud_connection)
+    return db_cloud_connection
+
+
+def get_full_cloud_storage_interface_from_db(
+        db_cloud_connection: DBCloudStorageConnection) -> FullCloudStorageConnectionInterface:
+    """
+    Convert a cloud storage connection from the database model to the interface model.
+    """
+    return FullCloudStorageConnectionInterface(
+        connection_name=db_cloud_connection.connection_name,
+        storage_type=db_cloud_connection.storage_type,
+        auth_method=db_cloud_connection.auth_method,
+        aws_allow_unsafe_html=db_cloud_connection.aws_allow_unsafe_html,
+        aws_region=db_cloud_connection.aws_region,
+        aws_access_key_id=db_cloud_connection.aws_access_key_id,
+        aws_role_arn=db_cloud_connection.aws_role_arn,
+        azure_account_name=db_cloud_connection.azure_account_name,
+        azure_tenant_id=db_cloud_connection.azure_tenant_id,
+        azure_client_id=db_cloud_connection.azure_client_id,
+        endpoint_url=db_cloud_connection.endpoint_url,
+        verify_ssl=db_cloud_connection.verify_ssl
+    )
+
+
+def get_cloud_connection_schema(db: Session, connection_name: str, user_id: int) -> FullCloudStorageConnection | None:
+    """
+    Retrieves a full cloud storage connection schema, including decrypted secrets, by its name and user ID.
+    """
+    db_connection = get_cloud_connection(db, connection_name, user_id)
+    if not db_connection:
+        return None
+
+    # Decrypt secrets associated with the connection
+    aws_secret_key = None
+    if db_connection.aws_secret_access_key_id:
+        secret_record = db.query(Secret).filter(Secret.id == db_connection.aws_secret_access_key_id).first()
+        if secret_record:
+            aws_secret_key = decrypt_secret(secret_record.encrypted_value)
+
+    azure_account_key = None
+    if db_connection.azure_account_key_id:
+        secret_record = db.query(Secret).filter(Secret.id == db_connection.azure_account_key_id).first()
+        if secret_record:
+            azure_account_key = decrypt_secret(secret_record.encrypted_value)
+
+    azure_client_secret = None
+    if db_connection.azure_client_secret_id:
+        secret_record = db.query(Secret).filter(Secret.id == db_connection.azure_client_secret_id).first()
+        if secret_record:
+            azure_client_secret = decrypt_secret(secret_record.encrypted_value)
+
+    # Construct the full Pydantic model
+    return FullCloudStorageConnection(
+        connection_name=db_connection.connection_name,
+        storage_type=db_connection.storage_type,
+        auth_method=db_connection.auth_method,
+        aws_allow_unsafe_html=db_connection.aws_allow_unsafe_html,
+        aws_region=db_connection.aws_region,
+        aws_access_key_id=db_connection.aws_access_key_id,
+        aws_secret_access_key=aws_secret_key,
+        aws_role_arn=db_connection.aws_role_arn,
+        azure_account_name=db_connection.azure_account_name,
+        azure_account_key=azure_account_key,
+        azure_tenant_id=db_connection.azure_tenant_id,
+        azure_client_id=db_connection.azure_client_id,
+        azure_client_secret=azure_client_secret,
+        endpoint_url=db_connection.endpoint_url,
+        verify_ssl=db_connection.verify_ssl
+    )
+
+
+def cloud_connection_interface_from_db_connection(
+        db_connection: DBCloudStorageConnection) -> FullCloudStorageConnectionInterface:
+    """
+    Converts a DBCloudStorageConnection model to a FullCloudStorageConnectionInterface model,
+    which safely exposes non-sensitive data.
+    """
+    return FullCloudStorageConnectionInterface(
+        connection_name=db_connection.connection_name,
+        storage_type=db_connection.storage_type,
+        auth_method=db_connection.auth_method,
+        aws_allow_unsafe_html=db_connection.aws_allow_unsafe_html,
+        aws_region=db_connection.aws_region,
+        aws_access_key_id=db_connection.aws_access_key_id,
+        aws_role_arn=db_connection.aws_role_arn,
+        azure_account_name=db_connection.azure_account_name,
+        azure_tenant_id=db_connection.azure_tenant_id,
+        azure_client_id=db_connection.azure_client_id,
+        endpoint_url=db_connection.endpoint_url,
+        verify_ssl=db_connection.verify_ssl
+    )
+
+
+def get_all_cloud_connections_interface(db: Session, user_id: int) -> list[FullCloudStorageConnectionInterface]:
+    """
+    Retrieves a list of all cloud storage connections for a user in a safe interface format (no secrets).
+    """
+    db_connections = db.query(DBCloudStorageConnection).filter(DBCloudStorageConnection.user_id == user_id).all()
+
+    return [cloud_connection_interface_from_db_connection(conn) for conn in db_connections]
+
+
+def delete_cloud_connection(db: Session, connection_name: str, user_id: int) -> None:
+    """
+    Deletes a cloud storage connection and all of its associated secrets from the database.
+    """
+    db_connection = get_cloud_connection(db, connection_name, user_id)
+
+    if db_connection:
+        # Collect all secret IDs associated with this connection
+        secret_ids_to_delete = [
+            db_connection.aws_secret_access_key_id,
+            db_connection.aws_session_token_id,
+            db_connection.azure_account_key_id,
+            db_connection.azure_client_secret_id,
+            db_connection.azure_sas_token_id
+        ]
+        # Filter out None values
+        secret_ids_to_delete = [id for id in secret_ids_to_delete if id is not None]
+
+        # Delete associated secrets if they exist
+        if secret_ids_to_delete:
+            db.query(Secret).filter(Secret.id.in_(secret_ids_to_delete)).delete(synchronize_session=False)
+
+        # Delete the connection record itself
+        db.delete(db_connection)
+        db.commit()

flowfile_core/flowfile/extensions.py

@@ -17,7 +17,7 @@ def get_instant_func_results(node_step: FlowNode, func_string: str) -> InstantFu
         return InstantFuncResult(result='No input data connected, so cannot evaluate the result', success=None)
     node_input = node_step.main_input[0]
     try:
-        if node_input.node_stats.has_run and node_input.is_setup and node_input.results.example_data_path:
+        if node_input.node_stats.has_run_with_current_setup and node_input.is_setup and node_input.results.example_data_path:
             df = get_first_row(node_input.results.example_data_path)
         else:
             df = node_input.get_predicted_resulting_data().data_frame.collect()

flowfile_core/flowfile/flow_data_engine/cloud_storage_reader.py

@@ -0,0 +1,259 @@
+import boto3
+from botocore.exceptions import ClientError
+from typing import Optional, Dict, Any, Callable, Literal
+
+from flowfile_core.schemas.cloud_storage_schemas import FullCloudStorageConnection
+
+
+def create_storage_options_from_boto_credentials(profile_name: Optional[str],
+                                                 region_name: Optional[str] = None) -> Dict[str, Any]:
+    """
+    Create a storage options dictionary from AWS credentials using a boto3 profile.
+    This is the most robust way to handle profile-based authentication as it
+    bypasses Polars' internal credential provider chain, avoiding conflicts.
+
+    Parameters
+    ----------
+    profile_name
+        The name of the AWS profile in ~/.aws/credentials.
+    region_name
+        The AWS region to use.
+
+    Returns
+    -------
+    Dict[str, Any]
+        A storage options dictionary for Polars with explicit credentials.
+    """
+    session = boto3.Session(profile_name=profile_name, region_name=region_name)
+    credentials = session.get_credentials()
+    frozen_creds = credentials.get_frozen_credentials()
+
+    storage_options = {
+        "aws_access_key_id": frozen_creds.access_key,
+        "aws_secret_access_key": frozen_creds.secret_key,
+        "aws_session_token": frozen_creds.token,
+    }
+    # Use the session's region if one was resolved, otherwise use the provided one
+    if session.region_name:
+        storage_options["aws_region"] = session.region_name
+
+    print("Boto3: Successfully created storage options with explicit credentials.")
+    return storage_options
+
+
+class CloudStorageReader:
+    """Helper class to handle different cloud storage authentication methods and read operations."""
+
+    @staticmethod
+    def get_storage_options(connection: FullCloudStorageConnection) -> Dict[str, Any]:
+        """
+        Build storage options dict based on the connection type and auth method.
+
+        Args:
+            connection: Full connection details with decrypted secrets
+
+        Returns:
+            Dict containing appropriate storage options for the provider
+        """
+        if connection.storage_type == "s3":
+            return CloudStorageReader._get_s3_storage_options(connection)
+        elif connection.storage_type == "adls":
+            return CloudStorageReader._get_adls_storage_options(connection)
+        elif connection.storage_type == "gcs":
+            return CloudStorageReader._get_gcs_storage_options(connection)
+        else:
+            raise ValueError(f"Unsupported storage type: {connection.storage_type}")
+
+    @staticmethod
+    def _get_s3_storage_options(connection: 'FullCloudStorageConnection') -> Dict[str, Any]:
+        """Build S3-specific storage options."""
+        auth_method = connection.auth_method
+        print(f"Building S3 storage options for auth_method: '{auth_method}'")
+        if auth_method == "aws-cli":
+            return create_storage_options_from_boto_credentials(
+                profile_name=connection.connection_name,
+                region_name=connection.aws_region
+            )
+
+        storage_options = {}
+        if connection.aws_region:
+            storage_options["aws_region"] = connection.aws_region
+        if connection.endpoint_url:
+            storage_options["endpoint_url"] = connection.endpoint_url
+        if not connection.verify_ssl:
+            storage_options["verify"] = "False"
+        if connection.aws_allow_unsafe_html: # Note: Polars uses aws_allow_http
+            storage_options["aws_allow_http"] = "true"
+
+        if auth_method == "access_key":
+            storage_options["aws_access_key_id"] = connection.aws_access_key_id
+            storage_options["aws_secret_access_key"] = connection.aws_secret_access_key.get_secret_value()
+            # Explicitly clear any session token from the environment
+            storage_options["aws_session_token"] = ""
+
+        elif auth_method == "iam_role":
+            # Correctly implement IAM role assumption using boto3 STS client.
+            sts_client = boto3.client('sts', region_name=connection.aws_region)
+            assumed_role_object = sts_client.assume_role(
+                RoleArn=connection.aws_role_arn,
+                RoleSessionName="PolarsCloudStorageReaderSession" # A descriptive session name
+            )
+            credentials = assumed_role_object['Credentials']
+            storage_options["aws_access_key_id"] = credentials['AccessKeyId']
+            storage_options["aws_secret_access_key"] = credentials['SecretAccessKey']
+            storage_options["aws_session_token"] = credentials['SessionToken']
+
+        return storage_options
+
+    @staticmethod
+    def _get_adls_storage_options(connection: 'FullCloudStorageConnection') -> Dict[str, Any]:
+        """Build Azure ADLS-specific storage options."""
+        storage_options = {}
+
+        if connection.auth_method == "access_key":
+            # Account key authentication
+            if connection.azure_account_name:
+                storage_options["account_name"] = connection.azure_account_name
+            if connection.azure_account_key:
+                storage_options["account_key"] = connection.azure_account_key.get_secret_value()
+
+        elif connection.auth_method == "service_principal":
+            # Service principal authentication
+            if connection.azure_tenant_id:
+                storage_options["tenant_id"] = connection.azure_tenant_id
+            if connection.azure_client_id:
+                storage_options["client_id"] = connection.azure_client_id
+            if connection.azure_client_secret:
+                storage_options["client_secret"] = connection.azure_client_secret.get_secret_value()
+
+        elif connection.auth_method == "sas_token":
+            # SAS token authentication
+            if connection.azure_sas_token:
+                storage_options["sas_token"] = connection.azure_sas_token.get_secret_value()
+
+        return storage_options
+
+    @staticmethod
+    def _get_gcs_storage_options(connection: 'FullCloudStorageConnection') -> Dict[str, Any]:
+        """Build GCS-specific storage options."""
+        # GCS typically uses service account authentication
+        # Implementation would depend on how credentials are stored
+        return {}
+
+    @staticmethod
+    def get_credential_provider(connection: 'FullCloudStorageConnection') -> Optional[Callable]:
+        """
+        Get a credential provider function if needed for the authentication method.
+
+        Args:
+            connection: Full connection details
+
+        Returns:
+            Credential provider function or None
+        """
+        if connection.storage_type == "s3" and connection.auth_method == "iam_role":
+            # For IAM role, create a credential provider
+            def aws_credential_provider():
+                # This would typically use boto3 to assume the role
+                # For now, returning a placeholder
+                return {
+                    "aws_access_key_id": "...",
+                    "aws_secret_access_key": "...",
+                    "aws_session_token": "...",
+                }, None  # expiry
+
+            return aws_credential_provider
+        return None
+
+
+def get_first_file_from_s3_dir(source: str, storage_options: Dict[str, Any] = None) -> str:
+    """
+    Get the first parquet file from an S3 directory path.
+
+    Parameters
+    ----------
+    source : str
+        S3 path with wildcards (e.g., 's3://bucket/prefix/**/*/*.parquet')
+
+    storage_options: FullCloudStorageConnection
+
+    Returns
+    -------
+    str
+        S3 URI of the first parquet file found
+
+    Raises
+    ------
+    ValueError
+        If source path is invalid or no parquet files found
+    ClientError
+        If S3 access fails
+    """
+    if not source.startswith('s3://'):
+        raise ValueError("Source must be a valid S3 URI starting with 's3://'")
+    bucket_name, prefix = _parse_s3_path(source)
+    file_extension = _get_file_extension(source)
+    base_prefix = _remove_wildcards_from_prefix(prefix)
+    s3_client = _create_s3_client(storage_options)
+
+    # Get parquet files
+    first_file = _get_first_file(s3_client, bucket_name, base_prefix, file_extension)
+
+    # Return first file URI
+    return f"s3://{bucket_name}/{first_file['Key']}"
+
+
+def _get_file_extension(source: str) -> str:
+    parts = source.split(".")
+    if len(parts) == 1:
+        raise ValueError("Source path does not contain a file extension")
+    return parts[-1].lower()
+
+
+def _parse_s3_path(source: str) -> tuple[str, str]:
+    """Parse S3 URI into bucket name and prefix."""
+    path_parts = source[5:].split('/', 1)  # Remove 's3://'
+    bucket_name = path_parts[0]
+    prefix = path_parts[1] if len(path_parts) > 1 else ''
+    return bucket_name, prefix
+
+
+def _remove_wildcards_from_prefix(prefix: str) -> str:
+    """Remove wildcard patterns from S3 prefix."""
+    return prefix.split('*')[0]
+
+
+def _create_s3_client(storage_options: Optional[Dict[str, Any]]):
+    """Create boto3 S3 client with optional credentials."""
+    if storage_options is None:
+        return boto3.client('s3')
+
+    # Handle both 'aws_region' and 'region_name' keys
+    client_options = storage_options.copy()
+    if 'aws_region' in client_options:
+        client_options['region_name'] = client_options.pop('aws_region')
+
+    return boto3.client('s3', **{k: v for k, v in client_options.items() if k != "aws_allow_http"})
+
+
+def _get_first_file(s3_client, bucket_name: str, base_prefix: str, file_extension: str) -> Dict[Any, Any]:
+    """List all parquet files in S3 bucket with given prefix."""
+    try:
+        paginator = s3_client.get_paginator('list_objects_v2')
+        pages = paginator.paginate(Bucket=bucket_name, Prefix=base_prefix)
+        for page in pages:
+            if 'Contents' in page:
+                for obj in page['Contents']:
+                    if obj['Key'].endswith(f".{file_extension}"):
+                        return obj
+            else:
+                raise ValueError(f"No objects found in s3://{bucket_name}/{base_prefix}")
+        raise ValueError(f"No {file_extension} files found in s3://{bucket_name}/{base_prefix}")
+    except ClientError as e:
+        raise ValueError(f"Failed to list files in s3://{bucket_name}/{base_prefix}: {e}")
+
+
+def ensure_path_has_wildcard_pattern(resource_path: str, file_format: Literal["csv", "parquet", "json"]):
+    if not resource_path.endswith(f"*.{file_format}"):
+        resource_path = resource_path.rstrip("/") + f"/**/*.{file_format}"
+    return resource_path

flowfile_core/flowfile/flow_data_engine/create/funcs.py

@@ -1,9 +1,9 @@
 import polars as pl
 import os
-
 from flowfile_core.schemas import input_schema
 from flowfile_core.flowfile.flow_data_engine.sample_data import create_fake_data
 from flowfile_core.flowfile.flow_data_engine.read_excel_tables import df_from_openpyxl, df_from_calamine_xlsx
+from polars._typing import CsvEncoding
 
 
 def create_from_json(received_table: input_schema.ReceivedCsvTable):
@@ -49,11 +49,21 @@ def create_from_json(received_table: input_schema.ReceivedCsvTable):
         return data
 
 
-def create_from_path_csv(received_table: input_schema.ReceivedCsvTable) -> pl.DataFrame:
+def standardize_utf8_encoding(non_standardized_encoding: str) -> CsvEncoding:
+    if non_standardized_encoding.upper() in ('UTF-8', 'UTF8'):
+        return 'utf8'
+    elif non_standardized_encoding.upper() in ('UTF-8-LOSSY', 'UTF8-LOSSY'):
+        return 'utf8-lossy'
+    else:
+        raise ValueError(f"Encoding {non_standardized_encoding} is not supported.")
+
+
+def create_from_path_csv(received_table: input_schema.ReceivedCsvTable) -> pl.LazyFrame:
     f = received_table.abs_file_path
     gbs_to_load = os.path.getsize(f) / 1024 / 1000 / 1000
     low_mem = gbs_to_load > 10
-    if received_table.encoding.upper() == 'UTF8' or received_table.encoding.upper() == 'UTF-8':
+    if received_table.encoding.upper() in ("UTF-8", "UTF8", 'UTF8-LOSSY', 'UTF-8-LOSSY'):
+        encoding: CsvEncoding = standardize_utf8_encoding(received_table.encoding)
         try:
             data = pl.scan_csv(f,
                                low_memory=low_mem,
@@ -61,11 +71,12 @@ def create_from_path_csv(received_table: input_schema.ReceivedCsvTable) -> pl.Da
                                separator=received_table.delimiter,
                                has_header=received_table.has_headers,
                                skip_rows=received_table.starting_from_line,
-                               encoding='utf8',
+                               encoding=encoding,
                                infer_schema_length=received_table.infer_schema_length)
             data.head(1).collect()
             return data
         except:
+
             try:
                 data = pl.scan_csv(f, low_memory=low_mem,
                                    separator=received_table.delimiter,
@@ -75,11 +86,11 @@ def create_from_path_csv(received_table: input_schema.ReceivedCsvTable) -> pl.Da
                                    ignore_errors=True)
                 return data
             except:
-                data = pl.scan_csv(f, low_memory=low_mem,
+                data = pl.scan_csv(f, low_memory=False,
                                    separator=received_table.delimiter,
                                    has_header=received_table.has_headers,
                                    skip_rows=received_table.starting_from_line,
-                                   encoding='utf8',
+                                   encoding=encoding,
                                    ignore_errors=True)
                 return data
     else:
@@ -90,14 +101,14 @@ def create_from_path_csv(received_table: input_schema.ReceivedCsvTable) -> pl.Da
                                    skip_rows=received_table.starting_from_line,
                                    encoding=received_table.encoding,
                                    ignore_errors=True, batch_size=2).next_batches(1)
-        return data[0]
+        return data[0].lazy()
 
 
 def create_random(number_of_records: int = 1000) -> pl.LazyFrame:
     return create_fake_data(number_of_records).lazy()
 
 
-def create_from_path_parquet(received_table: input_schema.ReceivedParquetTable):
+def create_from_path_parquet(received_table: input_schema.ReceivedParquetTable) -> pl.LazyFrame:
     low_mem = (os.path.getsize(received_table.abs_file_path) / 1024 / 1000 / 1000) > 2
     return pl.scan_parquet(source=received_table.abs_file_path, low_memory=low_mem)