feat(oracle): Add support for BULK COLLECT INTO (#4181)

* feat(oracle): Add support for BULK COLLECT INTO * feat(oracle): Add support for BULK COLLECT INTO * feat(oracle): Add support for BULK COLLECT INTO * feat(oracle): Add support for BULK COLLECT INTO * feat(oracle): Add support for BULK COLLECT INTO --------- Co-authored-by: mmoisen <[email protected]>
tobymao · Oct 2, 2024 · 6ba2bb0 · 6ba2bb0
1 parent 1d16ee2
commit 6ba2bb0
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 1 deletion.
diff --git a/sqlglot/dialects/oracle.py b/sqlglot/dialects/oracle.py
@@ -95,6 +95,7 @@ class Tokenizer(tokens.Tokenizer):
             "(+)": TokenType.JOIN_MARKER,
             "BINARY_DOUBLE": TokenType.DOUBLE,
             "BINARY_FLOAT": TokenType.FLOAT,
+            "BULK COLLECT INTO": TokenType.BULK_COLLECT_INTO,
             "COLUMNS": TokenType.COLUMN,
             "MATCH_RECOGNIZE": TokenType.MATCH_RECOGNIZE,
             "MINUS": TokenType.EXCEPT,
@@ -241,6 +242,25 @@ def _parse_json_exists(self) -> exp.JSONExists:
                 on_condition=self._parse_on_condition(),
             )
 
+        def _parse_into(self) -> t.Optional[exp.Into]:
+            # https://docs.oracle.com/en/database/oracle/oracle-database/19/lnpls/SELECT-INTO-statement.html
+            bulk_collect_into = self._match(TokenType.BULK_COLLECT_INTO)
+
+            if not bulk_collect_into and not self._match(TokenType.INTO):
+                return None
+
+            expressions = self._parse_expressions()
+            this = None
+            if len(expressions) == 1:
+                self._advance(-1)
+                expressions = []
+                self._match(TokenType.TABLE)
+                this = self._parse_table(schema=True)
+
+            return self.expression(
+                exp.Into, this=this, bulk_collect=bulk_collect_into, expressions=expressions
+            )
+
     class Generator(generator.Generator):
         LOCKING_READS_SUPPORTED = True
         JOIN_HINTS = False
@@ -344,3 +364,10 @@ def queryoption_sql(self, expression: exp.QueryOption) -> str:
         def coalesce_sql(self, expression: exp.Coalesce) -> str:
             func_name = "NVL" if expression.args.get("is_nvl") else "COALESCE"
             return rename_func(func_name)(self, expression)
+
+        def into_sql(self, expression: exp.Into) -> str:
+            into = "INTO" if not expression.args.get("bulk_collect") else "BULK COLLECT INTO"
+            if expression.this:
+                return f"{self.seg(into)} {self.sql(expression, 'this')}"
+
+            return f"{self.seg(into)} {self.expressions(expression)}"
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -2106,7 +2106,13 @@ class PrimaryKey(Expression):
 # https://www.postgresql.org/docs/9.1/sql-selectinto.html
 # https://docs.aws.amazon.com/redshift/latest/dg/r_SELECT_INTO.html#r_SELECT_INTO-examples
 class Into(Expression):
-    arg_types = {"this": True, "temporary": False, "unlogged": False}
+    arg_types = {
+        "this": False,
+        "temporary": False,
+        "unlogged": False,
+        "bulk_collect": False,
+        "expressions": False,
+    }
 
 
 class From(Expression):

diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
@@ -221,6 +221,7 @@ class TokenType(AutoName):
     AUTO_INCREMENT = auto()
     BEGIN = auto()
     BETWEEN = auto()
+    BULK_COLLECT_INTO = auto()
     CACHE = auto()
     CASE = auto()
     CHARACTER_SET = auto()

diff --git a/tests/dialects/test_oracle.py b/tests/dialects/test_oracle.py
@@ -66,6 +66,18 @@ def test_oracle(self):
         self.validate_identity(
             "SELECT COUNT(1) INTO V_Temp FROM TABLE(CAST(somelist AS data_list)) WHERE col LIKE '%contact'"
         )
+        self.validate_identity(
+            "SELECT department_id INTO v_department_id FROM departments FETCH FIRST 1 ROWS ONLY"
+        )
+        self.validate_identity(
+            "SELECT department_id, department_name INTO v_department_id, v_department_name FROM departments FETCH FIRST 1 ROWS ONLY"
+        )
+        self.validate_identity(
+            "SELECT department_id BULK COLLECT INTO v_department_ids FROM departments"
+        )
+        self.validate_identity(
+            "SELECT department_id, department_name BULK COLLECT INTO v_department_ids, v_department_names FROM departments"
+        )
         self.validate_identity(
             "SELECT MIN(column_name) KEEP (DENSE_RANK FIRST ORDER BY column_name DESC) FROM table_name"
         )