diff --git a/base.py b/base.py index 75ac272f..8b66aaf3 100644 --- a/base.py +++ b/base.py @@ -418,14 +418,16 @@ def visit_create_file_format(self, file_format, **kw): This visitor will create the SQL representation for a CREATE FILE FORMAT command. """ - return "CREATE {}FILE FORMAT {} TYPE='{}' {}".format( - "OR REPLACE " if file_format.replace_if_exists else "", + return "CREATE {}FILE FORMAT {}{} TYPE='{}' {}{}".format( + "OR REPLACE " if file_format.or_replace else "", + "IF NOT EXISTS " if file_format.if_not_exists else "", file_format.format_name, file_format.formatter.file_format, " ".join( ["{} = {}".format(name, file_format.formatter.value_repr(name, value)) for name, value - in file_format.formatter.options.items()]) + in file_format.formatter.options.items()]), + " COMMENT = '{}'".format(file_format.comment) if file_format.comment else "" ) def visit_drop_table_comment(self, drop, **kw): diff --git a/custom_commands.py b/custom_commands.py index 2b2f510c..dbf17f17 100644 --- a/custom_commands.py +++ b/custom_commands.py @@ -434,11 +434,17 @@ class CreateFileFormat(DDLElement): """ __visit_name__ = "create_file_format" - def __init__(self, format_name, formatter, replace_if_exists=False): + def __init__(self, format_name, formatter, replace_if_exists=False, if_not_exists=False, comment=None): super().__init__() self.format_name = format_name self.formatter = formatter - self.replace_if_exists = replace_if_exists + self.or_replace = replace_if_exists + self.if_not_exists = if_not_exists + self.comment = comment + + if self.or_replace and self.if_not_exists: + raise ValueError("Can't replace and create if not exists in the same time {} file format.".format( + format_name)) class CreateStage(DDLElement): diff --git a/test/test_create.py b/test/test_create.py index 2fda4f0a..978cd1fa 100644 --- a/test/test_create.py +++ b/test/test_create.py @@ -77,6 +77,27 @@ def test_create_csv_format(sql_compiler): "TYPE='csv' FIELD_DELIMITER = ','" assert actual == expected + create_format_if_not_exists = CreateFileFormat( + format_name="ML_POC.PUBLIC.CSV_FILE_FORMAT", + formatter=CSVFormatter().field_delimiter(","), + if_not_exists=True + ) + actual = sql_compiler(create_format_if_not_exists) + expected = "CREATE FILE FORMAT IF NOT EXISTS ML_POC.PUBLIC.CSV_FILE_FORMAT " \ + "TYPE='csv' FIELD_DELIMITER = ','" + assert actual == expected + + create_format_with_comment = CreateFileFormat( + format_name="ML_POC.PUBLIC.CSV_FILE_FORMAT", + formatter=CSVFormatter().field_delimiter(","), + replace_if_exists=True, + comment='my nice file format comment' + ) + actual = sql_compiler(create_format_with_comment) + expected = "CREATE OR REPLACE FILE FORMAT ML_POC.PUBLIC.CSV_FILE_FORMAT " \ + "TYPE='csv' FIELD_DELIMITER = ',' COMMENT = 'my nice file format comment'" + assert actual == expected + def test_create_parquet_format(sql_compiler): """