diff --git a/python-sdk/src/astro/databases/base.py b/python-sdk/src/astro/databases/base.py index 52e53ed9af..5d6d7c6b39 100644 --- a/python-sdk/src/astro/databases/base.py +++ b/python-sdk/src/astro/databases/base.py @@ -19,6 +19,7 @@ from astro.exceptions import DatabaseCustomError, NonExistentTableException from astro.files import File, resolve_file_path_pattern from astro.files.types import create_file_type +from astro.files.types.base import FileType as FileTypeConstants from astro.settings import LOAD_TABLE_AUTODETECT_ROWS_COUNT, SCHEMA from astro.sql.table import BaseTable, Metadata from pandas.io.sql import SQLDatabase @@ -744,11 +745,14 @@ def check_schema_autodetection_is_supported(self, source_file: File) -> bool: source_file.location.location_type ) - source_filetype = create_file_type( - path=source_file.path, filetype=source_file.type.name + source_filetype = ( + source_file + if isinstance(source_file.type, FileTypeConstants) + else create_file_type(path=source_file.path, filetype=source_file.type) # type: ignore ) + is_source_filetype_supported = ( - (source_filetype in filetype_supported.get("filetype")) # type: ignore + (source_filetype.type.name in filetype_supported.get("filetype")) # type: ignore if filetype_supported else None ) diff --git a/python-sdk/tests/databases/test_base_database.py b/python-sdk/tests/databases/test_base_database.py index 899a13b2ae..0ceab38dea 100644 --- a/python-sdk/tests/databases/test_base_database.py +++ b/python-sdk/tests/databases/test_base_database.py @@ -59,6 +59,26 @@ def test_create_table_using_columns_raises_exception(): assert exc_info.match("To use this method, table.columns must be defined") +def test_check_schema_autodetection_is_supported(): + """ + Test the condition native schema autodetection for files and prefixes + """ + db = create_database("gcp_conn") + assert db.check_schema_autodetection_is_supported( + source_file=File(path="gs://bucket/prefix", filetype=FileType.CSV) + ) + + assert db.check_schema_autodetection_is_supported( + source_file=File(path="gs://bucket/prefix/key.csv") + ) + + assert not ( + db.check_schema_autodetection_is_supported( + source_file=File(path="s3://bucket/prefix/key.csv") + ) + ) + + def test_subclass_missing_append_table_raises_exception(): db = DatabaseSubclass(conn_id="fake_conn_id") source_table = Table()