summaryrefslogtreecommitdiffstats
path: root/tests/test_schema.py
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2022-10-21 09:29:26 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2022-10-21 09:29:26 +0000
commit8b4272814fb4585be120f183eb7c26bb8acde974 (patch)
tree85d56a8f5ac4ac94ab924d5bbc578586eeb2a998 /tests/test_schema.py
parentReleasing debian version 7.1.3-1. (diff)
downloadsqlglot-8b4272814fb4585be120f183eb7c26bb8acde974.tar.xz
sqlglot-8b4272814fb4585be120f183eb7c26bb8acde974.zip
Merging upstream version 9.0.1.
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'tests/test_schema.py')
-rw-r--r--tests/test_schema.py290
1 files changed, 290 insertions, 0 deletions
diff --git a/tests/test_schema.py b/tests/test_schema.py
new file mode 100644
index 0000000..bab97d8
--- /dev/null
+++ b/tests/test_schema.py
@@ -0,0 +1,290 @@
+import unittest
+
+from sqlglot import table
+from sqlglot.dataframe.sql import types as df_types
+from sqlglot.schema import MappingSchema, ensure_schema
+
+
+class TestSchema(unittest.TestCase):
+ def test_schema(self):
+ schema = ensure_schema(
+ {
+ "x": {
+ "a": "uint64",
+ }
+ }
+ )
+ self.assertEqual(
+ schema.column_names(
+ table(
+ "x",
+ )
+ ),
+ ["a"],
+ )
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x", db="db", catalog="c"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x", db="db"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x2"))
+
+ with self.assertRaises(ValueError):
+ schema.add_table(table("y", db="db"), {"b": "string"})
+ with self.assertRaises(ValueError):
+ schema.add_table(table("y", db="db", catalog="c"), {"b": "string"})
+
+ schema.add_table(table("y"), {"b": "string"})
+ schema_with_y = {
+ "x": {
+ "a": "uint64",
+ },
+ "y": {
+ "b": "string",
+ },
+ }
+ self.assertEqual(schema.schema, schema_with_y)
+
+ new_schema = schema.copy()
+ new_schema.add_table(table("z"), {"c": "string"})
+ self.assertEqual(schema.schema, schema_with_y)
+ self.assertEqual(
+ new_schema.schema,
+ {
+ "x": {
+ "a": "uint64",
+ },
+ "y": {
+ "b": "string",
+ },
+ "z": {
+ "c": "string",
+ },
+ },
+ )
+ schema.add_table(table("m"), {"d": "string"})
+ schema.add_table(table("n"), {"e": "string"})
+ schema_with_m_n = {
+ "x": {
+ "a": "uint64",
+ },
+ "y": {
+ "b": "string",
+ },
+ "m": {
+ "d": "string",
+ },
+ "n": {
+ "e": "string",
+ },
+ }
+ self.assertEqual(schema.schema, schema_with_m_n)
+ new_schema = schema.copy()
+ new_schema.add_table(table("o"), {"f": "string"})
+ new_schema.add_table(table("p"), {"g": "string"})
+ self.assertEqual(schema.schema, schema_with_m_n)
+ self.assertEqual(
+ new_schema.schema,
+ {
+ "x": {
+ "a": "uint64",
+ },
+ "y": {
+ "b": "string",
+ },
+ "m": {
+ "d": "string",
+ },
+ "n": {
+ "e": "string",
+ },
+ "o": {
+ "f": "string",
+ },
+ "p": {
+ "g": "string",
+ },
+ },
+ )
+
+ schema = ensure_schema(
+ {
+ "db": {
+ "x": {
+ "a": "uint64",
+ }
+ }
+ }
+ )
+ self.assertEqual(schema.column_names(table("x", db="db")), ["a"])
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x", db="db", catalog="c"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x", db="db2"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x2", db="db"))
+
+ with self.assertRaises(ValueError):
+ schema.add_table(table("y"), {"b": "string"})
+ with self.assertRaises(ValueError):
+ schema.add_table(table("y", db="db", catalog="c"), {"b": "string"})
+
+ schema.add_table(table("y", db="db"), {"b": "string"})
+ self.assertEqual(
+ schema.schema,
+ {
+ "db": {
+ "x": {
+ "a": "uint64",
+ },
+ "y": {
+ "b": "string",
+ },
+ }
+ },
+ )
+
+ schema = ensure_schema(
+ {
+ "c": {
+ "db": {
+ "x": {
+ "a": "uint64",
+ }
+ }
+ }
+ }
+ )
+ self.assertEqual(schema.column_names(table("x", db="db", catalog="c")), ["a"])
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x", db="db"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x", db="db", catalog="c2"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x", db="db2"))
+ with self.assertRaises(ValueError):
+ schema.column_names(table("x2", db="db"))
+
+ with self.assertRaises(ValueError):
+ schema.add_table(table("x"), {"b": "string"})
+ with self.assertRaises(ValueError):
+ schema.add_table(table("x", db="db"), {"b": "string"})
+
+ schema.add_table(table("y", db="db", catalog="c"), {"a": "string", "b": "int"})
+ self.assertEqual(
+ schema.schema,
+ {
+ "c": {
+ "db": {
+ "x": {
+ "a": "uint64",
+ },
+ "y": {
+ "a": "string",
+ "b": "int",
+ },
+ }
+ }
+ },
+ )
+ schema.add_table(table("z", db="db2", catalog="c"), {"c": "string", "d": "int"})
+ self.assertEqual(
+ schema.schema,
+ {
+ "c": {
+ "db": {
+ "x": {
+ "a": "uint64",
+ },
+ "y": {
+ "a": "string",
+ "b": "int",
+ },
+ },
+ "db2": {
+ "z": {
+ "c": "string",
+ "d": "int",
+ }
+ },
+ }
+ },
+ )
+ schema.add_table(table("m", db="db2", catalog="c2"), {"e": "string", "f": "int"})
+ self.assertEqual(
+ schema.schema,
+ {
+ "c": {
+ "db": {
+ "x": {
+ "a": "uint64",
+ },
+ "y": {
+ "a": "string",
+ "b": "int",
+ },
+ },
+ "db2": {
+ "z": {
+ "c": "string",
+ "d": "int",
+ }
+ },
+ },
+ "c2": {
+ "db2": {
+ "m": {
+ "e": "string",
+ "f": "int",
+ }
+ }
+ },
+ },
+ )
+
+ schema = ensure_schema(
+ {
+ "x": {
+ "a": "uint64",
+ }
+ }
+ )
+ self.assertEqual(schema.column_names(table("x")), ["a"])
+
+ schema = MappingSchema()
+ schema.add_table(table("x"), {"a": "string"})
+ self.assertEqual(
+ schema.schema,
+ {
+ "x": {
+ "a": "string",
+ }
+ },
+ )
+ schema.add_table(table("y"), df_types.StructType([df_types.StructField("b", df_types.StringType())]))
+ self.assertEqual(
+ schema.schema,
+ {
+ "x": {
+ "a": "string",
+ },
+ "y": {
+ "b": "string",
+ },
+ },
+ )
+
+ def test_schema_add_table_with_and_without_mapping(self):
+ schema = MappingSchema()
+ schema.add_table("test")
+ self.assertEqual(schema.column_names("test"), [])
+ schema.add_table("test", {"x": "string"})
+ self.assertEqual(schema.column_names("test"), ["x"])
+ schema.add_table("test", {"x": "string", "y": "int"})
+ self.assertEqual(schema.column_names("test"), ["x", "y"])
+ schema.add_table("test")
+ self.assertEqual(schema.column_names("test"), ["x", "y"])