change db/qry constructor API to inject a new factory

twitter-archive · Apr 27, 2011 · 25f8ada · 25f8ada
1 parent 894fe49
commit 25f8ada
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 6 deletions.
diff --git a/project/build.properties b/project/build.properties
@@ -3,6 +3,6 @@
 project.organization=com.twitter
 project.name=querulous
 sbt.version=0.7.4
-project.version=2.0.2-SNAPSHOT
+project.version=2.0.2-stats-SNAPSHOT
 build.scala.versions=2.8.1
 project.initialize=false
diff --git a/src/main/scala/com/twitter/querulous/config/Database.scala b/src/main/scala/com/twitter/querulous/config/Database.scala
@@ -35,7 +35,11 @@ class Database {
   def timeout_=(t: TimingOutDatabase) { timeout = Some(t) }
   var memoize: Boolean = true
 
-  def apply(stats: StatsCollector): DatabaseFactory = {
+  def apply(stats: StatsCollector): DatabaseFactory = apply(stats, None)
+
+  def apply(stats: StatsCollector, statsFactory: DatabaseFactory => DatabaseFactory): DatabaseFactory = apply(stats, Some(statsFactory))
+
+  def apply(stats: StatsCollector, statsFactory: Option[DatabaseFactory => DatabaseFactory]): DatabaseFactory = {
     var factory: DatabaseFactory = pool.map(apacheConfig =>
       new ApachePoolingDatabaseFactory(
         apacheConfig.sizeMin,
@@ -54,6 +58,10 @@ class Database {
         timeoutConfig.poolSize)
     }
 
+    statsFactory.foreach { f =>
+      factory = f(factory)
+    }
+
     if (stats ne NullStatsCollector) {
       factory = new StatsCollectingDatabaseFactory(factory, stats)
     }

diff --git a/src/main/scala/com/twitter/querulous/config/Query.scala b/src/main/scala/com/twitter/querulous/config/Query.scala
@@ -29,7 +29,11 @@ class Query {
   var retries: Int = 0
   var debug: (String => Unit) = NoDebugOutput
 
-  def apply(statsCollector: StatsCollector): QueryFactory = {
+  def apply(statsCollector: StatsCollector): QueryFactory = apply(statsCollector, None)
+
+  def apply(statsCollector: StatsCollector, statsFactory: QueryFactory => QueryFactory): QueryFactory = apply(statsCollector, Some(statsFactory))
+
+  def apply(statsCollector: StatsCollector, statsFactory: Option[QueryFactory => QueryFactory]): QueryFactory = {
     var queryFactory: QueryFactory = new SqlQueryFactory
 
     if (!timeouts.isEmpty) {
@@ -40,6 +44,7 @@ class Query {
       queryFactory = new PerQueryTimingOutQueryFactory(new SqlQueryFactory, tupleTimeout)
     }
 
+    statsFactory.foreach { f => queryFactory = f(queryFactory) }
     if (statsCollector ne NullStatsCollector) {
       queryFactory = new StatsCollectingQueryFactory(queryFactory, statsCollector)
     }

diff --git a/src/main/scala/com/twitter/querulous/config/QueryEvaluator.scala b/src/main/scala/com/twitter/querulous/config/QueryEvaluator.scala
@@ -1,6 +1,8 @@
 package com.twitter.querulous.config
 
 import com.twitter.querulous._
+import com.twitter.querulous.database.DatabaseFactory
+import com.twitter.querulous.query.QueryFactory
 import com.twitter.util.Duration
 import evaluator._
 
@@ -16,10 +18,11 @@ class QueryEvaluator {
   var autoDisable: Option[AutoDisablingQueryEvaluator] = None
   def autoDisable_=(a: AutoDisablingQueryEvaluator) { autoDisable = Some(a) }
 
-
-  def apply(stats: StatsCollector): QueryEvaluatorFactory = {
+  def apply(stats: StatsCollector): QueryEvaluatorFactory = apply(stats, None, None)
+  def apply(stats: StatsCollector, dbStatsFactory: DatabaseFactory => DatabaseFactory, queryStatsFactory: QueryFactory => QueryFactory): QueryEvaluatorFactory = apply(stats, Some(dbStatsFactory), Some(queryStatsFactory))
+  def apply(stats: StatsCollector, dbStatsFactory: Option[DatabaseFactory => DatabaseFactory], queryStatsFactory: Option[QueryFactory => QueryFactory]): QueryEvaluatorFactory = {
     var factory: QueryEvaluatorFactory =
-      new StandardQueryEvaluatorFactory(database(stats), query(stats))
+      new StandardQueryEvaluatorFactory(database(stats, dbStatsFactory), query(stats, queryStatsFactory))
 
     autoDisable.foreach { disable =>
       factory = new AutoDisablingQueryEvaluatorFactory(