Merging upstream version 9.0.1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

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"])