fix(flink): customize the list of base idenitifers

ibis-project · Nov 22, 2023 · 969b9d7 · 969b9d7
1 parent 903cf5b
commit 969b9d7
Show file tree

Hide file tree

Showing 7 changed files with 573 additions and 24 deletions.
diff --git a/ibis/backends/base/sql/registry/helpers.py b/ibis/backends/base/sql/registry/helpers.py
@@ -14,9 +14,11 @@ def format_call(translator, func, *args):
     return "{}({})".format(func, ", ".join(formatted_args))
 
 
-def quote_identifier(name, quotechar="`", force=False):
+def quote_identifier(
+    name, quotechar="`", force=False, base_identifiers=identifiers.base_identifiers
+):
     """Add quotes to the `name` identifier if needed."""
-    if force or name.count(" ") or name in identifiers.base_identifiers:
+    if force or name.count(" ") or name in base_identifiers:
         return f"{quotechar}{name}{quotechar}"
     else:
         return name

diff --git a/ibis/backends/base/sql/registry/identifiers.py b/ibis/backends/base/sql/registry/identifiers.py
@@ -1,21 +1,5 @@
 from __future__ import annotations
 
-# Copyright 2014 Cloudera Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Base identifiers
-
 base_identifiers = [
     "add",
     "aggregate",

diff --git a/ibis/backends/flink/compiler/core.py b/ibis/backends/flink/compiler/core.py
@@ -14,10 +14,14 @@
     TableSetFormatter,
 )
 from ibis.backends.base.sql.registry import quote_identifier
+from ibis.backends.flink import identifiers
 from ibis.backends.flink.translator import FlinkExprTranslator
 
 
 class FlinkTableSetFormatter(TableSetFormatter):
+    def _quote_identifier(self, name):
+        return quote_identifier(name, base_identifiers=identifiers.base_identifiers)
+
     def _format_in_memory_table(self, op):
         names = op.schema.names
         raw_rows = []
@@ -121,7 +125,7 @@ def _tumble_window_params(
         filter(
             None,
             [
-                f"TABLE {quote_identifier(op.table.name)}",
+                f"TABLE {formatter._quote_identifier(op.table.name)}",
                 f"DESCRIPTOR({formatter._translate(op.time_col)})",
                 formatter._translate(op.window_size),
                 formatter._translate(op.offset) if op.offset else None,
@@ -136,7 +140,7 @@ def _hop_window_params(op: ops.HopWindowingTVF, formatter: TableSetFormatter) ->
         filter(
             None,
             [
-                f"TABLE {quote_identifier(op.table.name)}",
+                f"TABLE {formatter._quote_identifier(op.table.name)}",
                 f"DESCRIPTOR({formatter._translate(op.time_col)})",
                 formatter._translate(op.window_slide),
                 formatter._translate(op.window_size),
@@ -154,7 +158,7 @@ def _cumulate_window_params(
         filter(
             None,
             [
-                f"TABLE {quote_identifier(op.table.name)}",
+                f"TABLE {formatter._quote_identifier(op.table.name)}",
                 f"DESCRIPTOR({formatter._translate(op.time_col)})",
                 formatter._translate(op.window_step),
                 formatter._translate(op.window_size),