omnata-plugin-runtime 0.8.0a193__tar.gz → 0.8.0a194__tar.gz

Sign up to get free protection for your applications and to get access to all the features.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: omnata-plugin-runtime
3
- Version: 0.8.0a193
3
+ Version: 0.8.0a194
4
4
  Summary: Classes and common runtime components for building and running Omnata Plugins
5
5
  Author: James Weakley
6
6
  Author-email: james.weakley@omnata.com
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "omnata-plugin-runtime"
3
- version = "0.8.0-a193"
3
+ version = "0.8.0-a194"
4
4
  description = "Classes and common runtime components for building and running Omnata Plugins"
5
5
  authors = ["James Weakley <james.weakley@omnata.com>"]
6
6
  readme = "README.md"
@@ -47,7 +47,7 @@ from snowflake.snowpark import Session
47
47
  from snowflake.snowpark.functions import col
48
48
  from tenacity import Retrying, stop_after_attempt, wait_fixed, retry_if_exception_message
49
49
 
50
- from .logging import OmnataPluginLogHandler
50
+ from .logging import OmnataPluginLogHandler, logger
51
51
 
52
52
  from .api import (
53
53
  PluginMessage,
@@ -87,8 +87,9 @@ from .rate_limiting import (
87
87
  RateLimitState,
88
88
  RateLimitedSession
89
89
  )
90
+ from opentelemetry import trace
91
+ tracer = trace.get_tracer('omnata_plugin_runtime')
90
92
 
91
- logger = getLogger(__name__)
92
93
  SortDirectionType = Literal["asc", "desc"]
93
94
 
94
95
 
@@ -810,22 +811,23 @@ class OutboundSyncRequest(SyncRequest):
810
811
  logger.debug("applying results to table")
811
812
  # use a random table name with a random string to avoid collisions
812
813
  with self._snowflake_query_lock:
813
- for attempt in Retrying(stop=stop_after_attempt(30),wait=wait_fixed(2),reraise=True,retry=retry_if_exception_message(match=".*(is being|was) committed.*")):
814
- with attempt:
815
- success, nchunks, nrows, _ = write_pandas(
816
- conn=self._session._conn._cursor.connection, # pylint: disable=protected-access
817
- df=self._preprocess_results_dataframe(results_df),
818
- quote_identifiers=False,
819
- table_name=self._full_results_table_name,
820
- auto_create_table=False
821
- )
822
- if not success:
823
- raise ValueError(
824
- f"Failed to write results to table {self._full_results_table_name}"
814
+ with tracer.start_as_current_span("apply_results"):
815
+ for attempt in Retrying(stop=stop_after_attempt(30),wait=wait_fixed(2),reraise=True,retry=retry_if_exception_message(match=".*(is being|was) committed.*")):
816
+ with attempt:
817
+ success, nchunks, nrows, _ = write_pandas(
818
+ conn=self._session._conn._cursor.connection, # pylint: disable=protected-access
819
+ df=self._preprocess_results_dataframe(results_df),
820
+ quote_identifiers=False,
821
+ table_name=self._full_results_table_name,
822
+ auto_create_table=False
823
+ )
824
+ if not success:
825
+ raise ValueError(
826
+ f"Failed to write results to table {self._full_results_table_name}"
827
+ )
828
+ logger.debug(
829
+ f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
825
830
  )
826
- logger.debug(
827
- f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
828
- )
829
831
 
830
832
  def __dataframe_wrapper(
831
833
  self, data_frame: pandas.DataFrame, render_jinja: bool = True
@@ -1465,36 +1467,37 @@ class InboundSyncRequest(SyncRequest):
1465
1467
  """
1466
1468
  if len(results_df) > 0:
1467
1469
  with self._snowflake_query_lock:
1468
- for attempt in Retrying(stop=stop_after_attempt(30),wait=wait_fixed(2),reraise=True,retry=retry_if_exception_message(match=".*(is being|was) committed.*")):
1469
- with attempt:
1470
- logger.debug(
1471
- f"Applying {len(results_df)} results to {self._full_results_table_name}"
1472
- )
1473
- # try setting parquet engine here, since the engine parameter does not seem to make it through to the write_pandas function
1474
- success, nchunks, nrows, _ = write_pandas(
1475
- conn=self._session._conn._cursor.connection, # pylint: disable=protected-access
1476
- df=results_df,
1477
- table_name=self._full_results_table_name,
1478
- quote_identifiers=False, # already done in get_temp_table_name
1479
- # schema='INBOUND_RAW', # it seems to be ok to provide schema in the table name
1480
- table_type="transient"
1481
- )
1482
- if not success:
1483
- raise ValueError(
1484
- f"Failed to write results to table {self._full_results_table_name}"
1470
+ with tracer.start_as_current_span("apply_results"):
1471
+ for attempt in Retrying(stop=stop_after_attempt(30),wait=wait_fixed(2),reraise=True,retry=retry_if_exception_message(match=".*(is being|was) committed.*")):
1472
+ with attempt:
1473
+ logger.debug(
1474
+ f"Applying {len(results_df)} results to {self._full_results_table_name}"
1485
1475
  )
1486
- logger.debug(
1487
- f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
1488
- )
1489
- # temp tables aren't allowed
1490
- # snowflake_df = self._session.create_dataframe(results_df)
1491
- # snowflake_df.write.save_as_table(table_name=temp_table,
1492
- # mode='append',
1493
- # column_order='index',
1494
- # #create_temp_table=True
1495
- # )
1496
- for stream_name in stream_names:
1497
- self._results_exist[stream_name] = True
1476
+ # try setting parquet engine here, since the engine parameter does not seem to make it through to the write_pandas function
1477
+ success, nchunks, nrows, _ = write_pandas(
1478
+ conn=self._session._conn._cursor.connection, # pylint: disable=protected-access
1479
+ df=results_df,
1480
+ table_name=self._full_results_table_name,
1481
+ quote_identifiers=False, # already done in get_temp_table_name
1482
+ # schema='INBOUND_RAW', # it seems to be ok to provide schema in the table name
1483
+ table_type="transient"
1484
+ )
1485
+ if not success:
1486
+ raise ValueError(
1487
+ f"Failed to write results to table {self._full_results_table_name}"
1488
+ )
1489
+ logger.debug(
1490
+ f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
1491
+ )
1492
+ # temp tables aren't allowed
1493
+ # snowflake_df = self._session.create_dataframe(results_df)
1494
+ # snowflake_df.write.save_as_table(table_name=temp_table,
1495
+ # mode='append',
1496
+ # column_order='index',
1497
+ # #create_temp_table=True
1498
+ # )
1499
+ for stream_name in stream_names:
1500
+ self._results_exist[stream_name] = True
1498
1501
  else:
1499
1502
  logger.debug("Results dataframe is empty, not applying")
1500
1503
 
@@ -2137,14 +2140,16 @@ def __managed_inbound_processing_worker(
2137
2140
  logger.debug(f"stream returned from queue: {stream}")
2138
2141
  # restore the first argument, was originally the dataframe/generator but now it's the appropriately sized dataframe
2139
2142
  try:
2140
- logger.debug(f"worker {worker_index} processing stream {stream.stream_name}, invoking plugin class method {method.__name__}")
2141
- result = method(plugin_class_obj, *(stream, *method_args), **method_kwargs)
2142
- logger.debug(f"worker {worker_index} completed processing stream {stream.stream_name}")
2143
- if result is not None and result is False:
2144
- logger.info(f"worker {worker_index} requested that {stream.stream_name} be not marked as complete")
2145
- else:
2146
- logger.info(f"worker {worker_index} marking stream {stream.stream_name} as complete")
2147
- plugin_class_obj._sync_request.mark_stream_complete(stream.stream_name)
2143
+ with tracer.start_as_current_span("managed_inbound_processing") as managed_inbound_processing_span:
2144
+ logger.debug(f"worker {worker_index} processing stream {stream.stream_name}, invoking plugin class method {method.__name__}")
2145
+ managed_inbound_processing_span.set_attribute("stream_name", stream.stream_name)
2146
+ result = method(plugin_class_obj, *(stream, *method_args), **method_kwargs)
2147
+ logger.debug(f"worker {worker_index} completed processing stream {stream.stream_name}")
2148
+ if result is not None and result is False:
2149
+ logger.info(f"worker {worker_index} requested that {stream.stream_name} be not marked as complete")
2150
+ else:
2151
+ logger.info(f"worker {worker_index} marking stream {stream.stream_name} as complete")
2152
+ plugin_class_obj._sync_request.mark_stream_complete(stream.stream_name)
2148
2153
  except InterruptedWhileWaitingException:
2149
2154
  # If an inbound run is cancelled while waiting for rate limiting, this should mean that
2150
2155
  # the cancellation is handled elsewhere, so we don't need to do anything special here other than stop waiting
@@ -39,7 +39,7 @@ from .rate_limiting import ApiLimits, RateLimitState
39
39
  from opentelemetry import trace
40
40
 
41
41
  IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
42
- tracer = trace.get_tracer(__name__)
42
+ tracer = trace.get_tracer('omnata_plugin_runtime')
43
43
 
44
44
  class PluginEntrypoint:
45
45
  """