[xy] Create TEMP Snowflake table in data integration destination (mage-ai#4578)

* [xy] Create temp table in Snowflake destination.

* [xy] Pass temp_table flag.

* [xy] Update create command.

* [xy] Not commit for temp table.

* [xy] Set table_type in write_pandas method.

* [xy] Set table_type in write_pandas method.

* [xy] Restrict langchain package version.

* [xy] Restrict langchain_community version.

* [xy] Fix unit test.

Author: wangxiaoyou1993

mage_integrations/mage_integrations/connections/sql/base.py

@@ -19,15 +19,37 @@ def execute(
         self,
         query_strings: List[str],
         commit=False,
+        connection=None,
     ) -> List[List[tuple]]:
-        connection = self.build_connection()
+        """
+        Execute the provided SQL queries using the given connection, or create a new one if not
+        provided. If a new connection is created, it'll be closed automatically at the end.
+
+        Args:
+            query_strings (List[str]): List of SQL queries to execute.
+            commit (bool, optional): Whether to commit the transaction. Defaults to False.
+            connection (Optional[Any], optional): An existing connection to use. If None, a new
+                connection will be created. Defaults to None.
+
+        Returns:
+            List[List[Tuple]]: A list of result sets for each query executed. Each result set is
+                represented as a list of tuples.
+
+        Raises:
+            Any exceptions raised during the execution process may propagate upward.
+        """
+        new_connection_created = False
+        if connection is None:
+            connection = self.build_connection()
+            new_connection_created = True
 
         data = self.execute_with_connection(connection, query_strings)
 
         if commit:
             connection.commit()
 
-        self.close_connection(connection)
+        if new_connection_created:
+            self.close_connection(connection)
 
         return data
 

mage_integrations/mage_integrations/destinations/snowflake/__init__.py

@@ -60,6 +60,7 @@ def build_create_table_commands(
         schema_name: str,
         stream: str,
         table_name: str,
+        temp_table: bool = False,
         database_name: str = None,
         unique_constraints: List[str] = None,
     ) -> List[str]:
@@ -69,14 +70,15 @@ def build_create_table_commands(
                 convert_column_type,
                 lambda item_type_converted: 'ARRAY',
             ),
+            column_identifier=self.quote,
             columns=schema['properties'].keys(),
             full_table_name=self.full_table_name(
                 database_name,
                 schema_name,
                 table_name,
             ),
+            create_temporary_table=temp_table,
             unique_constraints=unique_constraints,
-            column_identifier=self.quote,
             use_lowercase=self.use_lowercase,
         )
 
@@ -341,6 +343,8 @@ def write_dataframe_to_table(
         database: str,
         schema: str,
         table: str,
+        connection=None,
+        temp_table: bool = False,
     ) -> List[List[tuple]]:
         """
         Write a Pandas DataFrame to a table in a Snowflake database.
@@ -363,19 +367,31 @@ def write_dataframe_to_table(
         """
         self.logger.info(
             f'write_pandas to: {database}.{schema}.{table}')
-        snowflake_connection = self.build_connection()
-        connection = snowflake_connection.build_connection()
+
+        new_connection_created = False
+        snowflake_connection = None
+        if connection is None:
+            snowflake_connection = self.build_connection()
+            connection = snowflake_connection.build_connection()
+            new_connection_created
         if self.disable_double_quotes:
             df.columns = [col.upper() for col in df.columns]
+
+        kwargs = dict(
+            database=database.upper() if self.disable_double_quotes else database,
+            schema=schema.upper() if self.disable_double_quotes else schema,
+            auto_create_table=False,
+        )
+        if temp_table:
+            kwargs['table_type'] = 'temp'
         success, num_chunks, num_rows, output = write_pandas(
             connection,
             df,
             table.upper() if self.disable_double_quotes else table,
-            database=database.upper() if self.disable_double_quotes else database,
-            schema=schema.upper() if self.disable_double_quotes else schema,
-            auto_create_table=False,
+            **kwargs,
         )
-        snowflake_connection.close_connection(connection)
+        if new_connection_created and snowflake_connection is not None:
+            snowflake_connection.close_connection(connection)
         self.logger.info(
             f'write_pandas completed: {success}, {num_chunks} chunks, {num_rows} rows.')
         self.logger.info(f'write_pandas output: {output}')
@@ -456,17 +472,30 @@ def process_queries(
                                 schema_name=schema,
                                 stream=None,
                                 table_name=f'temp_{table}',
+                                temp_table=True,
                                 database_name=database,
                                 unique_constraints=unique_constraints,
                             )
-
-                results += self.build_connection().execute(
-                    drop_temp_table_command + create_temp_table_command, commit=True)
+                # Run commands in one Snowflake session to leverage TEMP table
+                snowflake_connection = self.build_connection()
+                connection = snowflake_connection.build_connection()
+                results += snowflake_connection.execute(
+                    drop_temp_table_command + create_temp_table_command,
+                    commit=False,
+                    connection=connection,
+                )
 
                 # Outputs of write_dataframe_to_table are for temporary table only, thus not added
                 # to results
                 # results += self.write_dataframe_to_table(df, database, schema, f'temp_{table}')
-                self.write_dataframe_to_table(df, database, schema, f'temp_{table}')
+                self.write_dataframe_to_table(
+                    df,
+                    database,
+                    schema,
+                    f'temp_{table}',
+                    connection=connection,
+                    temp_table=True,
+                )
                 self.logger.info(
                     f'write_dataframe_to_table completed to: {full_table_name_temp}')
 
@@ -480,8 +509,13 @@ def process_queries(
 
                 self.logger.info(f'Merging {full_table_name_temp} into {full_table_name}')
                 self.logger.info(f'Dropping temporary table: {full_table_name_temp}')
-                results += self.build_connection().execute(
-                    merge_command + drop_temp_table_command, commit=True)
+                results += snowflake_connection.execute(
+                    merge_command + drop_temp_table_command,
+                    commit=True,
+                    connection=connection,
+                )
+                # Close connection after finishing running all commands
+                snowflake_connection.close_connection(connection)
                 self.logger.info(f'Merged and dropped temporary table: {full_table_name_temp}')
             else:
                 results += self.write_dataframe_to_table(df, database, schema, table)

mage_integrations/mage_integrations/tests/destinations/snowflake/test_snowflake.py

@@ -66,7 +66,7 @@ def test_create_table_commands(self):
                                                                  SCHEMA_NAME,
                                                                  STREAM,
                                                                  TABLE_NAME,
-                                                                 DATABASE_NAME)
+                                                                 database_name=DATABASE_NAME)
         self.assertEqual(
             table_commands,
             ['CREATE TABLE "test_db"."test"."test_table" ("ID" VARCHAR, "_USER" VARIANT)']

requirements.txt

@@ -65,6 +65,7 @@ google-cloud-storage~=2.5.0
 gspread==5.7.2
 kubernetes>=28.1.0
 langchain>=0.0.222; python_version >= '3.8'
+langchain_community<0.0.20
 mysql-connector-python~=8.2.0
 openai>=0.27.8, <1.0.0
 opentelemetry-exporter-prometheus==0.43b0

setup.py

@@ -43,6 +43,7 @@ def readme():
         'ai': [
             'astor>=0.8.1',
             'langchain>=0.0.222',
+            'langchain_community<0.0.20',
             'openai>=0.27.8, <1.0.0',
         ],
         'azure': [
@@ -166,6 +167,7 @@ def readme():
             'kafka-python==2.0.2',
             'kubernetes>=28.1.0',
             'langchain>=0.0.222',
+            'langchain_community<0.0.20',
             'ldap3==2.9.1',
             'nats-py==2.6.0',
             'nkeys~=0.1.0',