Skip to content

Commit 9488115

Browse files
cpcloudcpcloud
authored
feat(api): add to_sqlglot method to Schema objects (#10063)
1 parent dfa55b6 commit 9488115

File tree

17 files changed

+180
-276
lines changed

17 files changed

+180
-276
lines changed

‎ibis/backends/bigquery/__init__.py

+4-15
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
)
3535
from ibis.backends.bigquery.datatypes import BigQuerySchema
3636
from ibis.backends.sql import SQLBackend
37-
from ibis.backends.sql.datatypes import BigQueryType
3837

3938
if TYPE_CHECKING:
4039
from collections.abc import Iterable, Mapping
@@ -1060,22 +1059,12 @@ def create_table(
10601059

10611060
table = _force_quote_table(table)
10621061

1063-
column_defs = [
1064-
sge.ColumnDef(
1065-
this=sg.to_identifier(name, quoted=self.compiler.quoted),
1066-
kind=BigQueryType.from_ibis(typ),
1067-
constraints=(
1068-
None
1069-
if typ.nullable or typ.is_array()
1070-
else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
1071-
),
1072-
)
1073-
for name, typ in (schema or {}).items()
1074-
]
1075-
10761062
stmt = sge.Create(
10771063
kind="TABLE",
1078-
this=sge.Schema(this=table, expressions=column_defs or None),
1064+
this=sge.Schema(
1065+
this=table,
1066+
expressions=schema.to_sqlglot(self.dialect) if schema else None,
1067+
),
10791068
replace=overwrite,
10801069
properties=sge.Properties(expressions=properties),
10811070
expression=None if obj is None else self.compile(obj),

‎ibis/backends/clickhouse/__init__.py

+1-7
Original file line numberDiff line numberDiff line change
@@ -674,13 +674,7 @@ def create_table(
674674

675675
this = sge.Schema(
676676
this=sg.table(name, db=database, quoted=self.compiler.quoted),
677-
expressions=[
678-
sge.ColumnDef(
679-
this=sg.to_identifier(name, quoted=self.compiler.quoted),
680-
kind=self.compiler.type_mapper.from_ibis(typ),
681-
)
682-
for name, typ in (schema or obj.schema()).items()
683-
],
677+
expressions=(schema or obj.schema()).to_sqlglot(self.dialect),
684678
)
685679
properties = [
686680
# the engine cannot be quoted, since clickhouse won't allow e.g.,

‎ibis/backends/datafusion/__init__.py

+4-14
Original file line numberDiff line numberDiff line change
@@ -657,20 +657,10 @@ def create_table(
657657
table_ident = sg.table(name, db=database, quoted=quoted)
658658

659659
if query is None:
660-
column_defs = [
661-
sge.ColumnDef(
662-
this=sg.to_identifier(colname, quoted=quoted),
663-
kind=self.compiler.type_mapper.from_ibis(typ),
664-
constraints=(
665-
None
666-
if typ.nullable
667-
else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
668-
),
669-
)
670-
for colname, typ in (schema or table.schema()).items()
671-
]
672-
673-
target = sge.Schema(this=table_ident, expressions=column_defs)
660+
target = sge.Schema(
661+
this=table_ident,
662+
expressions=(schema or table.schema()).to_sqlglot(self.dialect),
663+
)
674664
else:
675665
target = table_ident
676666

‎ibis/backends/duckdb/__init__.py

+4-14
Original file line numberDiff line numberDiff line change
@@ -175,19 +175,6 @@ def create_table(
175175
else:
176176
query = None
177177

178-
column_defs = [
179-
sge.ColumnDef(
180-
this=sg.to_identifier(colname, quoted=self.compiler.quoted),
181-
kind=self.compiler.type_mapper.from_ibis(typ),
182-
constraints=(
183-
None
184-
if typ.nullable
185-
else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
186-
),
187-
)
188-
for colname, typ in (schema or table.schema()).items()
189-
]
190-
191178
if overwrite:
192179
temp_name = util.gen_name("duckdb_table")
193180
else:
@@ -196,7 +183,10 @@ def create_table(
196183
initial_table = sg.table(
197184
temp_name, catalog=catalog, db=database, quoted=self.compiler.quoted
198185
)
199-
target = sge.Schema(this=initial_table, expressions=column_defs)
186+
target = sge.Schema(
187+
this=initial_table,
188+
expressions=(schema or table.schema()).to_sqlglot(self.dialect),
189+
)
200190

201191
create_stmt = sge.Create(
202192
kind="TABLE",

‎ibis/backends/exasol/__init__.py

+7-32
Original file line numberDiff line numberDiff line change
@@ -256,27 +256,13 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
256256
# only register if we haven't already done so
257257
if (name := op.name) not in self.list_tables():
258258
quoted = self.compiler.quoted
259-
column_defs = [
260-
sg.exp.ColumnDef(
261-
this=sg.to_identifier(colname, quoted=quoted),
262-
kind=self.compiler.type_mapper.from_ibis(typ),
263-
constraints=(
264-
None
265-
if typ.nullable
266-
else [
267-
sg.exp.ColumnConstraint(
268-
kind=sg.exp.NotNullColumnConstraint()
269-
)
270-
]
271-
),
272-
)
273-
for colname, typ in schema.items()
274-
]
275259

276260
ident = sg.to_identifier(name, quoted=quoted)
277261
create_stmt = sg.exp.Create(
278262
kind="TABLE",
279-
this=sg.exp.Schema(this=ident, expressions=column_defs),
263+
this=sg.exp.Schema(
264+
this=ident, expressions=schema.to_sqlglot(self.dialect)
265+
),
280266
)
281267
create_stmt_sql = create_stmt.sql(self.name)
282268

@@ -366,27 +352,16 @@ def create_table(
366352
else:
367353
query = None
368354

369-
type_mapper = self.compiler.type_mapper
370-
column_defs = [
371-
sge.ColumnDef(
372-
this=sg.to_identifier(colname, quoted=quoted),
373-
kind=type_mapper.from_ibis(typ),
374-
constraints=(
375-
None
376-
if typ.nullable
377-
else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
378-
),
379-
)
380-
for colname, typ in (schema or table.schema()).items()
381-
]
382-
383355
if overwrite:
384356
temp_name = util.gen_name(f"{self.name}_table")
385357
else:
386358
temp_name = name
387359

360+
if not schema:
361+
schema = table.schema()
362+
388363
table = sg.table(temp_name, catalog=database, quoted=quoted)
389-
target = sge.Schema(this=table, expressions=column_defs)
364+
target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
390365

391366
create_stmt = sge.Create(kind="TABLE", this=target)
392367

‎ibis/backends/mssql/__init__.py

+4-14
Original file line numberDiff line numberDiff line change
@@ -639,29 +639,19 @@ def create_table(
639639
else:
640640
query = None
641641

642-
column_defs = [
643-
sge.ColumnDef(
644-
this=sg.to_identifier(colname, quoted=self.compiler.quoted),
645-
kind=self.compiler.type_mapper.from_ibis(typ),
646-
constraints=(
647-
None
648-
if typ.nullable
649-
else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
650-
),
651-
)
652-
for colname, typ in (schema or table.schema()).items()
653-
]
654-
655642
if overwrite:
656643
temp_name = util.gen_name(f"{self.name}_table")
657644
else:
658645
temp_name = name
659646

647+
if not schema:
648+
schema = table.schema()
649+
660650
table = sg.table(
661651
"#" * temp + temp_name, catalog=catalog, db=db, quoted=self.compiler.quoted
662652
)
663653
raw_table = sg.table(temp_name, catalog=catalog, db=db, quoted=False)
664-
target = sge.Schema(this=table, expressions=column_defs)
654+
target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
665655

666656
create_stmt = sge.Create(
667657
kind="TABLE",

‎ibis/backends/mysql/__init__.py

+4-14
Original file line numberDiff line numberDiff line change
@@ -417,26 +417,16 @@ def create_table(
417417
else:
418418
query = None
419419

420-
column_defs = [
421-
sge.ColumnDef(
422-
this=sg.to_identifier(colname, quoted=self.compiler.quoted),
423-
kind=self.compiler.type_mapper.from_ibis(typ),
424-
constraints=(
425-
None
426-
if typ.nullable
427-
else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
428-
),
429-
)
430-
for colname, typ in (schema or table.schema()).items()
431-
]
432-
433420
if overwrite:
434421
temp_name = util.gen_name(f"{self.name}_table")
435422
else:
436423
temp_name = name
437424

425+
if not schema:
426+
schema = table.schema()
427+
438428
table = sg.table(temp_name, catalog=database, quoted=self.compiler.quoted)
439-
target = sge.Schema(this=table, expressions=column_defs)
429+
target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
440430

441431
create_stmt = sge.Create(
442432
kind="TABLE",

‎ibis/backends/oracle/__init__.py

+6-32
Original file line numberDiff line numberDiff line change
@@ -433,26 +433,16 @@ def create_table(
433433
else:
434434
query = None
435435

436-
column_defs = [
437-
sge.ColumnDef(
438-
this=sg.to_identifier(colname, quoted=self.compiler.quoted),
439-
kind=self.compiler.type_mapper.from_ibis(typ),
440-
constraints=(
441-
None
442-
if typ.nullable
443-
else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
444-
),
445-
)
446-
for colname, typ in (schema or table.schema()).items()
447-
]
448-
449436
if overwrite:
450437
temp_name = util.gen_name(f"{self.name}_table")
451438
else:
452439
temp_name = name
453440

454441
initial_table = sg.table(temp_name, db=database, quoted=self.compiler.quoted)
455-
target = sge.Schema(this=initial_table, expressions=column_defs)
442+
target = sge.Schema(
443+
this=initial_table,
444+
expressions=(schema or table.schema()).to_sqlglot(self.dialect),
445+
)
456446

457447
create_stmt = sge.Create(
458448
kind="TABLE",
@@ -518,27 +508,11 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
518508
# only register if we haven't already done so
519509
if (name := op.name) not in self.list_tables():
520510
quoted = self.compiler.quoted
521-
column_defs = [
522-
sge.ColumnDef(
523-
this=sg.to_identifier(colname, quoted=quoted),
524-
kind=self.compiler.type_mapper.from_ibis(typ),
525-
constraints=(
526-
None
527-
if typ.nullable
528-
else [
529-
sg.exp.ColumnConstraint(
530-
kind=sg.exp.NotNullColumnConstraint()
531-
)
532-
]
533-
),
534-
)
535-
for colname, typ in schema.items()
536-
]
537-
538511
create_stmt = sge.Create(
539512
kind="TABLE",
540513
this=sg.exp.Schema(
541-
this=sg.to_identifier(name, quoted=quoted), expressions=column_defs
514+
this=sg.to_identifier(name, quoted=quoted),
515+
expressions=schema.to_sqlglot(self.dialect),
542516
),
543517
properties=sge.Properties(expressions=[sge.TemporaryProperty()]),
544518
).sql(self.name)

‎ibis/backends/postgres/__init__.py

+6-32
Original file line numberDiff line numberDiff line change
@@ -102,27 +102,11 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
102102
# only register if we haven't already done so
103103
if (name := op.name) not in self.list_tables():
104104
quoted = self.compiler.quoted
105-
column_defs = [
106-
sg.exp.ColumnDef(
107-
this=sg.to_identifier(colname, quoted=quoted),
108-
kind=self.compiler.type_mapper.from_ibis(typ),
109-
constraints=(
110-
None
111-
if typ.nullable
112-
else [
113-
sg.exp.ColumnConstraint(
114-
kind=sg.exp.NotNullColumnConstraint()
115-
)
116-
]
117-
),
118-
)
119-
for colname, typ in schema.items()
120-
]
121-
122105
create_stmt = sg.exp.Create(
123106
kind="TABLE",
124107
this=sg.exp.Schema(
125-
this=sg.to_identifier(name, quoted=quoted), expressions=column_defs
108+
this=sg.to_identifier(name, quoted=quoted),
109+
expressions=schema.to_sqlglot(self.dialect),
126110
),
127111
properties=sg.exp.Properties(expressions=[sge.TemporaryProperty()]),
128112
)
@@ -672,26 +656,16 @@ def create_table(
672656
else:
673657
query = None
674658

675-
column_defs = [
676-
sge.ColumnDef(
677-
this=sg.to_identifier(colname, quoted=self.compiler.quoted),
678-
kind=self.compiler.type_mapper.from_ibis(typ),
679-
constraints=(
680-
None
681-
if typ.nullable
682-
else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
683-
),
684-
)
685-
for colname, typ in (schema or table.schema()).items()
686-
]
687-
688659
if overwrite:
689660
temp_name = util.gen_name(f"{self.name}_table")
690661
else:
691662
temp_name = name
692663

664+
if not schema:
665+
schema = table.schema()
666+
693667
table = sg.table(temp_name, db=database, quoted=self.compiler.quoted)
694-
target = sge.Schema(this=table, expressions=column_defs)
668+
target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
695669

696670
create_stmt = sge.Create(
697671
kind="TABLE",

0 commit comments

Comments
 (0)