sqlframe 1.1.3__py3-none-any.whl

sqlframe/redshift/catalog.py

@@ -0,0 +1,127 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import fnmatch
+import typing as t
+
+from sqlglot import exp, parse_one
+
+from sqlframe.base.catalog import Function, _BaseCatalog
+from sqlframe.base.mixins.catalog_mixins import (
+    GetCurrentCatalogFromFunctionMixin,
+    GetCurrentDatabaseFromFunctionMixin,
+    ListCatalogsFromInfoSchemaMixin,
+    ListColumnsFromInfoSchemaMixin,
+    ListDatabasesFromInfoSchemaMixin,
+    ListTablesFromInfoSchemaMixin,
+    SetCurrentDatabaseFromSearchPathMixin,
+)
+from sqlframe.base.util import schema_, to_schema
+
+if t.TYPE_CHECKING:
+    from sqlframe.redshift.dataframe import RedshiftDataFrame
+    from sqlframe.redshift.session import RedshiftSession
+
+
+class RedshiftCatalog(
+    GetCurrentCatalogFromFunctionMixin["RedshiftSession", "RedshiftDataFrame"],
+    GetCurrentDatabaseFromFunctionMixin["RedshiftSession", "RedshiftDataFrame"],
+    ListDatabasesFromInfoSchemaMixin["RedshiftSession", "RedshiftDataFrame"],
+    ListCatalogsFromInfoSchemaMixin["RedshiftSession", "RedshiftDataFrame"],
+    SetCurrentDatabaseFromSearchPathMixin["RedshiftSession", "RedshiftDataFrame"],
+    ListTablesFromInfoSchemaMixin["RedshiftSession", "RedshiftDataFrame"],
+    ListColumnsFromInfoSchemaMixin["RedshiftSession", "RedshiftDataFrame"],
+    _BaseCatalog["RedshiftSession", "RedshiftDataFrame"],
+):
+    CURRENT_CATALOG_EXPRESSION: exp.Expression = exp.func("current_database")
+
+    def listFunctions(
+        self, dbName: t.Optional[str] = None, pattern: t.Optional[str] = None
+    ) -> t.List[Function]:
+        """
+        Returns a t.List of functions registered in the specified database.
+
+        .. versionadded:: 3.4.0
+
+        Parameters
+        ----------
+        dbName : str
+            name of the database to t.List the functions.
+            ``dbName`` can be qualified with catalog name.
+        pattern : str
+            The pattern that the function name needs to match.
+
+            .. versionchanged: 3.5.0
+                Adds ``pattern`` argument.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Function`.
+
+        Notes
+        -----
+        If no database is specified, the current database and catalog
+        are used. This API includes all temporary functions.
+
+        Examples
+        --------
+        >>> spark.catalog.t.listFunctions()
+        [Function(name=...
+
+        >>> spark.catalog.t.listFunctions(pattern="to_*")
+        [Function(name=...
+
+        >>> spark.catalog.t.listFunctions(pattern="*not_existing_func*")
+        []
+        """
+        if dbName is None:
+            schema = schema_(
+                db=exp.parse_identifier(self.currentDatabase(), dialect=self.session.input_dialect),
+                catalog=exp.parse_identifier(
+                    self.currentCatalog(), dialect=self.session.input_dialect
+                ),
+            )
+        else:
+            schema = to_schema(dbName, dialect=self.session.input_dialect)
+            if not schema.catalog:
+                schema.set("catalog", exp.parse_identifier(self.currentCatalog()))
+        query = parse_one(
+            f"""SELECT database_name as catalog, schema_name as namespace, function_name as name
+FROM svv_redshift_functions
+WHERE database_name = '{schema.catalog}'
+and schema_name = '{schema.db}'
+ORDER BY function_name;
+        """,
+            dialect=self.session.input_dialect,
+        )
+        functions = self.session._fetch_rows(query)
+        if pattern:
+            functions = [x for x in functions if fnmatch.fnmatch(x["name"], pattern)]
+        return [
+            Function(
+                name=x["name"],
+                catalog=x["catalog"],
+                namespace=[x["namespace"]],
+                description=None,
+                className="",
+                isTemporary=False,
+            )
+            for x in functions
+        ]
+
+    # def get_columns(self, table_name: t.Union[exp.Table, str]) -> t.Dict[str, exp.DataType]:
+    #     table = self.ensure_table(table_name)
+    #     if not table.catalog:
+    #         table.set("catalog", exp.parse_identifier(self.currentCatalog()))
+    #     if not table.db:
+    #         table.set("db", exp.parse_identifier(self.currentDatabase()))
+    #     sql = f"SHOW COLUMNS FROM TABLE {table.sql(dialect=self.session.input_dialect)}"
+    #     results = sorted(self.session._fetch_rows(sql), key=lambda x: x["ordinal_position"])
+    #     return {
+    #         row["column_name"]: exp.DataType.build(
+    #             row["data_type"], dialect=self.session.input_dialect, udt=True
+    #         )
+    #         for row in results
+    #     }

sqlframe/redshift/column.py

@@ -0,0 +1 @@
+from sqlframe.base.column import Column

sqlframe/redshift/dataframe.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import logging
+import sys
+import typing as t
+
+from sqlframe.base.dataframe import (
+    _BaseDataFrame,
+    _BaseDataFrameNaFunctions,
+    _BaseDataFrameStatFunctions,
+)
+from sqlframe.redshift.group import RedshiftGroupedData
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+if t.TYPE_CHECKING:
+    from sqlframe.redshift.readwriter import RedshiftDataFrameWriter
+    from sqlframe.redshift.session import RedshiftSession
+
+
+logger = logging.getLogger(__name__)
+
+
+class RedshiftDataFrameNaFunctions(_BaseDataFrameNaFunctions["RedshiftDataFrame"]):
+    pass
+
+
+class RedshiftDataFrameStatFunctions(_BaseDataFrameStatFunctions["RedshiftDataFrame"]):
+    pass
+
+
+class RedshiftDataFrame(
+    _BaseDataFrame[
+        "RedshiftSession",
+        "RedshiftDataFrameWriter",
+        "RedshiftDataFrameNaFunctions",
+        "RedshiftDataFrameStatFunctions",
+        "RedshiftGroupedData",
+    ]
+):
+    _na = RedshiftDataFrameNaFunctions
+    _stat = RedshiftDataFrameStatFunctions
+    _group_data = RedshiftGroupedData
+
+    def cache(self) -> Self:
+        logger.warning("Redshift does not support caching. Ignoring cache() call.")
+        return self
+
+    def persist(self) -> Self:
+        logger.warning("Redshift does not support persist. Ignoring persist() call.")
+        return self

sqlframe/redshift/functions.py

@@ -0,0 +1,18 @@
+import inspect
+import sys
+
+import sqlframe.base.functions
+
+module = sys.modules["sqlframe.base.functions"]
+globals().update(
+    {
+        name: func
+        for name, func in inspect.getmembers(module, inspect.isfunction)
+        if hasattr(func, "unsupported_engines")
+        and "redshift" not in func.unsupported_engines
+        and "*" not in func.unsupported_engines
+    }
+)
+
+
+from sqlframe.base.function_alternatives import e_literal as e  # noqa

sqlframe/redshift/group.py

@@ -0,0 +1,14 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.group import _BaseGroupedData
+
+if t.TYPE_CHECKING:
+    from sqlframe.redshift.dataframe import RedshiftDataFrame
+
+
+class RedshiftGroupedData(_BaseGroupedData["RedshiftDataFrame"]):
+    pass

sqlframe/redshift/readwriter.py

@@ -0,0 +1,29 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.mixins.readwriter_mixins import PandasLoaderMixin, PandasWriterMixin
+from sqlframe.base.readerwriter import (
+    _BaseDataFrameReader,
+    _BaseDataFrameWriter,
+)
+
+if t.TYPE_CHECKING:
+    from sqlframe.redshift.session import RedshiftSession  # noqa
+    from sqlframe.redshift.dataframe import RedshiftDataFrame  # noqa
+
+
+class RedshiftDataFrameReader(
+    PandasLoaderMixin["RedshiftSession", "RedshiftDataFrame"],
+    _BaseDataFrameReader["RedshiftSession", "RedshiftDataFrame"],
+):
+    pass
+
+
+class RedshiftDataFrameWriter(
+    PandasWriterMixin["RedshiftSession", "RedshiftDataFrame"],
+    _BaseDataFrameWriter["RedshiftSession", "RedshiftDataFrame"],
+):
+    pass

sqlframe/redshift/session.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import typing as t
+import warnings
+
+from sqlframe.base.session import _BaseSession
+from sqlframe.redshift.catalog import RedshiftCatalog
+from sqlframe.redshift.dataframe import RedshiftDataFrame
+from sqlframe.redshift.readwriter import (
+    RedshiftDataFrameReader,
+    RedshiftDataFrameWriter,
+)
+
+if t.TYPE_CHECKING:
+    from redshift_connector.core import Connection as RedshiftConnection
+else:
+    RedshiftConnection = t.Any
+
+
+class RedshiftSession(
+    _BaseSession[  # type: ignore
+        RedshiftCatalog,
+        RedshiftDataFrameReader,
+        RedshiftDataFrameWriter,
+        RedshiftDataFrame,
+        RedshiftConnection,
+    ],
+):
+    _catalog = RedshiftCatalog
+    _reader = RedshiftDataFrameReader
+    _writer = RedshiftDataFrameWriter
+    _df = RedshiftDataFrame
+
+    def __init__(self, conn: t.Optional[RedshiftConnection] = None):
+        warnings.warn(
+            "RedshiftSession is still in active development. Functions may not work as expected."
+        )
+        if not hasattr(self, "_conn"):
+            super().__init__(conn)
+
+    class Builder(_BaseSession.Builder):
+        DEFAULT_INPUT_DIALECT = "redshift"
+        DEFAULT_OUTPUT_DIALECT = "redshift"
+
+        @property
+        def session(self) -> RedshiftSession:
+            return RedshiftSession(**self._session_kwargs)
+
+        def getOrCreate(self) -> RedshiftSession:
+            self._set_session_properties()
+            return self.session
+
+    builder = Builder()

sqlframe/redshift/types.py

@@ -0,0 +1 @@
+from sqlframe.base.types import *

sqlframe/redshift/window.py

@@ -0,0 +1 @@
+from sqlframe.base.window import *

sqlframe/snowflake/__init__.py

@@ -0,0 +1,26 @@
+from sqlframe.snowflake.catalog import SnowflakeCatalog
+from sqlframe.snowflake.column import Column
+from sqlframe.snowflake.dataframe import (
+    SnowflakeDataFrame,
+    SnowflakeDataFrameNaFunctions,
+)
+from sqlframe.snowflake.group import SnowflakeGroupedData
+from sqlframe.snowflake.readwriter import (
+    SnowflakeDataFrameReader,
+    SnowflakeDataFrameWriter,
+)
+from sqlframe.snowflake.session import SnowflakeSession
+from sqlframe.snowflake.window import Window, WindowSpec
+
+__all__ = [
+    "SnowflakeCatalog",
+    "Column",
+    "SnowflakeDataFrame",
+    "SnowflakeDataFrameNaFunctions",
+    "SnowflakeGroupedData",
+    "SnowflakeDataFrameReader",
+    "SnowflakeDataFrameWriter",
+    "SnowflakeSession",
+    "Window",
+    "WindowSpec",
+]

sqlframe/snowflake/catalog.py

@@ -0,0 +1,134 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import fnmatch
+import json
+import typing as t
+
+from sqlglot import exp, parse_one
+
+from sqlframe.base.catalog import Function, _BaseCatalog
+from sqlframe.base.decorators import normalize
+from sqlframe.base.mixins.catalog_mixins import (
+    GetCurrentCatalogFromFunctionMixin,
+    GetCurrentDatabaseFromFunctionMixin,
+    ListCatalogsFromInfoSchemaMixin,
+    ListColumnsFromInfoSchemaMixin,
+    ListDatabasesFromInfoSchemaMixin,
+    ListTablesFromInfoSchemaMixin,
+    SetCurrentCatalogFromUseMixin,
+    SetCurrentDatabaseFromUseMixin,
+)
+from sqlframe.base.util import schema_, to_schema
+
+if t.TYPE_CHECKING:
+    from sqlframe.snowflake.session import SnowflakeSession  # noqa
+    from sqlframe.snowflake.dataframe import SnowflakeDataFrame  # noqa
+
+
+class SnowflakeCatalog(
+    SetCurrentCatalogFromUseMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    GetCurrentCatalogFromFunctionMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    GetCurrentDatabaseFromFunctionMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    ListDatabasesFromInfoSchemaMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    ListCatalogsFromInfoSchemaMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    SetCurrentDatabaseFromUseMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    ListTablesFromInfoSchemaMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    ListColumnsFromInfoSchemaMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    _BaseCatalog["SnowflakeSession", "SnowflakeDataFrame"],
+):
+    CURRENT_CATALOG_EXPRESSION: exp.Expression = exp.func("current_database")
+
+    @normalize(["dbName"])
+    def listFunctions(
+        self, dbName: t.Optional[str] = None, pattern: t.Optional[str] = None
+    ) -> t.List[Function]:
+        """
+        Returns a t.List of functions registered in the specified database.
+
+        .. versionadded:: 3.4.0
+
+        Parameters
+        ----------
+        dbName : str
+            name of the database to t.List the functions.
+            ``dbName`` can be qualified with catalog name.
+        pattern : str
+            The pattern that the function name needs to match.
+
+            .. versionchanged: 3.5.0
+                Adds ``pattern`` argument.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Function`.
+
+        Notes
+        -----
+        If no database is specified, the current database and catalog
+        are used. This API includes all temporary functions.
+
+        Examples
+        --------
+        >>> spark.catalog.t.listFunctions()
+        [Function(name=...
+
+        >>> spark.catalog.t.listFunctions(pattern="to_*")
+        [Function(name=...
+
+        >>> spark.catalog.t.listFunctions(pattern="*not_existing_func*")
+        []
+        """
+        if dbName is None:
+            schema = schema_(
+                db=exp.parse_identifier(self.currentDatabase(), dialect=self.session.input_dialect),
+                catalog=exp.parse_identifier(
+                    self.currentCatalog(), dialect=self.session.input_dialect
+                ),
+            )
+        else:
+            schema = to_schema(dbName, dialect=self.session.input_dialect)
+            if not schema.catalog:
+                schema.set("catalog", exp.parse_identifier(self.currentCatalog()))
+        query = parse_one(
+            f"""SHOW USER FUNCTIONS IN {schema.sql(dialect=self.session.input_dialect)}""",
+            dialect=self.session.input_dialect,
+        )
+        functions = self.session._fetch_rows(query)
+        if pattern:
+            functions = [x for x in functions if fnmatch.fnmatch(x["name"], pattern)]
+        return [
+            Function(
+                name=x["name"],
+                catalog=x["catalog_name"],
+                namespace=[x["schema_name"]],
+                description=None,
+                className="",
+                isTemporary=False,
+            )
+            for x in functions
+        ]
+
+    @normalize(["table_name"])
+    def get_columns(self, table_name: str) -> t.Dict[str, exp.DataType]:
+        table = exp.to_table(table_name)
+        if not table.catalog:
+            table.set(
+                "catalog",
+                exp.parse_identifier(self.currentCatalog(), dialect=self.session.input_dialect),
+            )
+        if not table.db:
+            table.set(
+                "db",
+                exp.parse_identifier(self.currentDatabase(), dialect=self.session.input_dialect),
+            )
+        sql = f"SHOW COLUMNS IN TABLE {table.sql(dialect=self.session.input_dialect)}"
+        results = self.session._fetch_rows(sql)
+        return {
+            row["column_name"]: exp.DataType.build(
+                json.loads(row["data_type"])["type"], dialect=self.session.input_dialect, udt=True
+            )
+            for row in results
+        }

sqlframe/snowflake/column.py

@@ -0,0 +1 @@
+from sqlframe.base.column import Column

sqlframe/snowflake/dataframe.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import logging
+import sys
+import typing as t
+
+from sqlframe.base.dataframe import (
+    _BaseDataFrame,
+    _BaseDataFrameNaFunctions,
+    _BaseDataFrameStatFunctions,
+)
+from sqlframe.snowflake.group import SnowflakeGroupedData
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+if t.TYPE_CHECKING:
+    from sqlframe.snowflake.readwriter import SnowflakeDataFrameWriter
+    from sqlframe.snowflake.session import SnowflakeSession
+
+
+logger = logging.getLogger(__name__)
+
+
+class SnowflakeDataFrameNaFunctions(_BaseDataFrameNaFunctions["SnowflakeDataFrame"]):
+    pass
+
+
+class SnowflakeDataFrameStatFunctions(_BaseDataFrameStatFunctions["SnowflakeDataFrame"]):
+    pass
+
+
+class SnowflakeDataFrame(
+    _BaseDataFrame[
+        "SnowflakeSession",
+        "SnowflakeDataFrameWriter",
+        "SnowflakeDataFrameNaFunctions",
+        "SnowflakeDataFrameStatFunctions",
+        "SnowflakeGroupedData",
+    ]
+):
+    _na = SnowflakeDataFrameNaFunctions
+    _stat = SnowflakeDataFrameStatFunctions
+    _group_data = SnowflakeGroupedData
+
+    def cache(self) -> Self:
+        logger.warning("Snowflake does not support caching. Ignoring cache() call.")
+        return self
+
+    def persist(self) -> Self:
+        logger.warning("Snowflake does not support persist. Ignoring persist() call.")
+        return self

sqlframe/snowflake/functions.py

@@ -0,0 +1,18 @@
+import inspect
+import sys
+
+import sqlframe.base.functions
+
+module = sys.modules["sqlframe.base.functions"]
+globals().update(
+    {
+        name: func
+        for name, func in inspect.getmembers(module, inspect.isfunction)
+        if hasattr(func, "unsupported_engines")
+        and "snowflake" not in func.unsupported_engines
+        and "*" not in func.unsupported_engines
+    }
+)
+
+
+from sqlframe.base.function_alternatives import e_literal as e  # noqa

sqlframe/snowflake/group.py

@@ -0,0 +1,14 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.group import _BaseGroupedData
+
+if t.TYPE_CHECKING:
+    from sqlframe.snowflake.dataframe import SnowflakeDataFrame
+
+
+class SnowflakeGroupedData(_BaseGroupedData["SnowflakeDataFrame"]):
+    pass

sqlframe/snowflake/readwriter.py

@@ -0,0 +1,29 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.mixins.readwriter_mixins import PandasLoaderMixin, PandasWriterMixin
+from sqlframe.base.readerwriter import (
+    _BaseDataFrameReader,
+    _BaseDataFrameWriter,
+)
+
+if t.TYPE_CHECKING:
+    from sqlframe.snowflake.session import SnowflakeSession  # noqa
+    from sqlframe.snowflake.dataframe import SnowflakeDataFrame  # noqa
+
+
+class SnowflakeDataFrameReader(
+    PandasLoaderMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    _BaseDataFrameReader["SnowflakeSession", "SnowflakeDataFrame"],
+):
+    pass
+
+
+class SnowflakeDataFrameWriter(
+    PandasWriterMixin["SnowflakeSession", "SnowflakeDataFrame"],
+    _BaseDataFrameWriter["SnowflakeSession", "SnowflakeDataFrame"],
+):
+    pass

sqlframe/snowflake/session.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import typing as t
+import warnings
+
+from sqlframe.base.session import _BaseSession
+from sqlframe.snowflake.catalog import SnowflakeCatalog
+from sqlframe.snowflake.dataframe import SnowflakeDataFrame
+from sqlframe.snowflake.readwriter import (
+    SnowflakeDataFrameReader,
+    SnowflakeDataFrameWriter,
+)
+
+if t.TYPE_CHECKING:
+    from snowflake.connector import SnowflakeConnection
+else:
+    SnowflakeConnection = t.Any
+
+
+class SnowflakeSession(
+    _BaseSession[  # type: ignore
+        SnowflakeCatalog,
+        SnowflakeDataFrameReader,
+        SnowflakeDataFrameWriter,
+        SnowflakeDataFrame,
+        SnowflakeConnection,
+    ],
+):
+    _catalog = SnowflakeCatalog
+    _reader = SnowflakeDataFrameReader
+    _writer = SnowflakeDataFrameWriter
+    _df = SnowflakeDataFrame
+
+    def __init__(self, conn: t.Optional[SnowflakeConnection] = None):
+        warnings.warn(
+            "SnowflakeSession is still in active development. Functions may not work as expected."
+        )
+        if not hasattr(self, "_conn"):
+            super().__init__(conn)
+
+    class Builder(_BaseSession.Builder):
+        DEFAULT_INPUT_DIALECT = "snowflake"
+        DEFAULT_OUTPUT_DIALECT = "snowflake"
+
+        @property
+        def session(self) -> SnowflakeSession:
+            return SnowflakeSession(**self._session_kwargs)
+
+        def getOrCreate(self) -> SnowflakeSession:
+            self._set_session_properties()
+            return self.session
+
+    builder = Builder()

sqlframe/snowflake/types.py

@@ -0,0 +1 @@
+from sqlframe.base.types import *

sqlframe/snowflake/window.py

@@ -0,0 +1 @@
+from sqlframe.base.window import *

sqlframe/spark/__init__.py

@@ -0,0 +1,23 @@
+from sqlframe.redshift.catalog import RedshiftCatalog
+from sqlframe.redshift.column import Column
+from sqlframe.redshift.dataframe import RedshiftDataFrame, RedshiftDataFrameNaFunctions
+from sqlframe.redshift.group import RedshiftGroupedData
+from sqlframe.redshift.readwriter import (
+    RedshiftDataFrameReader,
+    RedshiftDataFrameWriter,
+)
+from sqlframe.redshift.session import RedshiftSession
+from sqlframe.redshift.window import Window, WindowSpec
+
+__all__ = [
+    "RedshiftCatalog",
+    "Column",
+    "RedshiftDataFrame",
+    "RedshiftDataFrameNaFunctions",
+    "RedshiftGroupedData",
+    "RedshiftDataFrameReader",
+    "RedshiftDataFrameWriter",
+    "RedshiftSession",
+    "Window",
+    "WindowSpec",
+]