From e47e44b05221c521d536d1b69c5f7238cea5eded Mon Sep 17 00:00:00 2001 From: Yuyang Wang Date: Tue, 29 Aug 2023 14:36:28 -0700 Subject: [PATCH 1/5] SNOW-900244 --- docs/source/dataframe.rst | 1 + docs/source/session.rst | 4 +++- src/snowflake/snowpark/dataframe.py | 7 +++++++ src/snowflake/snowpark/session.py | 12 ++++++++++++ 4 files changed, 23 insertions(+), 1 deletion(-) diff --git a/docs/source/dataframe.rst b/docs/source/dataframe.rst index 93ba7afabc6..cab7c071b2c 100644 --- a/docs/source/dataframe.rst +++ b/docs/source/dataframe.rst @@ -125,3 +125,4 @@ DataFrame DataFrame.stat DataFrame.write DataFrame.is_cached + DataFrame.session diff --git a/docs/source/session.rst b/docs/source/session.rst index 0fe3a6d1b9b..468aa7a4a34 100644 --- a/docs/source/session.rst +++ b/docs/source/session.rst @@ -70,4 +70,6 @@ Snowpark Session Session.sql_simplifier_enabled Session.telemetry_enabled Session.udf - Session.udtf \ No newline at end of file + Session.udtf + Session.session_id + Session.connection \ No newline at end of file diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index 0f0ec72ba0c..54e1276e061 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -3320,6 +3320,13 @@ def na(self) -> DataFrameNaFunctions: """ return self._na + @property + def session(self) -> "snowflake.snowpark.Session": + """ + Returns a :class:`snowflake.snowpark.Session` object that provides access to the session current dataframe is relying on. + """ + return self._session + def describe(self, *cols: Union[str, List[str]]) -> "DataFrame": """ Computes basic statistics for numeric columns, which includes diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index ee28252bc22..52303e97da7 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -1681,6 +1681,18 @@ def read(self) -> "DataFrameReader": supported sources (e.g. a file in a stage) as a DataFrame.""" return DataFrameReader(self) + @property + def session_id(self) -> int: + """Returns an integer that represent the session id of this session.""" + return self._session_id + + @property + def connection(self) -> "ServerConnection": + """Returns a :class:`ServerConnection` object that allow you to access the connection between current session + and snowflake server""" + return self._conn + + def _run_query( self, query: str, From ee7332217182b6c5b8cddec8e13806d0727e330b Mon Sep 17 00:00:00 2001 From: Yuyang Wang Date: Tue, 29 Aug 2023 15:26:43 -0700 Subject: [PATCH 2/5] SNOW-900244 changelog --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ee02fe4730e..0e74655d616 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,10 @@ ### New Features - Accept `Iterable` objects input for `schema` when creating dataframes using `Session.create_dataframe`. +- Added the property `DataFrame.session` to return a session object. +- Added the property `Session.session_id` to return a integer represent session id. +- Added the property `Session.connection` to return a ServerConnection object . + ## 1.7.0 (2023-08-28) From fd275c6bdd621499ea8dbceb56427e3484df58bf Mon Sep 17 00:00:00 2001 From: Yuyang Wang Date: Tue, 5 Sep 2023 11:05:31 -0700 Subject: [PATCH 3/5] address comments --- CHANGELOG.md | 2 +- src/snowflake/snowpark/dataframe.py | 2 +- src/snowflake/snowpark/session.py | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e74655d616..34a25909db3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ - Accept `Iterable` objects input for `schema` when creating dataframes using `Session.create_dataframe`. - Added the property `DataFrame.session` to return a session object. -- Added the property `Session.session_id` to return a integer represent session id. +- Added the property `Session.session_id` to return an integer that represents session id. - Added the property `Session.connection` to return a ServerConnection object . diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index 54e1276e061..ed79befd53d 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -3323,7 +3323,7 @@ def na(self) -> DataFrameNaFunctions: @property def session(self) -> "snowflake.snowpark.Session": """ - Returns a :class:`snowflake.snowpark.Session` object that provides access to the session current dataframe is relying on. + Returns a :class:`snowflake.snowpark.Session` object that provides access to the session the current DataFrame is relying on. """ return self._session diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index 52303e97da7..caee7fea396 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -1683,14 +1683,14 @@ def read(self) -> "DataFrameReader": @property def session_id(self) -> int: - """Returns an integer that represent the session id of this session.""" + """Returns an integer that represents the session id of this session.""" return self._session_id @property - def connection(self) -> "ServerConnection": - """Returns a :class:`ServerConnection` object that allow you to access the connection between current session + def connection(self) -> "SnowflakeConnection": + """Returns a :class:`SnowflakeConnection` object that allow you to access the connection between current session and snowflake server""" - return self._conn + return self._conn._conn def _run_query( From 1354deb08b2d5dca2271e3ca4823bb7adee4ef3e Mon Sep 17 00:00:00 2001 From: Yuyang Wang Date: Mon, 11 Sep 2023 17:49:49 -0700 Subject: [PATCH 4/5] dataframe mock test --- CHANGELOG.md | 4 ++-- tests/unit/test_dataframe.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9e51164883..8b1cf5b3a06 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +7,9 @@ - Added support for VOLATILE/IMMUTABLE keyword when registering UDFs. - Added support for specifying clustering keys when saving dataframes using `DataFrame.save_as_table`. - Accept `Iterable` objects input for `schema` when creating dataframes using `Session.create_dataframe`. -- Added the property `DataFrame.session` to return a session object. +- Added the property `DataFrame.session` to return a `Session` object. - Added the property `Session.session_id` to return an integer that represents session id. -- Added the property `Session.connection` to return a ServerConnection object . +- Added the property `Session.connection` to return a `ServerConnection` object . - Added support for creating a Snowpark session from a configuration file or environment variables. diff --git a/tests/unit/test_dataframe.py b/tests/unit/test_dataframe.py index a58f1c55956..be78b8745c1 100644 --- a/tests/unit/test_dataframe.py +++ b/tests/unit/test_dataframe.py @@ -8,6 +8,7 @@ import pytest import snowflake.snowpark.session +from snowflake.snowpark.session import Session from snowflake.snowpark import ( DataFrame, DataFrameNaFunctions, @@ -292,3 +293,14 @@ def test_dataFrame_printSchema(capfd): out == "root\n |-- A: IntegerType() (nullable = False)\n |-- B: StringType() (nullable = True)\n" ) + + +def test_session(): + fake_session = mock.create_autospec(Session, _session_id=123456) + fake_session._analyzer = mock.Mock() + df = DataFrame(fake_session) + + assert(df.session == fake_session) + assert(df.session._session_id == fake_session._session_id) + + From c89e55a1d7c00d65b1ee86d7daf809631d9c99e2 Mon Sep 17 00:00:00 2001 From: Yuyang Wang Date: Tue, 12 Sep 2023 10:20:07 -0700 Subject: [PATCH 5/5] session mock unit test --- CHANGELOG.md | 4 ++-- src/snowflake/snowpark/session.py | 6 +++--- tests/unit/test_session.py | 24 +++++++++++++++++++++++- 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b1cf5b3a06..cd5a5259bb1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,8 +8,8 @@ - Added support for specifying clustering keys when saving dataframes using `DataFrame.save_as_table`. - Accept `Iterable` objects input for `schema` when creating dataframes using `Session.create_dataframe`. - Added the property `DataFrame.session` to return a `Session` object. -- Added the property `Session.session_id` to return an integer that represents session id. -- Added the property `Session.connection` to return a `ServerConnection` object . +- Added the property `Session.session_id` to return an integer that represents session ID. +- Added the property `Session.connection` to return a `SnowflakeConnection` object . - Added support for creating a Snowpark session from a configuration file or environment variables. diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index 5603a638882..339e83af2e6 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -1697,13 +1697,13 @@ def read(self) -> "DataFrameReader": @property def session_id(self) -> int: - """Returns an integer that represents the session id of this session.""" + """Returns an integer that represents the session ID of this session.""" return self._session_id @property def connection(self) -> "SnowflakeConnection": - """Returns a :class:`SnowflakeConnection` object that allow you to access the connection between current session - and snowflake server""" + """Returns a :class:`SnowflakeConnection` object that allows you to access the connection between the current session + and Snowflake server.""" return self._conn._conn diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 9a76d0604ee..032c926f1c5 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -3,7 +3,7 @@ # import json import os -from typing import Optional +from typing import Optional, Dict, Union from unittest import mock from unittest.mock import MagicMock @@ -368,3 +368,25 @@ def test_parse_table_name(): assert parse_table_name('*&^."abc".abc') # unsupported chars in unquoted ids with pytest.raises(SnowparkInvalidObjectNameException): assert parse_table_name('."abc".') # unsupported semantic + + +def test_session_id(): + fake_server_connection = mock.create_autospec(ServerConnection) + fake_server_connection.get_session_id = mock.Mock(return_value=123456) + session = Session(fake_server_connection) + + assert(session.session_id == 123456) + + +def test_connection(): + fake_snowflake_connection = mock.create_autospec(SnowflakeConnection) + fake_snowflake_connection._telemetry = mock.Mock() + fake_snowflake_connection._session_parameters = mock.Mock() + fake_snowflake_connection.is_closed = mock.Mock(return_value=False) + fake_options = {"": ""} + server_connection = ServerConnection(fake_options, fake_snowflake_connection) + session = Session(server_connection) + + assert(session.connection == session._conn._conn) + assert(session.connection == fake_snowflake_connection) +