feat(api): add to_sqlglot method to Schema objects (#10063)

ibis-project · Sep 9, 2024 · 9488115 · 9488115
1 parent dfa55b6
commit 9488115
Show file tree

Hide file tree

Showing 17 changed files with 180 additions and 276 deletions.
diff --git a/ibis/backends/bigquery/__init__.py b/ibis/backends/bigquery/__init__.py
@@ -34,7 +34,6 @@
 )
 from ibis.backends.bigquery.datatypes import BigQuerySchema
 from ibis.backends.sql import SQLBackend
-from ibis.backends.sql.datatypes import BigQueryType
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping
@@ -1060,22 +1059,12 @@ def create_table(
 
         table = _force_quote_table(table)
 
-        column_defs = [
-            sge.ColumnDef(
-                this=sg.to_identifier(name, quoted=self.compiler.quoted),
-                kind=BigQueryType.from_ibis(typ),
-                constraints=(
-                    None
-                    if typ.nullable or typ.is_array()
-                    else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
-                ),
-            )
-            for name, typ in (schema or {}).items()
-        ]
-
         stmt = sge.Create(
             kind="TABLE",
-            this=sge.Schema(this=table, expressions=column_defs or None),
+            this=sge.Schema(
+                this=table,
+                expressions=schema.to_sqlglot(self.dialect) if schema else None,
+            ),
             replace=overwrite,
             properties=sge.Properties(expressions=properties),
             expression=None if obj is None else self.compile(obj),

diff --git a/ibis/backends/clickhouse/__init__.py b/ibis/backends/clickhouse/__init__.py
@@ -674,13 +674,7 @@ def create_table(
 
         this = sge.Schema(
             this=sg.table(name, db=database, quoted=self.compiler.quoted),
-            expressions=[
-                sge.ColumnDef(
-                    this=sg.to_identifier(name, quoted=self.compiler.quoted),
-                    kind=self.compiler.type_mapper.from_ibis(typ),
-                )
-                for name, typ in (schema or obj.schema()).items()
-            ],
+            expressions=(schema or obj.schema()).to_sqlglot(self.dialect),
         )
         properties = [
             # the engine cannot be quoted, since clickhouse won't allow e.g.,

diff --git a/ibis/backends/datafusion/__init__.py b/ibis/backends/datafusion/__init__.py
@@ -657,20 +657,10 @@ def create_table(
         table_ident = sg.table(name, db=database, quoted=quoted)
 
         if query is None:
-            column_defs = [
-                sge.ColumnDef(
-                    this=sg.to_identifier(colname, quoted=quoted),
-                    kind=self.compiler.type_mapper.from_ibis(typ),
-                    constraints=(
-                        None
-                        if typ.nullable
-                        else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
-                    ),
-                )
-                for colname, typ in (schema or table.schema()).items()
-            ]
-
-            target = sge.Schema(this=table_ident, expressions=column_defs)
+            target = sge.Schema(
+                this=table_ident,
+                expressions=(schema or table.schema()).to_sqlglot(self.dialect),
+            )
         else:
             target = table_ident
 

diff --git a/ibis/backends/duckdb/__init__.py b/ibis/backends/duckdb/__init__.py
@@ -175,19 +175,6 @@ def create_table(
         else:
             query = None
 
-        column_defs = [
-            sge.ColumnDef(
-                this=sg.to_identifier(colname, quoted=self.compiler.quoted),
-                kind=self.compiler.type_mapper.from_ibis(typ),
-                constraints=(
-                    None
-                    if typ.nullable
-                    else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
-                ),
-            )
-            for colname, typ in (schema or table.schema()).items()
-        ]
-
         if overwrite:
             temp_name = util.gen_name("duckdb_table")
         else:
@@ -196,7 +183,10 @@ def create_table(
         initial_table = sg.table(
             temp_name, catalog=catalog, db=database, quoted=self.compiler.quoted
         )
-        target = sge.Schema(this=initial_table, expressions=column_defs)
+        target = sge.Schema(
+            this=initial_table,
+            expressions=(schema or table.schema()).to_sqlglot(self.dialect),
+        )
 
         create_stmt = sge.Create(
             kind="TABLE",

diff --git a/ibis/backends/exasol/__init__.py b/ibis/backends/exasol/__init__.py
@@ -256,27 +256,13 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
         # only register if we haven't already done so
         if (name := op.name) not in self.list_tables():
             quoted = self.compiler.quoted
-            column_defs = [
-                sg.exp.ColumnDef(
-                    this=sg.to_identifier(colname, quoted=quoted),
-                    kind=self.compiler.type_mapper.from_ibis(typ),
-                    constraints=(
-                        None
-                        if typ.nullable
-                        else [
-                            sg.exp.ColumnConstraint(
-                                kind=sg.exp.NotNullColumnConstraint()
-                            )
-                        ]
-                    ),
-                )
-                for colname, typ in schema.items()
-            ]
 
             ident = sg.to_identifier(name, quoted=quoted)
             create_stmt = sg.exp.Create(
                 kind="TABLE",
-                this=sg.exp.Schema(this=ident, expressions=column_defs),
+                this=sg.exp.Schema(
+                    this=ident, expressions=schema.to_sqlglot(self.dialect)
+                ),
             )
             create_stmt_sql = create_stmt.sql(self.name)
 
@@ -366,27 +352,16 @@ def create_table(
         else:
             query = None
 
-        type_mapper = self.compiler.type_mapper
-        column_defs = [
-            sge.ColumnDef(
-                this=sg.to_identifier(colname, quoted=quoted),
-                kind=type_mapper.from_ibis(typ),
-                constraints=(
-                    None
-                    if typ.nullable
-                    else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
-                ),
-            )
-            for colname, typ in (schema or table.schema()).items()
-        ]
-
         if overwrite:
             temp_name = util.gen_name(f"{self.name}_table")
         else:
             temp_name = name
 
+        if not schema:
+            schema = table.schema()
+
         table = sg.table(temp_name, catalog=database, quoted=quoted)
-        target = sge.Schema(this=table, expressions=column_defs)
+        target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
 
         create_stmt = sge.Create(kind="TABLE", this=target)
 

diff --git a/ibis/backends/mssql/__init__.py b/ibis/backends/mssql/__init__.py
@@ -639,29 +639,19 @@ def create_table(
         else:
             query = None
 
-        column_defs = [
-            sge.ColumnDef(
-                this=sg.to_identifier(colname, quoted=self.compiler.quoted),
-                kind=self.compiler.type_mapper.from_ibis(typ),
-                constraints=(
-                    None
-                    if typ.nullable
-                    else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
-                ),
-            )
-            for colname, typ in (schema or table.schema()).items()
-        ]
-
         if overwrite:
             temp_name = util.gen_name(f"{self.name}_table")
         else:
             temp_name = name
 
+        if not schema:
+            schema = table.schema()
+
         table = sg.table(
             "#" * temp + temp_name, catalog=catalog, db=db, quoted=self.compiler.quoted
         )
         raw_table = sg.table(temp_name, catalog=catalog, db=db, quoted=False)
-        target = sge.Schema(this=table, expressions=column_defs)
+        target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
 
         create_stmt = sge.Create(
             kind="TABLE",

diff --git a/ibis/backends/mysql/__init__.py b/ibis/backends/mysql/__init__.py
@@ -417,26 +417,16 @@ def create_table(
         else:
             query = None
 
-        column_defs = [
-            sge.ColumnDef(
-                this=sg.to_identifier(colname, quoted=self.compiler.quoted),
-                kind=self.compiler.type_mapper.from_ibis(typ),
-                constraints=(
-                    None
-                    if typ.nullable
-                    else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
-                ),
-            )
-            for colname, typ in (schema or table.schema()).items()
-        ]
-
         if overwrite:
             temp_name = util.gen_name(f"{self.name}_table")
         else:
             temp_name = name
 
+        if not schema:
+            schema = table.schema()
+
         table = sg.table(temp_name, catalog=database, quoted=self.compiler.quoted)
-        target = sge.Schema(this=table, expressions=column_defs)
+        target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
 
         create_stmt = sge.Create(
             kind="TABLE",

diff --git a/ibis/backends/oracle/__init__.py b/ibis/backends/oracle/__init__.py
@@ -433,26 +433,16 @@ def create_table(
         else:
             query = None
 
-        column_defs = [
-            sge.ColumnDef(
-                this=sg.to_identifier(colname, quoted=self.compiler.quoted),
-                kind=self.compiler.type_mapper.from_ibis(typ),
-                constraints=(
-                    None
-                    if typ.nullable
-                    else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
-                ),
-            )
-            for colname, typ in (schema or table.schema()).items()
-        ]
-
         if overwrite:
             temp_name = util.gen_name(f"{self.name}_table")
         else:
             temp_name = name
 
         initial_table = sg.table(temp_name, db=database, quoted=self.compiler.quoted)
-        target = sge.Schema(this=initial_table, expressions=column_defs)
+        target = sge.Schema(
+            this=initial_table,
+            expressions=(schema or table.schema()).to_sqlglot(self.dialect),
+        )
 
         create_stmt = sge.Create(
             kind="TABLE",
@@ -518,27 +508,11 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
         # only register if we haven't already done so
         if (name := op.name) not in self.list_tables():
             quoted = self.compiler.quoted
-            column_defs = [
-                sge.ColumnDef(
-                    this=sg.to_identifier(colname, quoted=quoted),
-                    kind=self.compiler.type_mapper.from_ibis(typ),
-                    constraints=(
-                        None
-                        if typ.nullable
-                        else [
-                            sg.exp.ColumnConstraint(
-                                kind=sg.exp.NotNullColumnConstraint()
-                            )
-                        ]
-                    ),
-                )
-                for colname, typ in schema.items()
-            ]
-
             create_stmt = sge.Create(
                 kind="TABLE",
                 this=sg.exp.Schema(
-                    this=sg.to_identifier(name, quoted=quoted), expressions=column_defs
+                    this=sg.to_identifier(name, quoted=quoted),
+                    expressions=schema.to_sqlglot(self.dialect),
                 ),
                 properties=sge.Properties(expressions=[sge.TemporaryProperty()]),
             ).sql(self.name)

diff --git a/ibis/backends/postgres/__init__.py b/ibis/backends/postgres/__init__.py
@@ -102,27 +102,11 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
         # only register if we haven't already done so
         if (name := op.name) not in self.list_tables():
             quoted = self.compiler.quoted
-            column_defs = [
-                sg.exp.ColumnDef(
-                    this=sg.to_identifier(colname, quoted=quoted),
-                    kind=self.compiler.type_mapper.from_ibis(typ),
-                    constraints=(
-                        None
-                        if typ.nullable
-                        else [
-                            sg.exp.ColumnConstraint(
-                                kind=sg.exp.NotNullColumnConstraint()
-                            )
-                        ]
-                    ),
-                )
-                for colname, typ in schema.items()
-            ]
-
             create_stmt = sg.exp.Create(
                 kind="TABLE",
                 this=sg.exp.Schema(
-                    this=sg.to_identifier(name, quoted=quoted), expressions=column_defs
+                    this=sg.to_identifier(name, quoted=quoted),
+                    expressions=schema.to_sqlglot(self.dialect),
                 ),
                 properties=sg.exp.Properties(expressions=[sge.TemporaryProperty()]),
             )
@@ -672,26 +656,16 @@ def create_table(
         else:
             query = None
 
-        column_defs = [
-            sge.ColumnDef(
-                this=sg.to_identifier(colname, quoted=self.compiler.quoted),
-                kind=self.compiler.type_mapper.from_ibis(typ),
-                constraints=(
-                    None
-                    if typ.nullable
-                    else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
-                ),
-            )
-            for colname, typ in (schema or table.schema()).items()
-        ]
-
         if overwrite:
             temp_name = util.gen_name(f"{self.name}_table")
         else:
             temp_name = name
 
+        if not schema:
+            schema = table.schema()
+
         table = sg.table(temp_name, db=database, quoted=self.compiler.quoted)
-        target = sge.Schema(this=table, expressions=column_defs)
+        target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
 
         create_stmt = sge.Create(
             kind="TABLE",