mage-ai 0.9.65__py3-none-any.whl → 0.9.67__py3-none-any.whl

mage_ai/data_preparation/models/block/dbt/block_sql.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pandas as pd
+import simplejson
 from jinja2 import Template
 
 from mage_ai.data_preparation.models.block import Block
@@ -27,6 +28,8 @@ from mage_ai.settings.utils import base_repo_path
 from mage_ai.shared.files import find_file_from_another_file_path
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.path_fixer import add_absolute_path
+from mage_ai.shared.strings import remove_extension_from_filename
+from mage_ai.shared.utils import clean_name
 
 
 class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
@@ -136,6 +139,167 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
 
         return search_paths.get(node_path_type_part)
 
+    def upstream_dbt_blocks(self, read_only=False) -> List['DBTBlockSQL']:
+        """
+        Get an up to date list, which represents the upstream dbt graph.
+        It is using `dbt list` to generate the list.
+        Args:
+            read_only (bool):
+                If True it does not read the Blocks from the model. Defaults to False
+        Returns:
+            List[DBTBlockSQL]: THe upstream dbt graph as DBTBlocksSQL objects
+        """
+        # Get upstream nodes via dbt list
+        with Profiles(
+            self.project_path,
+            self.pipeline.variables if self.pipeline else {},
+        ) as profiles:
+            try:
+                args = [
+                    'list',
+                    # project-dir
+                    # /home/src/default_repo/default_platform/default_repo/dbt/demo
+                    '--project-dir', self.project_path,
+                    '--profiles-dir', str(profiles.profiles_dir),
+                    '--select', '+' + Path(self.file_path).stem,
+                    '--output', 'json',
+                    '--output-keys', 'unique_id original_file_path depends_on',
+                    '--resource-type', 'model',
+                    '--resource-type', 'snapshot'
+                ]
+                res = DBTCli().invoke(args)
+
+                if res:
+                    nodes_init = [simplejson.loads(node) for node in res.result]
+                else:
+                    return []
+            except Exception as err:
+                print(f'[ERROR] DBTBlockSQL.upstream_dbt_blocks: {err}.')
+                return [
+                    self.build_dbt_block(
+                        block_class=DBTBlock,
+                        block_dict=dict(
+                            block_type=self.type,
+                            configuration=self.configuration,
+                            language=self.language,
+                            name=self.uuid,
+                            pipeline=self.pipeline,
+                            uuid=self.uuid,
+                        ),
+                        hydrate_configuration=False,
+                    )
+                ]
+
+        # transform List into dict and remove unnecessary fields
+        file_path = self.file_path
+        path_parts = file_path.split(os.sep)
+        project_dir = path_parts[0]
+
+        nodes_default = {
+            node['unique_id']: {
+                # file_path:
+                # default_repo/dbt/demo/models/example/model.sql
+                # default_platform/default_repo/dbt/demo/models/example/model.sql
+                'file_path': os.path.join(project_dir, node['original_file_path']),
+                'original_file_path': node['original_file_path'],
+                'upstream_nodes': set(node['depends_on']['nodes'])
+            }
+            for node in nodes_init
+        }
+
+        nodes = self.hydrate_dbt_nodes(nodes_default, nodes_init)
+
+        # calculate downstream_nodes
+        for unique_id, node in nodes.items():
+            for upstream_node in node['upstream_nodes']:
+                if nodes.get(upstream_node):
+                    downstream_nodes = nodes[upstream_node].get('downstream_nodes', set())
+                    downstream_nodes.add(unique_id)
+                    nodes[upstream_node]['downstream_nodes'] = downstream_nodes
+
+        # map dbt unique_id to mage uuid
+        uuids_default = {
+            unique_id: clean_name(
+                remove_extension_from_filename(node['file_path']),
+                allow_characters=[os.sep]
+            )
+            for unique_id, node in nodes.items()
+        }
+        uuids = self.node_uuids_mapping(uuids_default, nodes)
+
+        # replace dbt unique_ids with mage uuids
+        nodes = {
+            uuids[unique_id]: {
+                'file_path': node['file_path'],
+                'original_file_path': node['original_file_path'],
+                'upstream_nodes': {
+                    uuids[upstream_node]
+                    for upstream_node in node.get('upstream_nodes', set())
+                    if uuids.get(upstream_node)
+                },
+                'downstream_nodes': {
+                    uuids[downstream_node]
+                    for downstream_node in node.get('downstream_nodes', set())
+                    if uuids.get(downstream_node)
+                }
+            }
+            for unique_id, node in nodes.items()
+        }
+
+        # create dict of blocks
+        blocks = {}
+        for uuid, node in nodes.items():
+            block = None
+            if not read_only:
+                if uuid == self.uuid:
+                    block = self
+                elif self.pipeline:
+                    block = self.pipeline.get_block(
+                        uuid,
+                        self.type,
+                    )
+            # if not found create the block
+            block = block or self.build_dbt_block(
+                block_class=DBTBlock,
+                block_dict=dict(
+                    block_type=self.type,
+                    configuration=dict(file_path=node['file_path']),
+                    language=self.language,
+                    name=uuid,
+                    pipeline=self.pipeline,
+                    uuid=uuid,
+                ),
+                node=node,
+            )
+            # reset upstream dbt blocks
+            block.upstream_blocks = [
+                block
+                for block in block.upstream_blocks
+                # there must not be any other upstream dbt blocks that are not part of
+                # the upstream nodes, therefore delete all upstream dbt blocks
+                if not isinstance(block, DBTBlockSQL)
+            ]
+            # reset downstream dbt blocks
+            block.downstream_blocks = [
+                block
+                for block in block.downstream_blocks
+                # remove all downstream dbt block, which are part of the upstream nodes,
+                # in order to fix errors
+                if block.uuid not in nodes.keys()
+            ]
+            blocks[uuid] = block
+
+        # Update upstream and downstream dbt blocks
+        for uuid, node in nodes.items():
+            for upstream_node in node.get('upstream_nodes', set()):
+                blocks[uuid].upstream_blocks.append(blocks[upstream_node])
+            for downstream_node in node.get('downstream_nodes', set()):
+                blocks[uuid].downstream_blocks.append(blocks[downstream_node])
+
+        # transform into list
+        blocks = [block for _, block in blocks.items()]
+        return blocks
+
     def _execute_block(
         self,
         outputs_from_input_vars,

mage_ai/data_preparation/models/block/dynamic/variables.py

@@ -522,9 +522,8 @@ def fetch_input_variables_for_dynamic_upstream_blocks(
                 global_vars=global_vars,
                 **kwargs,
             )
-
             input_vars.append(ir[0])
-            kwargs_vars.append(kr[0])
+            kwargs_vars.append(kr[0] if len(kr) > 0 else None)
             upstream_block_uuids.append(up[0])
 
     return input_vars, kwargs_vars, upstream_block_uuids

mage_ai/data_preparation/models/pipeline.py

@@ -1443,7 +1443,7 @@ class Pipeline:
                 block['downstream_blocks'] = hierarchy_list[b_index]['downstream_blocks']
 
         # dump new config information back in temp folder
-        with open(config_zip_path, 'w') as pipeline_config:
+        with open(config_zip_path, 'w', encoding='utf-8') as pipeline_config:
             yaml.dump(config, pipeline_config)
 
         return files_to_be_written, config
@@ -2141,12 +2141,12 @@ class Pipeline:
         content = yaml.dump(pipeline_dict)
 
         test_path = f'{self.config_path}.test'
-        async with aiofiles.open(test_path, mode='w') as fp:
+        async with aiofiles.open(test_path, mode='w', encoding='utf-8') as fp:
             await fp.write(content)
 
         if os.path.isfile(test_path):
             success = True
-            with open(test_path, mode='r') as fp:
+            with open(test_path, mode='r', encoding='utf-8') as fp:
                 try:
                     yaml.full_load(fp)
                 except yaml.scanner.ScannerError:

mage_ai/data_preparation/models/triggers/__init__.py

@@ -69,6 +69,7 @@ class Trigger(BaseConfig):
     sla: int = None     # in seconds
     settings: Dict = field(default_factory=dict)
     envs: List = field(default_factory=list)
+    repo_path: str = None
 
     def __post_init__(self):
         if self.schedule_type and type(self.schedule_type) is str:
@@ -163,6 +164,7 @@ def build_triggers(
     trigger_configs: Dict,
     pipeline_uuid: str = None,
     raise_exception: bool = False,
+    repo_path: str = None,
 ) -> List[Trigger]:
     triggers = []
     for trigger_config in trigger_configs:
@@ -170,6 +172,8 @@ def build_triggers(
             trigger_config['pipeline_uuid'] = pipeline_uuid
         try:
             trigger = Trigger.load(config=trigger_config)
+            if trigger.repo_path is None:
+                trigger.repo_path = repo_path
 
             # Add flag to settings so frontend can detect triggers with invalid cron expressions
             if not trigger.has_valid_schedule_interval:
@@ -194,7 +198,8 @@ def load_trigger_configs(
     yaml_config = yaml.safe_load(content) or {}
     trigger_configs = yaml_config.get('triggers') or {}
 
-    return build_triggers(trigger_configs, pipeline_uuid, raise_exception)
+    return build_triggers(
+        trigger_configs, pipeline_uuid, raise_exception, repo_path=get_repo_path())
 
 
 def add_or_update_trigger_for_pipeline_and_persist(

mage_ai/data_preparation/preferences.py

@@ -6,51 +6,49 @@ import yaml
 
 from mage_ai.data_preparation.models.constants import PREFERENCES_FILE
 from mage_ai.orchestration.db.models.oauth import User
-from mage_ai.settings import get_bool_value
+from mage_ai.settings import get_bool_value, get_settings_value
+from mage_ai.settings.keys import (
+    GIT_AUTH_TYPE,
+    GIT_BRANCH,
+    GIT_EMAIL,
+    GIT_ENABLE_GIT_INTEGRATION,
+    GIT_OVERWRITE_WITH_PROJECT_SETTINGS,
+    GIT_REPO_LINK,
+    GIT_REPO_PATH,
+    GIT_SYNC_ON_EXECUTOR_START,
+    GIT_SYNC_ON_PIPELINE_RUN,
+    GIT_SYNC_ON_START,
+    GIT_SYNC_SUBMODULES,
+    GIT_USERNAME,
+)
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.hash import merge_dict
 
-# Git environment variables
-GIT_REPO_LINK_VAR = 'GIT_REPO_LINK'
-GIT_REPO_PATH_VAR = 'GIT_REPO_PATH'
-GIT_USERNAME_VAR = 'GIT_USERNAME'
-GIT_EMAIL_VAR = 'GIT_EMAIL'
-GIT_AUTH_TYPE_VAR = 'GIT_AUTH_TYPE'
-GIT_BRANCH_VAR = 'GIT_BRANCH'
-GIT_SYNC_ON_PIPELINE_RUN_VAR = 'GIT_SYNC_ON_PIPELINE_RUN'
-GIT_SYNC_ON_START_VAR = 'GIT_SYNC_ON_START'
-GIT_SYNC_ON_EXECUTOR_START_VAR = 'GIT_SYNC_ON_EXECUTOR_START'
-GIT_SYNC_SUBMODULES_VAR = 'GIT_SYNC_SUBMODULES'
-
-GIT_ENABLE_GIT_INTEGRATION_VAR = 'GIT_ENABLE_GIT_INTEGRATION'
-GIT_OVERWRITE_WITH_PROJECT_SETTINGS_VAR = 'GIT_OVERWRITE_WITH_PROJECT_SETTINGS'
-
 ENV_VAR_TO_CONFIG_KEY = {
-    GIT_REPO_LINK_VAR: 'remote_repo_link',
-    GIT_REPO_PATH_VAR: 'repo_path',
-    GIT_USERNAME_VAR: 'username',
-    GIT_EMAIL_VAR: 'email',
-    GIT_AUTH_TYPE_VAR: 'auth_type',
-    GIT_BRANCH_VAR: 'branch',
-    GIT_SYNC_ON_PIPELINE_RUN_VAR: 'sync_on_pipeline_run',
-    GIT_SYNC_ON_START_VAR: 'sync_on_start',
-    GIT_SYNC_ON_EXECUTOR_START_VAR: 'sync_on_executor_start',
-    GIT_SYNC_SUBMODULES_VAR: 'sync_submodules',
+    GIT_REPO_LINK: 'remote_repo_link',
+    GIT_REPO_PATH: 'repo_path',
+    GIT_USERNAME: 'username',
+    GIT_EMAIL: 'email',
+    GIT_AUTH_TYPE: 'auth_type',
+    GIT_BRANCH: 'branch',
+    GIT_SYNC_ON_PIPELINE_RUN: 'sync_on_pipeline_run',
+    GIT_SYNC_ON_START: 'sync_on_start',
+    GIT_SYNC_ON_EXECUTOR_START: 'sync_on_executor_start',
+    GIT_SYNC_SUBMODULES: 'sync_submodules',
 }
 
 BOOLEAN_ENV_VARS = set(
     [
-        GIT_SYNC_ON_PIPELINE_RUN_VAR,
-        GIT_SYNC_ON_START_VAR,
-        GIT_SYNC_ON_EXECUTOR_START_VAR,
-        GIT_SYNC_SUBMODULES_VAR,
-        GIT_ENABLE_GIT_INTEGRATION_VAR,
+        GIT_SYNC_ON_PIPELINE_RUN,
+        GIT_SYNC_ON_START,
+        GIT_SYNC_ON_EXECUTOR_START,
+        GIT_SYNC_SUBMODULES,
+        GIT_ENABLE_GIT_INTEGRATION,
     ]
 )
 
 
-def get_value_for_sync_config(env_var) -> bool:
-    value = os.getenv(env_var)
+def get_value_for_sync_config(env_var, value) -> bool:
     return get_bool_value(value) if env_var in BOOLEAN_ENV_VARS else value
 
 
@@ -68,21 +66,28 @@ def build_sync_config(project_sync_config: Dict) -> Dict:
         if v is not None and (not isinstance(v, str) or v)
     }
 
-    if any([k in os.environ for k in ENV_VAR_TO_CONFIG_KEY]):
+    # Read settings from settings backend
+    config_from_settings = {
+        k: get_settings_value(k)
+        for k in ENV_VAR_TO_CONFIG_KEY.keys()
+    }
+
+    if any([v is not None for v in config_from_settings.values()]):
+        # Map sync config keys to settings values
         sync_config = {
-            v: get_value_for_sync_config(k)
+            v: get_value_for_sync_config(k, config_from_settings.get(k))
             for k, v in ENV_VAR_TO_CONFIG_KEY.items()
         }
 
         # If the GIT_OVERWRITE_WITH_PROJECT_SETTINGS environment variable is set to true,
         # overwrite the sync_config with the config from the project preferences file.
-        if get_bool_value(os.getenv(GIT_OVERWRITE_WITH_PROJECT_SETTINGS_VAR)):
+        if get_bool_value(get_settings_value(GIT_OVERWRITE_WITH_PROJECT_SETTINGS)):
             sync_config = merge_dict(sync_config, project_sync_config)
     else:
         sync_config = project_sync_config
 
     # Set the enable_git_integration field from environment variables
-    enable_git_integration = os.getenv(GIT_ENABLE_GIT_INTEGRATION_VAR)
+    enable_git_integration = get_settings_value(GIT_ENABLE_GIT_INTEGRATION)
     if enable_git_integration is not None:
         sync_config['enable_git_integration'] = get_bool_value(enable_git_integration)
 

mage_ai/data_preparation/repo_manager.py

@@ -1,3 +1,5 @@
+import json
+import logging
 import os
 import traceback
 import uuid
@@ -11,6 +13,7 @@ from jinja2 import Template
 
 from mage_ai.cluster_manager.constants import ClusterType
 from mage_ai.data_preparation.templates.utils import copy_template_directory
+from mage_ai.settings import INITIAL_METADATA, settings
 from mage_ai.settings.repo import DEFAULT_MAGE_DATA_DIR, MAGE_DATA_DIR_ENV_VAR
 from mage_ai.settings.repo import get_data_dir as get_data_dir_new
 from mage_ai.settings.repo import get_metadata_path
@@ -25,6 +28,8 @@ yml.preserve_quotes = True
 yml.width = float('inf')
 yml.indent(mapping=2, sequence=2, offset=0)
 
+logger = logging.getLogger(__name__)
+
 
 class ProjectType(str, Enum):
     MAIN = 'main'
@@ -64,9 +69,12 @@ class RepoConfig:
         self.ldap_config = None
         self.s3_bucket = None
         self.s3_path_prefix = None
+        self.settings_backend = None
         self.logging_config = None
         self.variables_dir = None
         self.variables_retention_period = None
+        self.workspace_initial_metadata = None
+        self.workspace_shared_config = None
 
         from mage_ai.data_preparation.shared.utils import get_template_vars
         try:
@@ -125,6 +133,7 @@ class RepoConfig:
             self.pipelines = repo_config.get('pipelines')
             self.retry_config = repo_config.get('retry_config')
             self.workspace_config_defaults = repo_config.get('workspace_config_defaults')
+            self.workspace_initial_metadata = repo_config.get('workspace_initial_metadata')
 
             self.ldap_config = repo_config.get('ldap_config')
 
@@ -136,6 +145,8 @@ class RepoConfig:
                 self.s3_bucket = path_parts.pop(0)
                 self.s3_path_prefix = '/'.join(path_parts)
 
+            self.settings_backend = repo_config.get('settings_backend', dict())
+
             self.logging_config = repo_config.get('logging_config', dict())
 
             self.variables_retention_period = repo_config.get('variables_retention_period')
@@ -233,6 +244,11 @@ def init_repo(
         raise FileExistsError(f'Repository {repo_path} already exists')
 
     new_config = dict()
+    if INITIAL_METADATA:
+        try:
+            new_config = json.loads(INITIAL_METADATA)
+        except Exception:
+            logger.exception('Error loading initial metadata.')
     if project_type == ProjectType.MAIN:
         copy_template_directory('main', repo_path)
         new_config.update(
@@ -300,6 +316,11 @@ def set_project_uuid_from_metadata() -> None:
             project_uuid = config.get('project_uuid')
 
 
+def update_settings_on_metadata_change() -> None:
+    repo_config = get_repo_config(root_project=True)
+    settings.set_settings_backend(**repo_config.settings_backend)
+
+
 def init_project_uuid(overwrite_uuid: str = None, root_project: bool = False) -> None:
     """
     Initialize the project_uuid constant. The project_uuid constant is used throughout

mage_ai/data_preparation/storage/gcs_storage.py

@@ -29,11 +29,36 @@ class GCSStorage(BaseStorage):
             path += '/'
         return self.path_exists(path)
 
-    def listdir(self, path: str, suffix: str = None) -> List[str]:
+    def listdir(
+        self,
+        path: str,
+        suffix: str = None,
+        max_results: int = None,
+    ) -> List[str]:
         if not path.endswith('/'):
             path += '/'
         path = gcs_url_path(path)
-        blobs = self.bucket.list_blobs(prefix=path)
+
+        """
+        NOTE: The number of results returned by the GCS bucket list_blobs method may
+        be lower than the number entered for the max_results argument, since this
+        method recursively fetches files inside of subdirectories at the path
+        specified, not just the directories themselves at the top level of the path.
+
+        As a result, when fetching block output folders in a remote variables directory,
+        the json or parquet files inside of the block output folders (e.g. the "type.json"
+        file in folder "output_0") are counted towards the "max_results" limit. If max_results
+        is set to 10, less than 10 output folders may be returned due to there being multiple
+        json or other files in each output folder. We do not increase the max_results or
+        include a "delimiter" argument because it would increase load times to unacceptable levels.
+
+        Example block output path in GCS:
+            gs://mage_demo/pipelines/example_pipeline/.variables/example_block/output_0/
+        """
+        blobs = self.bucket.list_blobs(
+            prefix=path,
+            max_results=max_results,
+        )
         keys = []
         for blob in blobs:
             # Avoid finding files recursevively in the dir path.

mage_ai/data_preparation/storage/local_storage.py

@@ -19,12 +19,27 @@ class LocalStorage(BaseStorage):
     def isdir(self, path: str) -> bool:
         return os.path.isdir(path)
 
-    def listdir(self, path: str, suffix: str = None) -> List[str]:
+    def listdir(
+        self,
+        path: str,
+        suffix: str = None,
+        max_results: int = None,
+    ) -> List[str]:
+        paths = []
         if not os.path.exists(path):
-            return []
-        paths = os.listdir(path)
+            return paths
+
+        if max_results is not None:
+            with os.scandir(path) as it:
+                for idx, entry in enumerate(it):
+                    paths.append(entry.name)
+                    if idx >= max_results - 1:
+                        break
+        else:
+            paths = os.listdir(path)
         if suffix is not None:
             paths = [p for p in paths if p.endswith(suffix)]
+
         return paths
 
     def makedirs(self, path: str, **kwargs) -> None:

mage_ai/data_preparation/storage/s3_storage.py

@@ -30,11 +30,16 @@ class S3Storage(BaseStorage):
             path += '/'
         return self.path_exists(path)
 
-    def listdir(self, path: str, suffix: str = None) -> List[str]:
+    def listdir(
+        self,
+        path: str,
+        suffix: str = None,
+        max_results: int = None,
+    ) -> List[str]:
         if not path.endswith('/'):
             path += '/'
         path = s3_url_path(path)
-        keys = self.client.listdir(path, suffix=suffix)
+        keys = self.client.listdir(path, suffix=suffix, max_results=max_results)
         return [k[len(path):].rstrip('/') for k in keys]
 
     def makedirs(self, path: str, **kwargs) -> None:

mage_ai/data_preparation/templates/data_loaders/streaming/generic_python.py

@@ -0,0 +1,23 @@
+from mage_ai.streaming.sources.base_python import BasePythonSource
+from typing import Callable
+
+if 'streaming_source' not in globals():
+    from mage_ai.data_preparation.decorators import streaming_source
+
+
+@streaming_source
+class CustomSource(BasePythonSource):
+    def init_client(self):
+        """
+        Implement the logic of initializing the client.
+        """
+
+    def batch_read(self, handler: Callable):
+        """
+        Batch read the messages from the source and use handler to process the messages.
+        """
+        while True:
+            records = []
+            # Implement the logic of fetching the records
+            if len(records) > 0:
+                handler(records)

mage_ai/data_preparation/templates/main/metadata.yaml

@@ -3,6 +3,12 @@ project_type: main
 variables_dir: ~/.mage_data
 # remote_variables_dir: s3://bucket/path_prefix
 
+# Project metadata that will be prepopulated for workspaces created from this project.
+workspace_initial_metadata: {}
+
+# Default workspace configuration. 
+workspace_default_config: {}
+
 variables_retention_period: '90d'
 
 emr_config:

mage_ai/data_preparation/templates/template.py

@@ -140,19 +140,23 @@ def __fetch_data_loader_templates(
     pipeline_type: PipelineType = PipelineType.PYTHON,
 ) -> str:
     data_source = config.get('data_source')
+    default_template_name = None
     file_extension = 'py'
     if pipeline_type == PipelineType.PYSPARK:
         template_folder = 'data_loaders/pyspark'
     elif pipeline_type == PipelineType.STREAMING:
         template_folder = 'data_loaders/streaming'
-        file_extension = 'yaml'
+        if language == BlockLanguage.YAML:
+            file_extension = 'yaml'
+        else:
+            default_template_name = 'generic_python.py'
     elif language == BlockLanguage.R:
         template_folder = 'data_loaders/r'
         file_extension = 'r'
     else:
         template_folder = 'data_loaders'
 
-    default_template = template_folder + '/' + 'default.jinja'
+    default_template = template_folder + '/' + (default_template_name or 'default.jinja')
     if data_source is None:
         template_path = default_template
     else:

mage_ai/data_preparation/variable_manager.py

@@ -271,6 +271,7 @@ class VariableManager:
         block_uuid: str,
         partition: str = None,
         clean_block_uuid: bool = True,
+        max_results: int = None,
     ) -> List[str]:
         variable_dir_path = os.path.join(
             self.pipeline_path(pipeline_uuid),
@@ -280,7 +281,7 @@ class VariableManager:
         )
         if not self.storage.path_exists(variable_dir_path):
             return []
-        variables = self.storage.listdir(variable_dir_path)
+        variables = self.storage.listdir(variable_dir_path, max_results=max_results)
         return sorted([v.split('.')[0] for v in variables])
 
     def pipeline_path(self, pipeline_uuid: str) -> str:

mage_ai/io/base.py

@@ -302,8 +302,11 @@ class BaseSQLDatabase(BaseIO):
         self,
         column_name: str,
         allow_reserved_words: bool = False,
+        auto_clean_name: bool = True,
         case_sensitive: bool = False,
     ) -> str:
+        if not auto_clean_name:
+            return False
         col_new = clean_name(column_name, case_sensitive=case_sensitive)
         if not allow_reserved_words and col_new.upper() in SQL_RESERVED_WORDS:
             col_new = f'_{col_new}'

mage_ai/io/bigquery.py

@@ -1,5 +1,6 @@
 from typing import Dict, List, Mapping, Union
 
+import numpy as np
 from google.cloud.bigquery import (
     Client,
     LoadJobConfig,
@@ -295,6 +296,7 @@ WHERE table_id = '{table_name}'
                 column_types = self.get_column_types(schema, table_name)
 
                 if df is not None:
+                    df.fillna(value=np.NaN, inplace=True)
                     for col in df.columns:
                         col_type = column_types.get(col)
                         if not col_type: