Skip to content

Commit

Permalink
Fix native snowflake load_file
Browse files Browse the repository at this point in the history
  • Loading branch information
tatiana committed Jul 19, 2022
1 parent 192fdf0 commit d5cd2cd
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions src/astro/databases/snowflake.py
Expand Up @@ -37,6 +37,9 @@
FileType.PARQUET: "MATCH_BY_COLUMN_NAME=CASE_INSENSITIVE",
}

NATIVE_LOAD_SUPPORTED_FILE_TYPES = (FileType.CSV, FileType.NDJSON, FileType.PARQUET)
NATIVE_LOAD_SUPPORTED_FILE_LOCATIONS = (FileLocation.GS, FileLocation.S3)


@dataclass
class SnowflakeStage:
Expand Down Expand Up @@ -292,6 +295,23 @@ def drop_stage(self, stage: SnowflakeStage) -> None:
# Table load methods
# ---------------------------------------------------------

def is_native_load_file_available(
self, source_file: File, target_table: Table
) -> bool:
"""
Check if there is an optimised path for source to destination.
:param source_file: File from which we need to transfer data
:param target_table: Table that needs to be populated with file data
"""
is_file_type_supported = (
source_file.type.name in NATIVE_LOAD_SUPPORTED_FILE_TYPES
)
is_file_location_supported = (
source_file.location.location_type in NATIVE_LOAD_SUPPORTED_FILE_LOCATIONS
)
return is_file_type_supported and is_file_location_supported

def load_file_to_table_natively(
self,
source_file: File,
Expand Down Expand Up @@ -325,6 +345,9 @@ def load_file_to_table_natively(
<https://docs.snowflake.com/en/sql-reference/sql/create-stage.html>`_
"""
import pdb

pdb.set_trace()
native_support_kwargs = native_support_kwargs or {}
storage_integration = native_support_kwargs.get("storage_integration")
stage = self.create_stage(
Expand Down

0 comments on commit d5cd2cd

Please sign in to comment.