refactor: Rename metadata_repo to metadata_store for consistency across tools and remove compatibility modules
This commit is contained in:
parent
5719fdbe18
commit
a7afdfac8b
7 changed files with 17 additions and 54 deletions
|
|
@ -122,14 +122,14 @@ async def optimize_query_core(
|
||||||
|
|
||||||
llm = LLMToolFactory.create(settings)
|
llm = LLMToolFactory.create(settings)
|
||||||
storage = LocalFileStore()
|
storage = LocalFileStore()
|
||||||
metadata_repo = JsonMetadataStore(settings.storage.metadata_file)
|
metadata_store = JsonMetadataStore(settings.storage.metadata_file)
|
||||||
generator = PromptGeneratorFactory.create(db_type)
|
generator = PromptGeneratorFactory.create(db_type)
|
||||||
|
|
||||||
sql_file = _write_sql(sql, output_dir=output_dir, filename_prefix=f"opt_{db_type.value}")
|
sql_file = _write_sql(sql, output_dir=output_dir, filename_prefix=f"opt_{db_type.value}")
|
||||||
flow = OptimizeQueryFlow(
|
flow = OptimizeQueryFlow(
|
||||||
llm_tool=llm,
|
llm_tool=llm,
|
||||||
storage_tool=storage,
|
storage_tool=storage,
|
||||||
metadata_store=metadata_repo,
|
metadata_store=metadata_store,
|
||||||
prompt_generator=generator,
|
prompt_generator=generator,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -166,13 +166,13 @@ async def compare_optimizations_core(
|
||||||
|
|
||||||
llm = LLMToolFactory.create(settings)
|
llm = LLMToolFactory.create(settings)
|
||||||
storage = LocalFileStore()
|
storage = LocalFileStore()
|
||||||
metadata_repo = JsonMetadataStore(settings.storage.metadata_file)
|
metadata_store = JsonMetadataStore(settings.storage.metadata_file)
|
||||||
|
|
||||||
sql_file = _write_sql(sql, output_dir=output_dir, filename_prefix="compare")
|
sql_file = _write_sql(sql, output_dir=output_dir, filename_prefix="compare")
|
||||||
flow = CompareOptimizationsFlow(
|
flow = CompareOptimizationsFlow(
|
||||||
llm_tool=llm,
|
llm_tool=llm,
|
||||||
storage_tool=storage,
|
storage_tool=storage,
|
||||||
metadata_store=metadata_repo,
|
metadata_store=metadata_store,
|
||||||
prompt_generator_factory=PromptGeneratorFactory,
|
prompt_generator_factory=PromptGeneratorFactory,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -47,7 +47,7 @@ class CollectAndOptimizeWorstQueryFlow:
|
||||||
) -> None:
|
) -> None:
|
||||||
self._llm = llm_tool
|
self._llm = llm_tool
|
||||||
self._storage_tool = storage_tool
|
self._storage_tool = storage_tool
|
||||||
self._metadata_repo = metadata_store
|
self._metadata_store = metadata_store
|
||||||
self._prompt_generator = prompt_generator
|
self._prompt_generator = prompt_generator
|
||||||
|
|
||||||
def _fetch_worst(
|
def _fetch_worst(
|
||||||
|
|
@ -216,7 +216,7 @@ class CollectAndOptimizeWorstQueryFlow:
|
||||||
output = await OptimizeQueryFlow(
|
output = await OptimizeQueryFlow(
|
||||||
llm_tool=self._llm,
|
llm_tool=self._llm,
|
||||||
storage_tool=self._storage_tool,
|
storage_tool=self._storage_tool,
|
||||||
metadata_store=self._metadata_repo,
|
metadata_store=self._metadata_store,
|
||||||
prompt_generator=self._prompt_generator,
|
prompt_generator=self._prompt_generator,
|
||||||
).execute(
|
).execute(
|
||||||
sql_file=worst_query_file,
|
sql_file=worst_query_file,
|
||||||
|
|
|
||||||
|
|
@ -27,7 +27,7 @@ class CompareOptimizationsFlow:
|
||||||
Attributes:
|
Attributes:
|
||||||
_llm: LLM tool
|
_llm: LLM tool
|
||||||
_storage: File storage tool
|
_storage: File storage tool
|
||||||
_metadata_repo: Metadata store tool
|
_metadata_store: Metadata store tool
|
||||||
_prompt_generator_factory: Factory for creating prompt generators
|
_prompt_generator_factory: Factory for creating prompt generators
|
||||||
|
|
||||||
Examples:
|
Examples:
|
||||||
|
|
@ -60,7 +60,7 @@ class CompareOptimizationsFlow:
|
||||||
"""
|
"""
|
||||||
self._llm = llm_tool
|
self._llm = llm_tool
|
||||||
self._storage = storage_tool
|
self._storage = storage_tool
|
||||||
self._metadata_repo = metadata_store
|
self._metadata_store = metadata_store
|
||||||
self._prompt_factory = prompt_generator_factory
|
self._prompt_factory = prompt_generator_factory
|
||||||
|
|
||||||
async def execute(
|
async def execute(
|
||||||
|
|
@ -151,7 +151,7 @@ class CompareOptimizationsFlow:
|
||||||
return await OptimizeQueryFlow(
|
return await OptimizeQueryFlow(
|
||||||
llm_tool=self._llm,
|
llm_tool=self._llm,
|
||||||
storage_tool=self._storage,
|
storage_tool=self._storage,
|
||||||
metadata_store=self._metadata_repo,
|
metadata_store=self._metadata_store,
|
||||||
prompt_generator=self._prompt_factory.create(database_type),
|
prompt_generator=self._prompt_factory.create(database_type),
|
||||||
).execute(
|
).execute(
|
||||||
sql_file,
|
sql_file,
|
||||||
|
|
|
||||||
|
|
@ -30,7 +30,7 @@ class OptimizeQueryFlow:
|
||||||
Attributes:
|
Attributes:
|
||||||
_llm: LLM tool
|
_llm: LLM tool
|
||||||
_storage: File storage tool
|
_storage: File storage tool
|
||||||
_metadata_repo: Metadata store tool
|
_metadata_store: Metadata store tool
|
||||||
_prompt_generator: Prompt generator tool for target database
|
_prompt_generator: Prompt generator tool for target database
|
||||||
|
|
||||||
Examples:
|
Examples:
|
||||||
|
|
@ -63,7 +63,7 @@ class OptimizeQueryFlow:
|
||||||
"""
|
"""
|
||||||
self._llm = llm_tool
|
self._llm = llm_tool
|
||||||
self._storage = storage_tool
|
self._storage = storage_tool
|
||||||
self._metadata_repo = metadata_store
|
self._metadata_store = metadata_store
|
||||||
self._prompt_generator = prompt_generator
|
self._prompt_generator = prompt_generator
|
||||||
|
|
||||||
async def execute(
|
async def execute(
|
||||||
|
|
@ -144,7 +144,7 @@ class OptimizeQueryFlow:
|
||||||
pair_review = None
|
pair_review = None
|
||||||
|
|
||||||
# Persist metadata and output file and return
|
# Persist metadata and output file and return
|
||||||
await self._metadata_repo.save(metadata)
|
await self._metadata_store.save(metadata)
|
||||||
await self._save_output_file(sql_file, result, pair_review=pair_review, source_object=source_object)
|
await self._save_output_file(sql_file, result, pair_review=pair_review, source_object=source_object)
|
||||||
return OptimizationOutput.from_result(result, source_object=source_object)
|
return OptimizationOutput.from_result(result, source_object=source_object)
|
||||||
|
|
||||||
|
|
@ -198,7 +198,7 @@ class OptimizeQueryFlow:
|
||||||
Returns:
|
Returns:
|
||||||
Existing or new metadata instance
|
Existing or new metadata instance
|
||||||
"""
|
"""
|
||||||
if existing := await self._metadata_repo.get(query_hash):
|
if existing := await self._metadata_store.get(query_hash):
|
||||||
return existing
|
return existing
|
||||||
|
|
||||||
return OptimizationMetadata.create_new(query, database_type)
|
return OptimizationMetadata.create_new(query, database_type)
|
||||||
|
|
|
||||||
|
|
@ -1,11 +0,0 @@
|
||||||
"""Compatibility module for output objects.
|
|
||||||
|
|
||||||
Keep this shim to avoid import breakages while exposing the new output types.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from sql_optimizer_team.tools.engine.outputs.optimization_output import (
|
|
||||||
ComparisonOutput,
|
|
||||||
OptimizationOutput,
|
|
||||||
)
|
|
||||||
|
|
||||||
__all__ = ["OptimizationOutput", "ComparisonOutput"]
|
|
||||||
|
|
@ -84,13 +84,13 @@ class SQLServerStatsTool:
|
||||||
List of dicts containing performance metrics and SQL text.
|
List of dicts containing performance metrics and SQL text.
|
||||||
|
|
||||||
Examples:
|
Examples:
|
||||||
>>> repo = SQLServerStatsTool(instance="sis_vida_prd")
|
>>> tool = SQLServerStatsTool(instance="sis_vida_prd")
|
||||||
>>> # Get queries with highest CPU consumption
|
>>> # Get queries with highest CPU consumption
|
||||||
>>> queries = repo.fetch_worst_queries(limit=10, metric=QueryMetric.CPU)
|
>>> queries = tool.fetch_worst_queries(limit=10, metric=QueryMetric.CPU)
|
||||||
>>> # Get queries executed at least 100 times
|
>>> # Get queries executed at least 100 times
|
||||||
>>> queries = repo.fetch_worst_queries(limit=10, metric=QueryMetric.CPU, min_executions=100)
|
>>> queries = tool.fetch_worst_queries(limit=10, metric=QueryMetric.CPU, min_executions=100)
|
||||||
>>> # Order by most executed first, then by CPU
|
>>> # Order by most executed first, then by CPU
|
||||||
>>> queries = repo.fetch_worst_queries(limit=10, metric=QueryMetric.CPU, order_by_executions=True)
|
>>> queries = tool.fetch_worst_queries(limit=10, metric=QueryMetric.CPU, order_by_executions=True)
|
||||||
"""
|
"""
|
||||||
# Ensure safety and avoid SQL injection
|
# Ensure safety and avoid SQL injection
|
||||||
try:
|
try:
|
||||||
|
|
|
||||||
|
|
@ -1,26 +0,0 @@
|
||||||
"""Compatibility module for tool exceptions.
|
|
||||||
|
|
||||||
Keep this shim to avoid import breakages while exposing the new exceptions.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from sql_optimizer_team.tools.engine.types.tool_exceptions import (
|
|
||||||
ToolError,
|
|
||||||
InvalidDatabaseTypeError,
|
|
||||||
InvalidQueryHashError,
|
|
||||||
InvalidVersionError,
|
|
||||||
LLMProviderError,
|
|
||||||
MetadataNotFoundError,
|
|
||||||
OptimizationError,
|
|
||||||
QueryValidationError,
|
|
||||||
)
|
|
||||||
|
|
||||||
__all__ = [
|
|
||||||
"ToolError",
|
|
||||||
"QueryValidationError",
|
|
||||||
"LLMProviderError",
|
|
||||||
"MetadataNotFoundError",
|
|
||||||
"InvalidDatabaseTypeError",
|
|
||||||
"InvalidQueryHashError",
|
|
||||||
"InvalidVersionError",
|
|
||||||
"OptimizationError",
|
|
||||||
]
|
|
||||||
Loading…
Reference in a new issue